[audit-high] Replay-safe idempotency keys across SDK and framework wrappers

gateway/tests/test_framework_inprocess_validator.py

@@ -271,7 +271,10 @@ async def test_governed_tool_sends_declared_defaults_and_owns_key_handling(
     Declared-default parameters must travel to the server (the stub's
     default, not the server implementation's); a supplied blank idempotency
     key is rejected rather than silently replaced; an omitted (or None) key
-    derives a fresh one per call, making each call a new invocation.
+    is derived deterministically from the permit, tool, and arguments, so
+    retrying the same logical call replays the original receipt instead of
+    executing and charging again; a distinct action needs a distinct key,
+    supplied either explicitly or via ``action_id``.
     """
     registry = get_service_registry()
     calls = 0
@@ -334,21 +337,25 @@ async def defaults_stub(
                     await defaults_stub(message="explicit", idempotency_key=blank)
             assert calls == 1  # the rejected calls never reached the server
 
-            # Omitted key: a fresh uuid per call, so each call is a new
-            # invocation with its own receipt. Explicit None behaves the same.
+            # Omitted key: derived deterministically from the permit, tool,
+            # and arguments, so repeating the same logical call replays the
+            # original receipt without re-executing or charging again.
+            # Explicit None behaves the same.
             await defaults_stub(message="explicit")
             first_receipt = defaults_stub.last_receipt
             await defaults_stub(message="explicit")
             second_receipt = defaults_stub.last_receipt
             await defaults_stub(message="explicit", idempotency_key=None)
             third_receipt = defaults_stub.last_receipt
-            assert calls == 4
-            receipt_ids = {
-                first_receipt.receipt_id,
-                second_receipt.receipt_id,
-                third_receipt.receipt_id,
-            }
-            assert len(receipt_ids) == 3
+            assert calls == 2
+            assert second_receipt.receipt_id == first_receipt.receipt_id
+            assert third_receipt.receipt_id == first_receipt.receipt_id
+
+            # A distinct action needs a distinct key: the same arguments
+            # with a caller action id execute and bill separately.
+            await defaults_stub(message="explicit", action_id="inproc-action-2")
+            assert calls == 3
+            assert defaults_stub.last_receipt.receipt_id != first_receipt.receipt_id
         finally:
             await sdk.close()
     finally:

integrations/framework_integrations/_governed.py

@@ -23,10 +23,12 @@
 import asyncio
 import contextvars
 import functools
+import hashlib
 import inspect
-import uuid
+import json
+from collections.abc import Callable
 from decimal import Decimal
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 try:
     from b2a_sdk.edge_client import GovernedEdgeSession
@@ -41,6 +43,57 @@
     from b2a_sdk.models import Receipt
 
 
+def derive_governed_idempotency_key(
+    *,
+    permit_id: str,
+    tool_name: str,
+    arguments: dict[str, Any],
+    action_id: str | None = None,
+) -> str:
+    """Derive the idempotency key for one logical governed tool call.
+
+    Key rule: the key is a deterministic ``gov-`` prefixed SHA-256 hash of
+    the stable identity of the logical call: the session's permit id, the
+    tool name, the canonical (key-sorted JSON) bound arguments, and the
+    caller-supplied action id when one is given. Retrying the same logical
+    call therefore reuses the same key, and the server replays the original
+    receipt instead of charging again. No randomness is used, so a retry
+    can never mint a fresh key by accident.
+
+    A different action needs a different key. Arguments alone do not always
+    distinguish two actions: pass ``action_id`` (for example the
+    framework's tool call id) whenever two calls can carry identical
+    arguments but must bill separately. Without ``action_id``, two calls
+    with identical permit, tool, and arguments share a key and the second
+    replays the first.
+
+    Raises:
+        ValueError: if ``permit_id`` or ``tool_name`` is blank, or a
+            supplied ``action_id`` is blank.
+    """
+    if not isinstance(permit_id, str) or not permit_id.strip():
+        raise ValueError("permit_id must not be blank")
+    if not isinstance(tool_name, str) or not tool_name.strip():
+        raise ValueError("tool_name must not be blank")
+    if action_id is not None and (
+        not isinstance(action_id, str) or not action_id.strip()
+    ):
+        raise ValueError("action_id must not be blank; omit it to derive without one")
+    canonical = json.dumps(
+        {
+            "permit_id": permit_id,
+            "tool": tool_name,
+            "args": arguments,
+            "action_id": action_id,
+        },
+        sort_keys=True,
+        separators=(",", ":"),
+        default=str,
+    )
+    digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
+    return f"gov-{digest}"
+
+
 class GovernedToolWrapper:
     """Callable façade over one governed async tool call.
 
@@ -99,12 +152,20 @@ def governed_tool(
     raises ``RuntimeError`` at call time (matching ``b2a_sdk.decorators
     .billable``'s established behavior). On each call the wrapper:
 
-    1. consumes an ``idempotency_key`` keyword: a supplied non-blank key is
-       used as-is (retries that must replay — same receipt, no double
-       charge — need a caller-owned key); a supplied blank/whitespace key
-       raises ``ValueError`` rather than being silently replaced; when the
-       keyword is absent or ``None`` a fresh ``uuid4().hex`` is derived,
-       making that call a new invocation;
+    1. consumes an ``idempotency_key`` keyword and an ``action_id``
+       keyword (both are reserved: a stub parameter with either name is
+       consumed as key material, never sent to the server). A supplied
+       non-blank ``idempotency_key`` is used as-is: retries that must
+       replay (same receipt, no double charge) pass the same caller-owned
+       key. A supplied blank key raises ``ValueError`` rather than being
+       silently replaced. When the keyword is absent or ``None`` the key
+       is derived deterministically via
+       :func:`derive_governed_idempotency_key` from the session's permit
+       id, the tool name, the canonical bound arguments, and ``action_id``
+       when given, so retrying the same logical tool call reuses the key
+       instead of minting a fresh random one. Pass ``action_id`` (for
+       example the framework's tool call id) when two distinct actions
+       can carry identical arguments but must bill separately;
     2. binds the remaining arguments to the stub's signature **with the
        stub's declared defaults applied**, so an omitted parameter travels
        to the server as the stub's contractual default rather than letting
@@ -128,22 +189,37 @@ def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
 
         @functools.wraps(func)
         async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            supplied_key: str | None = None
             if "idempotency_key" in kwargs:
                 supplied = kwargs.pop("idempotency_key")
-                if supplied is None:
-                    idempotency_key = uuid.uuid4().hex
-                else:
-                    idempotency_key = str(supplied)
-                    if not idempotency_key.strip():
+                if supplied is not None:
+                    supplied_key = str(supplied)
+                    if not supplied_key.strip():
                         raise ValueError(
                             "idempotency_key must not be blank; omit it to "
-                            "derive a fresh key"
+                            "derive a stable key"
+                        )
+            action_id: str | None = None
+            if "action_id" in kwargs:
+                supplied_action = kwargs.pop("action_id")
+                if supplied_action is not None:
+                    action_id = str(supplied_action)
+                    if not action_id.strip():
+                        raise ValueError(
+                            "action_id must not be blank; omit it to derive without one"
                         )
-            else:
-                idempotency_key = uuid.uuid4().hex
             bound = signature.bind(*args, **kwargs)
             bound.apply_defaults()
             arguments = dict(bound.arguments)
+            if supplied_key is not None:
+                idempotency_key = supplied_key
+            else:
+                idempotency_key = derive_governed_idempotency_key(
+                    permit_id=session.permit_id,
+                    tool_name=name,
+                    arguments=arguments,
+                    action_id=action_id,
+                )
             result = await session.invoke(
                 name,
                 arguments,
@@ -169,4 +245,4 @@ def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
     return decorator
 
 
-__all__ = ["GovernedToolWrapper", "governed_tool"]
+__all__ = ["GovernedToolWrapper", "derive_governed_idempotency_key", "governed_tool"]

integrations/framework_integrations/tests/test_governed_idempotency_keys.py

@@ -0,0 +1,180 @@
+"""Stable idempotency keys in the framework ``governed_tool`` wrapper.
+
+The wrapper must derive the key from the stable identity of the logical
+tool call (permit id, tool name, canonical arguments, caller action id),
+never a fresh random value per invocation, so an ordinary retry of the
+same tool call reuses the key and replays instead of billing again.
+
+These tests use a fake session (no server, no database): the key choice
+happens in the wrapper before ``session.invoke`` is called.
+"""
+
+from typing import Any
+
+import pytest
+from framework_integrations._governed import (
+    derive_governed_idempotency_key,
+    governed_tool,
+)
+
+
+class _FakeResult:
+    def __init__(self) -> None:
+        self.receipt = object()
+        self.structured_content = {"echo": True}
+        self.content = [{"type": "text", "text": "echo"}]
+
+
+class _FakeSession:
+    """Records the idempotency key of every governed invocation."""
+
+    permit_id = "permit-test-1"
+    wallet_id = "wallet-test-1"
+
+    def __init__(self) -> None:
+        self.keys: list[str] = []
+        self.calls: list[tuple[str, dict[str, Any]]] = []
+
+    async def invoke(
+        self,
+        name: str,
+        arguments: dict[str, Any],
+        *,
+        idempotency_key: str,
+        estimated_credits=None,
+    ) -> _FakeResult:
+        self.keys.append(idempotency_key)
+        self.calls.append((name, dict(arguments)))
+        return _FakeResult()
+
+
+def _echo_tool(session: _FakeSession):
+    @governed_tool(session, tool_name="echo-tool")
+    async def echo(message: str) -> dict[str, Any]:
+        raise AssertionError("stub body must never run locally")
+
+    return echo
+
+
+def test_derive_helper_is_deterministic_and_bounded():
+    """Same inputs give the same short key twice."""
+    first = derive_governed_idempotency_key(
+        permit_id="permit-1", tool_name="echo", arguments={"b": 2, "a": 1}
+    )
+    second = derive_governed_idempotency_key(
+        permit_id="permit-1", tool_name="echo", arguments={"a": 1, "b": 2}
+    )
+    assert first == second
+    assert first.startswith("gov-")
+    assert len(first) <= 128
+
+
+def test_derive_helper_separates_distinct_actions():
+    """Different tool, args, permit, or action id give different keys."""
+    base = {"permit_id": "permit-1", "tool_name": "echo", "arguments": {"a": 1}}
+    assert derive_governed_idempotency_key(**base) != derive_governed_idempotency_key(
+        **{**base, "tool_name": "other"}
+    )
+    assert derive_governed_idempotency_key(**base) != derive_governed_idempotency_key(
+        **{**base, "arguments": {"a": 2}}
+    )
+    assert derive_governed_idempotency_key(**base) != derive_governed_idempotency_key(
+        **{**base, "permit_id": "permit-2"}
+    )
+    assert derive_governed_idempotency_key(**base) != derive_governed_idempotency_key(
+        **{**base, "action_id": "call-1"}
+    )
+
+
+def test_derive_helper_rejects_blank_inputs():
+    """Blank permit, tool, or action id fail closed before any dispatch."""
+    with pytest.raises(ValueError, match="permit_id"):
+        derive_governed_idempotency_key(permit_id="  ", tool_name="echo", arguments={})
+    with pytest.raises(ValueError, match="tool_name"):
+        derive_governed_idempotency_key(
+            permit_id="permit-1", tool_name="", arguments={}
+        )
+    with pytest.raises(ValueError, match="action_id"):
+        derive_governed_idempotency_key(
+            permit_id="permit-1",
+            tool_name="echo",
+            arguments={},
+            action_id="   ",
+        )
+
+
+async def test_wrapper_retry_of_same_tool_call_reuses_key():
+    """Retrying the same logical call sends the same derived key twice."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    await echo(message="hello")
+    await echo(message="hello")
+
+    assert len(session.keys) == 2
+    assert session.keys[0] == session.keys[1]
+    assert session.keys[0] == derive_governed_idempotency_key(
+        permit_id="permit-test-1",
+        tool_name="echo-tool",
+        arguments={"message": "hello"},
+    )
+
+
+async def test_wrapper_distinct_arguments_get_distinct_keys():
+    """Two different actions get different keys: no accidental replay."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    await echo(message="hello")
+    await echo(message="goodbye")
+
+    assert len(session.keys) == 2
+    assert session.keys[0] != session.keys[1]
+
+
+async def test_wrapper_action_id_disambiguates_identical_arguments():
+    """Identical arguments with different action ids bill separately."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    await echo(message="hello", action_id="framework-call-1")
+    await echo(message="hello", action_id="framework-call-2")
+    await echo(message="hello", action_id="framework-call-1")
+
+    assert len(session.keys) == 3
+    assert session.keys[0] != session.keys[1]
+    assert session.keys[0] == session.keys[2]
+
+
+async def test_wrapper_explicit_key_wins_over_derivation():
+    """A caller-supplied key is used as-is, even with an action id present."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    await echo(message="hello", idempotency_key="caller-key-1", action_id="call-9")
+
+    assert session.keys == ["caller-key-1"]
+
+
+async def test_wrapper_rejects_blank_key_and_blank_action_id():
+    """Blank key material is a caller bug: reject before dispatch."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    with pytest.raises(ValueError, match="idempotency_key"):
+        await echo(message="hello", idempotency_key="   ")
+    with pytest.raises(ValueError, match="action_id"):
+        await echo(message="hello", action_id="  ")
+    assert session.keys == []
+
+
+async def test_wrapper_coerces_non_string_action_id():
+    """A numeric framework run id still derives a stable key."""
+    session = _FakeSession()
+    echo = _echo_tool(session)
+
+    await echo(message="hello", action_id=7)
+    await echo(message="hello", action_id=7)
+
+    assert len(session.keys) == 2
+    assert session.keys[0] == session.keys[1]