diff --git a/integrations/openai-agent-middleware/README.md b/integrations/openai-agent-middleware/README.md index b401c161..7c42c478 100644 --- a/integrations/openai-agent-middleware/README.md +++ b/integrations/openai-agent-middleware/README.md @@ -20,9 +20,10 @@ plane's `Idempotency-Key` exists to capture, so `GovernedToolRunner`: 1. derives the key from the id (`oai-`) and **never invents one** — a tool call without an id is refused, because an invented key would turn the retry into a second charged action; -2. writes the derivation to an `OperationKeyStore` **before the first network - call**, so a crash between "the model asked" and "the receipt came back" - resumes with the same key; +2. writes the derivation, including the canonical JSON arguments, to an + `OperationKeyStore` **before the first network call**, so a crash between + "the model asked" and "the receipt came back" resumes with the same key + and the same arguments; 3. issues one permit per `(run_id, tool)` under a stable key, recording its complete JSON request before sending it. Retries reuse the original wallet, budget, scopes, and expiry, so the server receives identical request bytes. @@ -125,7 +126,10 @@ already saved. An incomplete legacy permit cannot safely reconstruct the original request and is refused. Legacy records that contain neither a wallet binding nor a request snapshot cannot be checked locally for a changed wallet; keep them with their original wallet. New operation records persist that -binding, including for completed-call replays. +binding, including for completed-call replays. Records written before +argument snapshots existed can still resume, but this runner cannot tell +locally if the arguments changed. Keep those stores with the original +transcript. ## What is refused client-side @@ -136,6 +140,10 @@ binding, including for completed-call replays. | blank, padded, non-ASCII, or non-printable run IDs or raw tool names; derived permit keys longer than 128 characters | invalid HTTP idempotency keys must fail before persistence or HTTP | | a function name that was not registered with this runner | the model cannot expand the application's allowed tool set | | a recorded tool call replayed under a different tool name | the operation identity is bound to one action | +| a recorded tool call replayed with different JSON arguments | the arguments are part of the action and are saved before the first network call, so a retry cannot commit a new body under the same id | +| a receipt whose wallet, tool, permit, or receipt id disagrees with the call | the runner must not treat another action's receipt as this call's result | +| an operation permit id that disagrees with the saved permit record | the call must use the permit that was issued for this run and tool | +| a blank wallet, a non-positive or unstorable permit budget, or a non-positive permit lifetime | those requests cannot succeed, and saving them would stick the failure to the run | | a recorded call resumed under a different run ID, or stored record IDs/keys that disagree with the call, run, or tool | persisted identity must match the deterministic operation and permit key derivations | | a persisted operation or permit request reused with a different wallet | the original issuer and subject remain bound to the action | | an incomplete legacy permit without a request snapshot | rebuilding its body could change the request under the original key | diff --git a/integrations/openai-agent-middleware/src/openai_b2a/__init__.py b/integrations/openai-agent-middleware/src/openai_b2a/__init__.py index 036a80d7..1baea6c2 100644 --- a/integrations/openai-agent-middleware/src/openai_b2a/__init__.py +++ b/integrations/openai-agent-middleware/src/openai_b2a/__init__.py @@ -3,8 +3,9 @@ The model's own ``tool_call.id`` is the operation identity: it is assigned once per tool call, lives in the conversation transcript, and therefore survives every retry of that call. :class:`GovernedToolRunner` derives the -trust plane's ``Idempotency-Key`` from it, persists the derivation before the -first network call, and never invents a key of its own. +trust plane's ``Idempotency-Key`` from it, persists that derivation and the +canonical arguments before the first network call, and never invents a key +of its own. """ from .client import B2AClient diff --git a/integrations/openai-agent-middleware/src/openai_b2a/runner.py b/integrations/openai-agent-middleware/src/openai_b2a/runner.py index ad09e667..d55756ee 100644 --- a/integrations/openai-agent-middleware/src/openai_b2a/runner.py +++ b/integrations/openai-agent-middleware/src/openai_b2a/runner.py @@ -22,6 +22,11 @@ request goes out, so any later attempt — a retry after a crash, or the next tool call in a resumed process — re-sends the identical request and the server replays the permit instead of rejecting a fresh timestamp under the same key. + +The arguments are part of the saved record. A retry whose JSON arguments are +not the same object is refused before any network call. Comparison uses +canonical JSON text, because Python treats ``True == 1``. A receipt is kept +only when its wallet, tool, permit, and receipt id match the call. """ from __future__ import annotations @@ -60,6 +65,42 @@ def _validate_identifier(value: str, label: str) -> str: return value +def _credit_amount_is_issuable(amount: Decimal) -> bool: + """A permit budget the trust plane will accept. + + Creation rejects zero, negative, and amounts Numeric(20, 8) cannot store + without rounding. Snapshotting one of those would stick the failure to + this run and tool, because retries resend the saved request. + """ + return ( + amount.is_finite() + and Decimal(0) < amount < Decimal(1000000000000) + and Decimal(f"{float(amount):.8f}") == amount + ) + + +def _canonical_arguments_json(arguments: Mapping[str, Any]) -> str: + """Stable JSON text for one tool call's arguments. + + Key order and insignificant whitespace are the same action. Python + equality is not: ``True == 1`` and ``1 == 1.0`` would hide a different body. + """ + try: + encoded = json.dumps( + arguments, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + except (TypeError, ValueError) as exc: + raise TypeError("tool call arguments must be JSON-serializable") from exc + parsed = json.loads(encoded) + if not isinstance(parsed, dict): + raise TypeError("tool call arguments must be a JSON object") + return encoded + + # ── Tool calls ─────────────────────────────────────────────────────────────── @@ -144,7 +185,11 @@ def function_name_for(tool_name: str) -> str: @dataclass(slots=True) class OperationRecord: - """What must survive a crash for a tool call to be retried as *one* action.""" + """What must survive a crash for a tool call to be retried as one action. + + ``arguments_json`` is the canonical JSON object for this call. It is + absent on records written before argument snapshots existed. + """ tool_call_id: str tool_name: str @@ -153,6 +198,7 @@ class OperationRecord: permit_id: str | None = None receipt_id: str | None = None wallet_id: str | None = None + arguments_json: str | None = None def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -354,6 +400,22 @@ def __init__( if not isinstance(run_id, str): raise TypeError("run_id must be a string") _validate_identifier(run_id, "run_id") + if not isinstance(wallet_id, str): + raise TypeError("wallet_id must be a string") + if not wallet_id.strip(): + raise ValueError("wallet_id must be a non-blank string") + if isinstance(permit_budget, bool) or not isinstance(permit_budget, Decimal | int): + raise TypeError("permit_budget must be a positive Decimal") + if not isinstance(permit_budget, Decimal): + permit_budget = Decimal(permit_budget) + if not _credit_amount_is_issuable(permit_budget): + raise ValueError("permit_budget must be a positive amount the trust plane can store") + if ( + isinstance(permit_ttl_minutes, bool) + or not isinstance(permit_ttl_minutes, int) + or permit_ttl_minutes <= 0 + ): + raise ValueError("permit_ttl_minutes must be a positive integer") self._client = client self._wallet_id = wallet_id self._run_id = run_id @@ -411,14 +473,15 @@ def permit_key_for(self, tool_name: str) -> str: # -- execution --------------------------------------------------------- async def run(self, tool_call: Any) -> GovernedToolResult: - """Execute one tool call as exactly one governed action. + """Execute one tool call as one governed action. - Retrying the same tool call — same id — reuses the recorded key and - permit, so the trust plane replays the original receipt instead of - charging again. + Retrying the same tool call, the same id and the same JSON arguments, + reuses the recorded key and permit, so the trust plane replays the + original receipt instead of charging again. """ call = normalize_tool_call(tool_call) tool_name = self.tool_name_for(call.name) + arguments_json = _canonical_arguments_json(call.arguments) operation_key = operation_key_for(call.id) permit_key = self.permit_key_for(tool_name) @@ -430,6 +493,7 @@ async def run(self, tool_call: Any) -> GovernedToolResult: idempotency_key=operation_key, permit_idempotency_key=permit_key, wallet_id=self._wallet_id, + arguments_json=arguments_json, ) elif record.tool_name != tool_name: raise ValueError( @@ -444,11 +508,23 @@ async def run(self, tool_call: Any) -> GovernedToolResult: raise ValueError("recorded operation identity does not match this call, run, and tool") if record.wallet_id is not None and record.wallet_id != self._wallet_id: raise ValueError("operation was recorded for a different wallet") + # Canonical text, not Python equality: True == 1 would hide a new body. + if record.arguments_json is not None and record.arguments_json != arguments_json: + raise ValueError( + f"tool call {call.id!r} was first recorded with different arguments; " + "refusing to replay it as a different action" + ) permit = self._key_store.get_permit(record.permit_idempotency_key) if permit is not None: if permit.permit_idempotency_key != permit_key or permit.tool_name != tool_name: raise ValueError("recorded permit identity does not match this run and tool") + if ( + record.permit_id is not None + and permit.permit_id is not None + and record.permit_id != permit.permit_id + ): + raise ValueError("recorded permit identity does not match this run and tool") if permit.request_payload is not None: if any( permit.request_payload.get(field) != self._wallet_id @@ -463,14 +539,29 @@ async def run(self, tool_call: Any) -> GovernedToolResult: # Durable before any network call, after local validation succeeds. self._key_store.put_operation(record) permit_id = await self._ensure_permit(record, permit) + if record.arguments_json is None: + # Legacy records have no argument snapshot. The server still + # rejects a changed body once it has seen the original call. + invoke_arguments = call.arguments + else: + invoke_arguments = json.loads(record.arguments_json) result = await self._client.invoke_tool( tool_name, - call.arguments, + invoke_arguments, wallet_id=self._wallet_id, permit_id=permit_id, idempotency_key=record.idempotency_key, ) - record.receipt_id = result.receipt.receipt_id + receipt = result.receipt + if ( + receipt.tool != tool_name + or receipt.wallet_id != self._wallet_id + or receipt.permit_id != permit_id + ): + raise ValueError("receipt does not match this wallet, tool, and permit") + if record.receipt_id is not None and record.receipt_id != receipt.receipt_id: + raise ValueError("receipt does not match the recorded action") + record.receipt_id = receipt.receipt_id self._key_store.put_operation(record) return GovernedToolResult( tool_call_id=call.id, diff --git a/integrations/openai-agent-middleware/tests/test_governed_flow.py b/integrations/openai-agent-middleware/tests/test_governed_flow.py index dccc60d4..8766987b 100644 --- a/integrations/openai-agent-middleware/tests/test_governed_flow.py +++ b/integrations/openai-agent-middleware/tests/test_governed_flow.py @@ -10,12 +10,12 @@ import json from datetime import UTC, datetime, timedelta -from decimal import Decimal +from decimal import Decimal, localcontext from types import SimpleNamespace import httpx import pytest -from b2a_sdk.errors import TransportError +from b2a_sdk.errors import APIError, TransportError from openai_b2a import ( B2AClient, @@ -108,9 +108,7 @@ def __call__(self, request: httpx.Request) -> httpx.Response: "isError": False, # Echo the invoked tool: the SDK refuses a receipt that # names a different call than the one just made. - "receipt": _receipt_payload( - key, receipt_id, tool=body["params"].get("name", TOOL) - ), + "receipt": _receipt_payload(key, receipt_id, body["params"]["name"]), }, }, ) @@ -697,3 +695,343 @@ async def test_loaded_permit_identity_must_match_its_lookup_key_and_tool( assert store_path.read_bytes() == before assert plane.permit_requests == [] and plane.invoke_requests == [] + + +# ── argument binding, receipts, and local config ──────────────────────────── + + +@pytest.mark.parametrize("raw", ["null", "1", '"hi"', "true"]) +def test_non_object_argument_json_is_refused(raw): + with pytest.raises(TypeError): + normalize_tool_call({"id": "call_1", "function": {"name": "f", "arguments": raw}}) + + +def test_invalid_argument_json_is_refused(): + with pytest.raises(ValueError): + normalize_tool_call({"id": "call_1", "function": {"name": "f", "arguments": "{"}}) + + +def test_tool_names_that_collapse_after_truncation_cannot_share_a_function(): + runner = GovernedToolRunner(_client(FakeTrustPlane()), wallet_id=WALLET, run_id="run-1") + stem = "n" * 64 + runner.register_tool(stem + ".one", description="first") + with pytest.raises(ValueError, match="already maps"): + runner.register_tool(stem + ".two", description="second") + + +async def test_different_arguments_are_refused_before_any_new_network_call(tmp_path): + plane = FakeTrustPlane() + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + _client(plane), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + await runner.run(_chat_tool_call("call_1", '{"text": "hi"}')) + before = store_path.read_bytes() + invokes = len(plane.invoke_requests) + + changed = _chat_tool_call("call_1", '{"text": "bye"}') + with pytest.raises(ValueError, match="arguments"): + await runner.run(changed) + + assert len(plane.invoke_requests) == invokes + assert store_path.read_bytes() == before + # A different call id is still its own action. + other = await runner.run(_chat_tool_call("call_2", '{"text": "bye"}')) + assert other.idempotency_key == "oai-call_2" + + +async def test_json_distinct_arguments_that_python_treats_as_equal_are_refused(): + """``True == 1`` in Python, but they are different JSON and a different action.""" + plane = FakeTrustPlane() + runner = GovernedToolRunner(_client(plane), wallet_id=WALLET, run_id="run-1") + runner.register_tool(TOOL, description="Append a note") + await runner.run(_chat_tool_call("call_1", '{"n": 1}')) + invokes = len(plane.invoke_requests) + changed = SimpleNamespace( + id="call_1", + function=SimpleNamespace(name=function_name_for(TOOL), arguments={"n": True}), + ) + with pytest.raises(ValueError, match="arguments"): + await runner.run(changed) + assert len(plane.invoke_requests) == invokes + + +async def test_argument_snapshot_is_fixed_before_the_permit_request(tmp_path): + """The body is saved before the permit call. A crash there must not let a + resume commit different arguments under the same tool call id.""" + store_path = tmp_path / "operations.json" + seen: list[str] = [] + + def failing(request: httpx.Request) -> httpx.Response: + seen.append(request.url.path) + raise httpx.ConnectError("boom", request=request) + + runner = GovernedToolRunner( + B2AClient( + api_key="test-key", + base_url="http://trust.test", + transport=httpx.MockTransport(failing), + ), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + with pytest.raises(TransportError): + await runner.run(_chat_tool_call("call_1", '{"b": 1, "a": 2}')) + + saved = json.loads(store_path.read_text()) + assert saved["operations"]["call_1"]["arguments_json"] == '{"a":2,"b":1}' + assert seen == ["/v1/permits"] + + seen.clear() + refused = GovernedToolRunner( + B2AClient( + api_key="test-key", + base_url="http://trust.test", + transport=httpx.MockTransport(failing), + ), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + refused.register_tool(TOOL, description="Append a note") + with pytest.raises(ValueError, match="arguments"): + await refused.run(_chat_tool_call("call_1", '{"a": 9}')) + assert seen == [] + + plane = FakeTrustPlane() + resumed = GovernedToolRunner( + _client(plane), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + resumed.register_tool(TOOL, description="Append a note") + result = await resumed.run(_chat_tool_call("call_1", '{"a": 2, "b": 1}')) + assert result.idempotency_key == "oai-call_1" + invoke = json.loads(plane.invoke_requests[0].content) + assert invoke["params"]["arguments"] == {"a": 2, "b": 1} + assert len(plane.permit_requests) == 1 + + +@pytest.mark.parametrize( + "arguments", + [{"bad": {1}}, {"n": float("nan")}], + ids=["set", "nan"], +) +async def test_unserializable_arguments_are_refused_before_persistence(tmp_path, arguments): + plane = FakeTrustPlane() + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + _client(plane), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + call = SimpleNamespace( + id="call_1", + function=SimpleNamespace(name=function_name_for(TOOL), arguments=arguments), + ) + with pytest.raises(TypeError): + await runner.run(call) + assert not store_path.exists() + assert plane.permit_requests == [] and plane.invoke_requests == [] + + +def _response_with_receipt_field(plane: FakeTrustPlane, field: str, value: str): + def handler(request: httpx.Request) -> httpx.Response: + response = plane(request) + if request.url.path != "/mcp/messages": + return response + payload = json.loads(response.content) + payload["result"]["receipt"][field] = value + return httpx.Response(200, json=payload) + + return handler + + +@pytest.mark.parametrize( + "field,value", + [ + ("wallet_id", "wallet-other"), + ("tool", "partner.notes.count"), + ("permit_id", "permit-other"), + ], +) +async def test_receipt_for_a_different_action_is_not_stored(tmp_path, field, value): + plane = FakeTrustPlane() + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + B2AClient( + api_key="test-key", + base_url="http://trust.test", + transport=httpx.MockTransport(_response_with_receipt_field(plane, field, value)), + ), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + with pytest.raises(APIError, match="mcp_receipt_does_not_match_request"): + await runner.run(_chat_tool_call("call_1")) + saved = json.loads(store_path.read_text()) + assert saved["operations"]["call_1"]["receipt_id"] is None + + +@pytest.mark.parametrize( + "field,value", + [ + ("wallet_id", "wallet-other"), + ("tool", "partner.notes.count"), + ("permit_id", "permit-other"), + ], +) +async def test_runner_rejects_mismatched_receipt_after_client_validation( + tmp_path, monkeypatch, field, value +): + plane = FakeTrustPlane() + client = _client(plane) + real_invoke = client.invoke_tool + + async def mismatched_invoke(*args, **kwargs): + result = await real_invoke(*args, **kwargs) + setattr(result.receipt, field, value) + return result + + monkeypatch.setattr(client, "invoke_tool", mismatched_invoke) + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + client, + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + with pytest.raises(ValueError, match="receipt does not match this wallet, tool, and permit"): + await runner.run(_chat_tool_call("call_1")) + saved = json.loads(store_path.read_text()) + assert saved["operations"]["call_1"]["receipt_id"] is None + + +async def test_replayed_receipt_id_must_match_the_recorded_action(tmp_path): + plane = FakeTrustPlane() + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + _client(plane), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + first = await runner.run(_chat_tool_call("call_1")) + assert first.receipt.receipt_id == "rcpt-1" + + def different_receipt(request: httpx.Request) -> httpx.Response: + response = plane(request) + if request.url.path != "/mcp/messages": + return response + payload = json.loads(response.content) + payload["result"]["receipt"]["receipt_id"] = "rcpt-other" + return httpx.Response(200, json=payload) + + resumed = GovernedToolRunner( + B2AClient( + api_key="test-key", + base_url="http://trust.test", + transport=httpx.MockTransport(different_receipt), + ), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + resumed.register_tool(TOOL, description="Append a note") + with pytest.raises(ValueError, match="receipt"): + await resumed.run(_chat_tool_call("call_1")) + saved = json.loads(store_path.read_text()) + assert saved["operations"]["call_1"]["receipt_id"] == "rcpt-1" + + +async def test_operation_permit_id_must_match_the_permit_record(tmp_path): + plane = FakeTrustPlane() + store_path = tmp_path / "operations.json" + runner = GovernedToolRunner( + _client(plane), + wallet_id=WALLET, + run_id="run-1", + key_store=JsonFileOperationKeyStore(store_path), + ) + runner.register_tool(TOOL, description="Append a note") + await runner.run(_chat_tool_call("call_1")) + persisted = json.loads(store_path.read_text()) + persisted["operations"]["call_1"]["permit_id"] = "permit-other" + store_path.write_text(json.dumps(persisted)) + before = store_path.read_bytes() + plane.invoke_requests.clear() + plane.permit_requests.clear() + + with pytest.raises(ValueError, match="permit"): + await runner.run(_chat_tool_call("call_1")) + + assert store_path.read_bytes() == before + assert plane.permit_requests == [] and plane.invoke_requests == [] + + +@pytest.mark.parametrize("wallet_id", ["", " "]) +def test_blank_wallet_is_refused(wallet_id): + with pytest.raises(ValueError, match="wallet"): + GovernedToolRunner(_client(FakeTrustPlane()), wallet_id=wallet_id, run_id="run-1") + + +def test_wallet_id_must_be_a_string(): + with pytest.raises(TypeError, match="wallet"): + GovernedToolRunner(_client(FakeTrustPlane()), wallet_id=None, run_id="run-1") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "budget", + [ + Decimal(0), + Decimal(-1), + Decimal("0.000000001"), + Decimal("100000000000.12345678"), + Decimal("999999999999.99999999"), + Decimal("NaN"), + True, + ], +) +def test_permit_budget_must_be_an_issuable_amount(budget): + with pytest.raises((TypeError, ValueError)): + GovernedToolRunner( + _client(FakeTrustPlane()), + wallet_id=WALLET, + run_id="run-1", + permit_budget=budget, # type: ignore[arg-type] + ) + + +def test_permit_budget_accepts_storable_amount_with_reduced_decimal_precision(): + with localcontext() as context: + context.prec = 10 + GovernedToolRunner( + _client(FakeTrustPlane()), + wallet_id=WALLET, + run_id="run-1", + permit_budget=Decimal(999_999_999_999), + ) + + +@pytest.mark.parametrize("ttl", [0, -1, True]) +def test_permit_ttl_must_be_a_positive_integer(ttl): + with pytest.raises(ValueError, match="permit_ttl"): + GovernedToolRunner( + _client(FakeTrustPlane()), + wallet_id=WALLET, + run_id="run-1", + permit_ttl_minutes=ttl, # type: ignore[arg-type] + )