diff --git a/arbiter/governance/__tests__/test_tool_execution_ledger.py b/arbiter/governance/__tests__/test_tool_execution_ledger.py new file mode 100644 index 00000000..ef2a9efe --- /dev/null +++ b/arbiter/governance/__tests__/test_tool_execution_ledger.py @@ -0,0 +1,463 @@ +"""Tests for arbiter/governance/tool_execution_ledger.py (PR1). + +The DynamoDB conditional write is NOT stubbed to always-succeed: ``FakeTable`` +faithfully evaluates the ``attribute_not_exists`` / ``status = :inflight`` / +``createdAt = :seen`` conditions this module emits and raises a real +``ConditionalCheckFailedException`` when they fail. The SaaS adapter IS +stubbed (a call counter). Together they let us prove exactly-once execution + +one recorded result, and — via a captured RED differential — that a +NON-conditional reserve lets both callers through. +""" +from __future__ import annotations + +import os +import sys + +import pytest +from botocore.exceptions import ClientError +from hypothesis import HealthCheck, given, settings, strategies as st + +_PROJECT_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..") +) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from arbiter.governance import tool_execution_ledger as ledger # noqa: E402 +from arbiter.governance.tool_execution_ledger import ( # noqa: E402 + OutcomeIndeterminateError, + RecordedToolFailure, + ReserveOutcome, + RetryableNoExecutionError, + ToolOutcomeError, + __reset_ledger_client_for_test, + execute_idempotent, + reserve, +) +from arbiter.workerWrapper.tool_idempotency import MODE_BYPASS, MODE_LEDGER, build_key # noqa: E402 + +TABLE_NAME = "citadel-tool-execution-ledger-test" + + +# --------------------------------------------------------------------------- +# Conditional-write fake (faithful, NOT always-succeed) +# --------------------------------------------------------------------------- + + +def _ccf(op: str) -> ClientError: + return ClientError( + {"Error": {"Code": "ConditionalCheckFailedException", "Message": "cond"}}, op + ) + + +class FakeTable: + """Stateful DynamoDB Table honoring the conditions this module emits.""" + + def __init__(self, name: str) -> None: + self.name = name + self.store: dict[tuple, dict] = {} + + @staticmethod + def _key(d: dict) -> tuple: + return (d[ledger.PK_ATTR], d[ledger.SK_ATTR]) + + def _resolve_name(self, token: str, names: dict) -> str: + return names.get(token, token) if token.startswith("#") else token + + def _eval_condition(self, expr, existing, names, values) -> bool: + if not expr: + return True + if "attribute_not_exists" in expr: + return existing is None + # Handle "LHS = :v [AND LHS = :v]" conjunctions. + for term in expr.split(" AND "): + lhs, rhs = [t.strip() for t in term.split("=")] + attr = self._resolve_name(lhs, names or {}) + expected = values[rhs] + if existing is None or existing.get(attr) != expected: + return False + return True + + def put_item(self, Item, ConditionExpression=None, **_kw): # noqa: N803 + key = self._key(Item) + existing = self.store.get(key) + if not self._eval_condition(ConditionExpression, existing, {}, {}): + raise _ccf("PutItem") + self.store[key] = dict(Item) + return {"ResponseMetadata": {"HTTPStatusCode": 200}} + + def get_item(self, Key, **_kw): # noqa: N803 + item = self.store.get((Key[ledger.PK_ATTR], Key[ledger.SK_ATTR])) + return {"Item": dict(item)} if item is not None else {} + + def update_item(self, Key, UpdateExpression, ConditionExpression=None, # noqa: N803 + ExpressionAttributeNames=None, ExpressionAttributeValues=None, **_kw): + key = (Key[ledger.PK_ATTR], Key[ledger.SK_ATTR]) + existing = self.store.get(key) + names = ExpressionAttributeNames or {} + values = ExpressionAttributeValues or {} + if not self._eval_condition(ConditionExpression, existing, names, values): + raise _ccf("UpdateItem") + assert UpdateExpression.startswith("SET ") + target = dict(existing) if existing else {ledger.PK_ATTR: key[0], ledger.SK_ATTR: key[1]} + for assignment in UpdateExpression[4:].split(","): + lhs, rhs = [t.strip() for t in assignment.split("=")] + attr = self._resolve_name(lhs, names) + target[attr] = values[rhs.strip()] + self.store[key] = target + return {"ResponseMetadata": {"HTTPStatusCode": 200}} + + def delete_item(self, Key, ConditionExpression=None, # noqa: N803 + ExpressionAttributeNames=None, ExpressionAttributeValues=None, **_kw): + key = (Key[ledger.PK_ATTR], Key[ledger.SK_ATTR]) + existing = self.store.get(key) + if not self._eval_condition(ConditionExpression, existing, + ExpressionAttributeNames or {}, ExpressionAttributeValues or {}): + raise _ccf("DeleteItem") + self.store.pop(key, None) + return {"ResponseMetadata": {"HTTPStatusCode": 200}} + + +class FakeResource: + def __init__(self): + self.tables: dict[str, FakeTable] = {} + + def Table(self, name): # noqa: N802 + return self.tables.setdefault(name, FakeTable(name)) + + +@pytest.fixture(autouse=True) +def _fake_ddb(monkeypatch): + __reset_ledger_client_for_test() + monkeypatch.setenv("TOOL_EXECUTION_LEDGER_TABLE", TABLE_NAME) + fake = FakeResource() + monkeypatch.setattr(ledger, "_get_dynamodb_resource", lambda: fake) + yield fake + __reset_ledger_client_for_test() + + +def _table(fake) -> FakeTable: + return fake.Table(TABLE_NAME) + + +def _key(): + return build_key("orgA", "exec1", "node1", 0, "createTicket", {"subject": "hi"}) + + +# Module-level alias: a dunder-prefixed name referenced bare inside a class +# body is name-mangled by Python (e.g. `__reset_ledger_client_for_test()` in +# a method becomes `_ClassName__reset_ledger_client_for_test`). Binding it to +# a plain-named module attribute here lets the property test call it safely. +_reset_ledger_client_for_test = __reset_ledger_client_for_test + + +# --------------------------------------------------------------------------- +# Acceptance: forced double delivery +# --------------------------------------------------------------------------- + + +class TestForcedDoubleDelivery: + def test_same_key_executes_once_one_recorded_result(self, _fake_ddb): + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + return {"status": "success", "ticketId": "T-1"} + + r1 = execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + r2 = execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + + assert calls["n"] == 1 # exactly one execution + assert r1["ticketId"] == "T-1" + assert r2["ticketId"] == "T-1" # loser gets recorded result + rows = [v for v in _table(_fake_ddb).store.values() if v.get("status") == "completed"] + assert len(rows) == 1 # one recorded result + + +# --------------------------------------------------------------------------- +# Concurrent race + RED proof +# --------------------------------------------------------------------------- + + +class TestConcurrentRace: + def test_loser_polls_then_retryable_never_executes(self, _fake_ddb): + pk, sk = _key() + # Winner A reserves and stays in_flight (does not finalize). + assert reserve(pk, sk, tool_name="createTicket").outcome == ReserveOutcome.WON + + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + return {"status": "success"} + + clock = {"t": 0.0} + wait_kwargs = { + "timeout": 0.3, + "interval": 0.1, + "clock": lambda: clock.__setitem__("t", clock["t"] + 0.2) or clock["t"], + "sleep": lambda _s: None, + } + with pytest.raises(RetryableNoExecutionError): + execute_idempotent( + pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, + run_tool=adapter, wait_kwargs=wait_kwargs, + ) + assert calls["n"] == 0 # loser NEVER executed + + def test_red_proof_nonconditional_reserve_lets_both_execute(self, _fake_ddb, monkeypatch): + # Captured RED: replace the conditional reserve with a NON-conditional + # one (always WON) and show both callers execute (the bug). This proves + # the conditional write — not app logic — is what enforces exactly-once. + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + return {"status": "success", "ticketId": f"T-{calls['n']}"} + + def nonconditional_reserve(_pk, _sk, *, tool_name, now=None): + _table(_fake_ddb).store[(_pk, _sk)] = { + ledger.PK_ATTR: _pk, ledger.SK_ATTR: _sk, "status": "in_flight", + } + return ledger.ReserveResult(ReserveOutcome.WON) + + monkeypatch.setattr(ledger, "reserve", nonconditional_reserve) + execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + assert calls["n"] == 2 # RED: non-conditional reserve double-executes + + def test_green_conditional_reserve_lets_one_execute(self, _fake_ddb): + # GREEN counterpart: the real conditional reserve yields exactly one. + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + return {"status": "success", "ticketId": f"T-{calls['n']}"} + + execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + execute_idempotent(pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, run_tool=adapter) + assert calls["n"] == 1 + + +# --------------------------------------------------------------------------- +# TTL + org scope +# --------------------------------------------------------------------------- + + +class TestTtlAndOrgScope: + def test_ttl_set_from_server_write_time(self, _fake_ddb, monkeypatch): + monkeypatch.setenv("TOOL_LEDGER_TTL_SECONDS", "172800") # 48h + pk, sk = _key() + reserve(pk, sk, tool_name="t", now=1_000_000.0) + row = _table(_fake_ddb).store[(pk, sk)] + assert row["ttl"] == 1_000_000 + 172800 # write-time + configured TTL + + def test_org_scope_isolates_partitions(self, _fake_ddb): + pk_a, sk = build_key("orgA", "exec1", "n", 0, "t", {"x": 1}) + pk_b, _ = build_key("orgB", "exec1", "n", 0, "t", {"x": 1}) + reserve(pk_a, sk, tool_name="t") + # org B's reserve for the "same" logical call is a different PK -> WON, + # never colliding with org A's row. + assert reserve(pk_b, sk, tool_name="t").outcome == ReserveOutcome.WON + assert ledger.get(pk_b, sk)[ledger.PK_ATTR] == "orgB#exec1" + # A cross-org read of the other org's key returns that org's row only. + assert ledger.get(pk_a, sk)[ledger.PK_ATTR] == "orgA#exec1" + + +# --------------------------------------------------------------------------- +# Failure matrix +# --------------------------------------------------------------------------- + + +class TestFailureMatrix: + def test_terminal_failure_recorded_and_replayed_without_reexec(self, _fake_ddb): + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + raise ToolOutcomeError("bad request", side_effect="applied", error_type="Http400") + + with pytest.raises(RecordedToolFailure): + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + # Replay: recorded terminal failure, NO re-execution. + with pytest.raises(RecordedToolFailure): + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + assert calls["n"] == 1 + + def test_retryable_not_sent_releases_and_reexecutes(self, _fake_ddb): + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + raise ToolOutcomeError("conn refused", side_effect="not_sent", retryable=True) + + with pytest.raises(RetryableNoExecutionError): + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + # Reservation released (status transition, not a delete) -> next + # attempt re-reserves (WON via conditional CAS) and may execute. + assert _table(_fake_ddb).store[(pk, sk)]["status"] == "released" + assert reserve(pk, sk, tool_name="t").outcome == ReserveOutcome.WON + + def test_unknown_outcome_is_fail_safe_indeterminate_never_reexec(self, _fake_ddb): + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + raise RuntimeError("5xx after send") # unclassified -> unknown + + with pytest.raises(OutcomeIndeterminateError): + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + row = _table(_fake_ddb).store[(pk, sk)] + assert row["status"] == "failed" + assert row["outcomeIndeterminate"] is True + assert row["retryable"] is False + # Replay: never re-executed, surfaced (not swallowed). + with pytest.raises(RecordedToolFailure): + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + assert calls["n"] == 1 + + def test_tool_error_result_recorded_and_returned(self, _fake_ddb): + pk, sk = _key() + + def adapter(): + return {"status": "error", "content": [{"text": "nope"}]} + + result = execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_LEDGER, run_tool=adapter) + assert result["status"] == "error" + assert _table(_fake_ddb).store[(pk, sk)]["status"] == "failed" + + +# --------------------------------------------------------------------------- +# Bypass path writes no ledger row +# --------------------------------------------------------------------------- + + +class TestBypassPath: + def test_bypass_skips_ledger_entirely(self, _fake_ddb): + pk, sk = _key() + calls = {"n": 0} + + def adapter(): + calls["n"] += 1 + return {"status": "success"} + + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_BYPASS, run_tool=adapter) + execute_idempotent(pk=pk, sk=sk, tool_name="t", mode=MODE_BYPASS, run_tool=adapter) + assert calls["n"] == 2 # no dedupe (read-only tool) + assert _table(_fake_ddb).store == {} # NO ledger row written + + +# --------------------------------------------------------------------------- +# Property: execute_idempotent never repeats the side effect (any JSON args) +# --------------------------------------------------------------------------- + +_json_scalars = st.one_of( + st.none(), + st.booleans(), + st.integers(min_value=-(2**53), max_value=2**53), + st.floats(allow_nan=False, allow_infinity=False, width=32), + st.text(max_size=20), +) +_json_values = st.recursive( + _json_scalars, + lambda children: st.one_of( + st.lists(children, max_size=5), + st.dictionaries(st.text(min_size=1, max_size=8), children, max_size=5), + ), + max_leaves=15, +) +# str-keyed JSON-serializable "args" object, as build_key's tool_input. +_tool_args = st.dictionaries(st.text(min_size=1, max_size=8), _json_values, max_size=6) + + +class TestExecuteIdempotentProperty: + """The mandated execution-level property test (checker finding). + + Drives ``execute_idempotent`` itself (not just canonicalization) under + Hypothesis: for ANY str-keyed JSON-serializable args, calling it twice + against a freshly derived key must invoke the adapter exactly once and + leave exactly one completed ledger row — never zero, never two. + """ + + @settings( + max_examples=100, + deadline=None, + suppress_health_check=[HealthCheck.function_scoped_fixture], + ) + @given(_tool_args) + def test_double_call_executes_adapter_exactly_once_per_example(self, tool_args): + # Fresh FakeTable + fresh counter + fresh ledger client cache for + # EVERY generated example — no state leaks across examples, so the + # assertion holds independently per example rather than only in + # aggregate. + _reset_ledger_client_for_test() + os.environ["TOOL_EXECUTION_LEDGER_TABLE"] = TABLE_NAME + fake = FakeResource() + original_get_resource = ledger._get_dynamodb_resource + ledger._get_dynamodb_resource = lambda: fake + try: + pk, sk = build_key("orgProp", "execProp", "nodeProp", 0, "createTicket", tool_args) + + calls = {"n": 0} + + def counting_adapter(): + calls["n"] += 1 + return {"status": "success", "ticketId": "T-prop"} + + r1 = execute_idempotent( + pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, + run_tool=counting_adapter, + ) + r2 = execute_idempotent( + pk=pk, sk=sk, tool_name="createTicket", mode=MODE_LEDGER, + run_tool=counting_adapter, + ) + + assert calls["n"] == 1, f"adapter invoked {calls['n']} times for args={tool_args!r}" + completed_rows = [ + v for v in _table(fake).store.values() if v.get("status") == ledger.STATUS_COMPLETED + ] + assert len(completed_rows) == 1, ( + f"expected exactly one completed ledger row for args={tool_args!r}, " + f"got {len(completed_rows)}" + ) + assert r1 == r2 == {"status": "success", "ticketId": "T-prop"} + finally: + ledger._get_dynamodb_resource = original_get_resource + _reset_ledger_client_for_test() + + +# --------------------------------------------------------------------------- +# Dead-holder reclaim +# --------------------------------------------------------------------------- + + +class TestDeadHolderReclaim: + def test_stale_inflight_is_reclaimed(self, _fake_ddb, monkeypatch): + monkeypatch.setenv("TOOL_LEDGER_LEASE_SECONDS", "900") + pk, sk = _key() + # A holder reserved long ago and died before finalizing. + reserve(pk, sk, tool_name="t", now=1000.0) + # A later attempt, well past the lease, reclaims via conditional CAS. + result = reserve(pk, sk, tool_name="t", now=1000.0 + 901) + assert result.outcome == ReserveOutcome.WON + assert result.reclaimed is True + + def test_fresh_inflight_is_not_reclaimed(self, _fake_ddb, monkeypatch): + monkeypatch.setenv("TOOL_LEDGER_LEASE_SECONDS", "900") + pk, sk = _key() + reserve(pk, sk, tool_name="t", now=1000.0) + result = reserve(pk, sk, tool_name="t", now=1000.0 + 10) + assert result.outcome == ReserveOutcome.IN_FLIGHT + + def test_missing_table_env_fails_closed(self, _fake_ddb, monkeypatch): + monkeypatch.delenv("TOOL_EXECUTION_LEDGER_TABLE", raising=False) + pk, sk = _key() + with pytest.raises(ledger.LedgerError): + reserve(pk, sk, tool_name="t") diff --git a/arbiter/governance/tool_execution_ledger.py b/arbiter/governance/tool_execution_ledger.py new file mode 100644 index 00000000..68ce8aa6 --- /dev/null +++ b/arbiter/governance/tool_execution_ledger.py @@ -0,0 +1,628 @@ +"""Tool-execution ledger — operational exactly-once dedupe (PR1 of 2). + +An org-scoped, TTL'd DynamoDB ledger that makes a governed tool call +exactly-once *within an attempt* and safe under a reservation race. It is the +operational backbone of tool-call idempotency; ``arbiter/workerWrapper/ +tool_idempotency.py`` derives the keys, this module coordinates the +reserve -> execute -> finalize protocol against them. + +Guarantee (state it precisely — do NOT collapse to a bare "exactly-once"): + +* **Exactly-once execution of the side effect is a GUARANTEE for calls that + resolve to the same key** — redelivery, same-attempt SDK/Strands retries, + and concurrent split-brain with identical keys. The reservation's + conditional write is what provides it: exactly one caller wins the reserve + and executes; every other caller is absorbed (recorded result) or bounced + with a retryable no-execution error, and NEVER executes. +* It is **best-effort across nondeterministic re-dispatch** (different keys), + which is closed only by the worker ``dispatchGeneration`` fence — DEFERRED + to PR2 and REQUIRED for the complete guarantee. Nothing here is the complete + guarantee on its own. +* Concurrent-loser *result delivery* is best-effort (a slow/dead holder may + yield a retryable error instead of the recorded result); side-effect + *execution* remains exactly-once. + +NOT an audit artifact: TTL is 48h operational (server-write-time based), +distinct from the 90-day ``arbiter/governance/ledger.py`` accountability +record. Do not conflate the two. + +Fail-closed: every failure path raises (never a bare ``except: pass``); a +ledger write/read failure must fail the call closed, never fall through to an +unprotected execution. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +import time +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable + +import boto3 +from botocore.exceptions import BotoCoreError, ClientError + +from arbiter.workerWrapper.tool_idempotency import MODE_BYPASS + +logger = logging.getLogger(__name__) + +# --- Attribute + status constants ------------------------------------------- + +PK_ATTR = "pk" # orgId#executionId +SK_ATTR = "sk" # nodeId#callIndex#toolName#argsHash + +STATUS_IN_FLIGHT = "in_flight" +STATUS_COMPLETED = "completed" +STATUS_FAILED = "failed" +# A reservation whose side effect provably did NOT happen is transitioned to +# 'released' (NOT deleted) so the worker IAM grant stays Put/Get/Update only +# (no dynamodb:DeleteItem — least privilege, per the design's IAM scope). A +# released row is re-reservable via a conditional CAS in reserve(). +STATUS_RELEASED = "released" + +# --- Tunables (env-overridable) --------------------------------------------- + +DEFAULT_TTL_SECONDS = 48 * 3600 # 48h operational dedupe window +DEFAULT_MAX_INLINE_BYTES = 300_000 # DDB 400KB item cap, headroom for attrs +DEFAULT_LEASE_SECONDS = 15 * 60 # dead-holder reclaim threshold +DEFAULT_POLL_TIMEOUT_SECONDS = 5.0 # concurrent-loser bounded poll ceiling +DEFAULT_POLL_INTERVAL_SECONDS = 0.1 + + +def _int_env(name: str, default: int) -> int: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = int(raw) + except (TypeError, ValueError): + return default + return value if value > 0 else default + + +def _float_env(name: str, default: float) -> float: + raw = os.environ.get(name) + if raw is None: + return default + try: + value = float(raw) + except (TypeError, ValueError): + return default + return value if value > 0 else default + + +def ttl_seconds() -> int: + return _int_env("TOOL_LEDGER_TTL_SECONDS", DEFAULT_TTL_SECONDS) + + +def max_inline_bytes() -> int: + return _int_env("TOOL_RESULT_MAX_INLINE_BYTES", DEFAULT_MAX_INLINE_BYTES) + + +def lease_seconds() -> int: + return _int_env("TOOL_LEDGER_LEASE_SECONDS", DEFAULT_LEASE_SECONDS) + + +# --- Exceptions -------------------------------------------------------------- + + +class LedgerError(Exception): + """Base for ledger failures. Callers MUST fail the call closed.""" + + +class RetryableNoExecutionError(LedgerError): + """Concurrent loser (or a released reservation) — no side effect occurred. + + Retryable: a later attempt may re-reserve and execute. The load-bearing + property is that the raising caller NEVER executed the side effect. + """ + + retryable = True + + +class OutcomeIndeterminateError(LedgerError): + """An un-tokened side-effecting call finished with an UNKNOWN outcome. + + Fail-safe toward no-duplicate (security consensus D2): the side effect may + or may not have applied, and there is no client token to dedupe a retry + end-to-end, so the call is NEVER auto-re-executed. Non-retryable; surfaced + upward, never swallowed. + """ + + retryable = False + + +class RecordedToolFailure(LedgerError): + """A prior terminal failure for this key, replayed without re-executing.""" + + retryable = False + + def __init__(self, message: str, *, recorded: dict[str, Any] | None = None): + super().__init__(message) + self.recorded = recorded or {} + + +class ToolOutcomeError(Exception): + """Adapter-raised classification of a failed tool call. + + Adapters that can classify their failure raise this so the coordinator + applies the correct failure-matrix branch: + + * ``side_effect='not_sent'`` + ``retryable=True`` — provably no side + effect (e.g. connection refused pre-send): reservation released, call + re-executable. + * ``side_effect='unknown'`` — outcome ambiguous (5xx/timeout after send): + fail-safe ``outcomeIndeterminate``, never re-executed. + * ``side_effect='applied'`` (terminal 4xx/validation): recorded failure, + non-retryable, returned on replay without re-execution. + + A bare ``Exception`` from the tool (not this type) is treated as + ``unknown`` — the fail-safe default. + """ + + def __init__( + self, + message: str, + *, + side_effect: str = "unknown", + retryable: bool = False, + error_type: str | None = None, + ): + super().__init__(message) + self.side_effect = side_effect + self.retryable = retryable + self.error_type = error_type or self.__class__.__name__ + + +# --- Reserve outcome --------------------------------------------------------- + + +class ReserveOutcome(Enum): + WON = "won" # caller may execute + HIT_COMPLETED = "hit_completed" # recorded success — return it, do NOT execute + HIT_FAILED = "hit_failed" # recorded terminal failure — replay, do NOT execute + IN_FLIGHT = "in_flight" # concurrent holder alive — poll then retryable error + + +@dataclass +class ReserveResult: + outcome: ReserveOutcome + row: dict[str, Any] | None = None + reclaimed: bool = False + + +# --- Lazy boto3 resource (QB-013-1: never construct at import time) ---------- + +_ddb_resource: Any = None + + +def _get_dynamodb_resource() -> Any: + global _ddb_resource + if _ddb_resource is None: + _ddb_resource = boto3.resource("dynamodb") + return _ddb_resource + + +def _table() -> Any: + table_name = os.environ.get("TOOL_EXECUTION_LEDGER_TABLE") + if not table_name: + raise LedgerError( + "TOOL_EXECUTION_LEDGER_TABLE not configured — cannot reserve a " + "tool execution (fail-closed)" + ) + return _get_dynamodb_resource().Table(table_name) + + +def __reset_ledger_client_for_test() -> None: + """Test-only: clear the cached boto3 resource.""" + global _ddb_resource + _ddb_resource = None + + +# --- Result preparation (inline only in PR1; S3 offload is PR2) ------------- + + +def _prepare_result(result: Any) -> tuple[Any, bool, str | None]: + """Return ``(inline_result_json, truncated, marker)`` for a tool result. + + Inline results only in PR1. When the serialized result exceeds + ``max_inline_bytes`` it is NOT stored inline and is NOT truncated into a + partial body (a partial replay would be unfaithful); instead we record a + deterministic content marker (``sha256`` of the canonical JSON) and set + ``truncated=True``. Faithful large-result replay via S3 offload is PR2 — + do NOT build it here. A caller that gets a truncated record must treat the + dedupe hit as "known-completed, body unavailable inline". + """ + try: + serialized = json.dumps(result, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str) + except (TypeError, ValueError) as exc: + raise LedgerError(f"tool result is not JSON-serializable: {exc}") from exc + encoded = serialized.encode("utf-8") + if len(encoded) > max_inline_bytes(): + marker = hashlib.sha256(encoded).hexdigest() + return None, True, marker + return serialized, False, None + + +# --- Reserve / get / finalize ------------------------------------------------ + + +def get(pk: str, sk: str) -> dict[str, Any] | None: + """Read the ledger row for a key, or ``None`` if absent.""" + try: + resp = _table().get_item(Key={PK_ATTR: pk, SK_ATTR: sk}) + except (ClientError, BotoCoreError) as exc: + raise LedgerError(f"ledger get_item failed for {pk!r}/{sk!r}: {exc}") from exc + item = resp.get("Item") + return item if isinstance(item, dict) else None + + +def _reclaim_stale(pk: str, sk: str, *, seen_created_at: Any, now: float) -> bool: + """Reclaim a dead holder's ``in_flight`` row via a conditional CAS. + + Conditional on ``(status = in_flight AND createdAt = :seen)`` so exactly + one reclaimer can win — two concurrent reclaimers cannot both re-reserve + (the second's ``createdAt`` guard fails). Returns True when this caller + won the reclaim (and may now execute), False otherwise. + """ + try: + _table().update_item( + Key={PK_ATTR: pk, SK_ATTR: sk}, + UpdateExpression="SET #s = :inflight, createdAt = :now, updatedAt = :now, ttl = :ttl", + ConditionExpression="#s = :inflight_guard AND createdAt = :seen", + ExpressionAttributeNames={"#s": "status"}, + ExpressionAttributeValues={ + ":inflight": STATUS_IN_FLIGHT, + ":inflight_guard": STATUS_IN_FLIGHT, + ":seen": seen_created_at, + ":now": now, + ":ttl": int(now) + ttl_seconds(), + }, + ) + return True + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + return False + raise LedgerError(f"ledger reclaim failed for {pk!r}/{sk!r}: {exc}") from exc + except BotoCoreError as exc: + raise LedgerError(f"ledger reclaim transport error for {pk!r}/{sk!r}: {exc}") from exc + + +def _reclaim_released(pk: str, sk: str, *, now: float) -> bool: + """Re-reserve a ``released`` row via a conditional CAS on the status. + + ``released -> in_flight`` guarded by ``status = released`` so two racing + re-reservers cannot both win. Returns True when this caller won. + """ + try: + _table().update_item( + Key={PK_ATTR: pk, SK_ATTR: sk}, + UpdateExpression="SET #s = :inflight, createdAt = :now, updatedAt = :now, ttl = :ttl", + ConditionExpression="#s = :released", + ExpressionAttributeNames={"#s": "status"}, + ExpressionAttributeValues={ + ":inflight": STATUS_IN_FLIGHT, + ":released": STATUS_RELEASED, + ":now": now, + ":ttl": int(now) + ttl_seconds(), + }, + ) + return True + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + return False + raise LedgerError(f"ledger re-reserve failed for {pk!r}/{sk!r}: {exc}") from exc + except BotoCoreError as exc: + raise LedgerError(f"ledger re-reserve transport error for {pk!r}/{sk!r}: {exc}") from exc + + +def reserve(pk: str, sk: str, *, tool_name: str, now: float | None = None) -> ReserveResult: + """Attempt to reserve a tool execution with a conditional first-write-wins. + + Returns a :class:`ReserveResult`: + + * ``WON`` — this caller holds the reservation and MUST proceed to execute + then :func:`finalize_success` / :func:`finalize_failure`. + * ``HIT_COMPLETED`` / ``HIT_FAILED`` — a prior terminal record exists; + return it, do NOT execute. + * ``IN_FLIGHT`` — a live concurrent holder; the caller should poll + (:func:`wait_for_terminal`) then raise :class:`RetryableNoExecutionError` + without executing. A *stale* in-flight row is reclaimed here (returns + ``WON`` with ``reclaimed=True``). + """ + now = time.time() if now is None else now + item = { + PK_ATTR: pk, + SK_ATTR: sk, + "status": STATUS_IN_FLIGHT, + "toolName": tool_name, + "createdAt": now, + "updatedAt": now, + # TTL derived from SERVER write-time (not a producer clock), so a + # skewed/malicious producer cannot force early expiry. + "ttl": int(now) + ttl_seconds(), + } + try: + _table().put_item(Item=item, ConditionExpression=f"attribute_not_exists({PK_ATTR})") + return ReserveResult(ReserveOutcome.WON) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") != "ConditionalCheckFailedException": + raise LedgerError(f"ledger reserve failed for {pk!r}/{sk!r}: {exc}") from exc + except BotoCoreError as exc: + raise LedgerError(f"ledger reserve transport error for {pk!r}/{sk!r}: {exc}") from exc + + # Reservation exists — resolve its current status. + row = get(pk, sk) + if row is None: + # Vanished between our failed put and this read (TTL/reclaim race); + # treat as retryable rather than executing unprotected. + return ReserveResult(ReserveOutcome.IN_FLIGHT) + status = row.get("status") + if status == STATUS_COMPLETED: + return ReserveResult(ReserveOutcome.HIT_COMPLETED, row=row) + if status == STATUS_FAILED: + return ReserveResult(ReserveOutcome.HIT_FAILED, row=row) + if status == STATUS_RELEASED: + # A prior attempt released the reservation (provably no side effect); + # re-reserve via conditional CAS so exactly one re-reserver wins. + if _reclaim_released(pk, sk, now=now): + return ReserveResult(ReserveOutcome.WON, reclaimed=True) + return ReserveResult(ReserveOutcome.IN_FLIGHT, row=row) + # in_flight — reclaim if the holder looks dead, else report live. + created_at = row.get("createdAt") + if isinstance(created_at, (int, float)) and (now - float(created_at)) > lease_seconds(): + if _reclaim_stale(pk, sk, seen_created_at=created_at, now=now): + return ReserveResult(ReserveOutcome.WON, reclaimed=True) + return ReserveResult(ReserveOutcome.IN_FLIGHT, row=row) + + +def wait_for_terminal( + pk: str, + sk: str, + *, + timeout: float | None = None, + interval: float | None = None, + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], None] = time.sleep, +) -> dict[str, Any] | None: + """Bounded poll for the in-flight row to reach a terminal status. + + Returns the terminal row (``completed``/``failed``) if it settles within + the timeout, else ``None`` (the caller then raises + :class:`RetryableNoExecutionError` and NEVER executes). ``clock``/``sleep`` + are injectable for deterministic tests. + """ + timeout = _float_env("TOOL_LEDGER_POLL_TIMEOUT_SECONDS", DEFAULT_POLL_TIMEOUT_SECONDS) if timeout is None else timeout + interval = _float_env("TOOL_LEDGER_POLL_INTERVAL_SECONDS", DEFAULT_POLL_INTERVAL_SECONDS) if interval is None else interval + deadline = clock() + timeout + while True: + row = get(pk, sk) + if row is not None and row.get("status") in (STATUS_COMPLETED, STATUS_FAILED): + return row + if clock() >= deadline: + return None + sleep(interval) + + +def _finalize(pk: str, sk: str, *, attributes: dict[str, Any], now: float) -> None: + """Transition an in-flight row to a terminal state (guarded).""" + names = {"#s": "status"} + values: dict[str, Any] = {":inflight": STATUS_IN_FLIGHT, ":now": now} + set_parts = ["#s = :status", "updatedAt = :now"] + for i, (key, value) in enumerate(attributes.items()): + placeholder = f":v{i}" + name_placeholder = f"#a{i}" + names[name_placeholder] = key + values[placeholder] = value + if key == "status": + values[":status"] = value + else: + set_parts.append(f"{name_placeholder} = {placeholder}") + if ":status" not in values: + raise LedgerError("_finalize requires a 'status' attribute") + try: + _table().update_item( + Key={PK_ATTR: pk, SK_ATTR: sk}, + UpdateExpression="SET " + ", ".join(set_parts), + ConditionExpression="#s = :inflight", + ExpressionAttributeNames=names, + ExpressionAttributeValues=values, + ) + except ClientError as exc: + code = exc.response.get("Error", {}).get("Code") + if code == "ConditionalCheckFailedException": + # Already finalized/reclaimed by another actor — benign; the + # terminal record stands. Do not swallow silently: log it. + logger.info( + "tool-ledger finalize no-op (row not in_flight) pk=%s sk=%s", pk, sk + ) + return + raise LedgerError(f"ledger finalize failed for {pk!r}/{sk!r}: {exc}") from exc + except BotoCoreError as exc: + raise LedgerError(f"ledger finalize transport error for {pk!r}/{sk!r}: {exc}") from exc + + +def finalize_success(pk: str, sk: str, *, result: Any, now: float | None = None) -> None: + """Record a successful execution result (in_flight -> completed).""" + now = time.time() if now is None else now + inline, truncated, marker = _prepare_result(result) + attrs: dict[str, Any] = {"status": STATUS_COMPLETED, "resultTruncated": truncated} + if truncated: + attrs["resultMarker"] = marker + else: + attrs["result"] = inline + _finalize(pk, sk, attributes=attrs, now=now) + + +def finalize_failure( + pk: str, + sk: str, + *, + error_type: str, + retryable: bool, + outcome_indeterminate: bool = False, + now: float | None = None, +) -> None: + """Record a terminal failure (in_flight -> failed).""" + now = time.time() if now is None else now + _finalize( + pk, + sk, + attributes={ + "status": STATUS_FAILED, + "errorType": error_type, + "retryable": retryable, + "outcomeIndeterminate": outcome_indeterminate, + }, + now=now, + ) + + +def release(pk: str, sk: str) -> None: + """Release a reservation whose side effect provably did NOT happen. + + Transitions ``in_flight -> released`` (NOT a delete — the worker IAM grant + is Put/Get/Update only, no ``dynamodb:DeleteItem``, per least privilege). + A released row is re-reservable by the next attempt via a conditional CAS + in :func:`reserve` (retryable-no-side-effect branch of the failure + matrix). Guarded by ``status = in_flight`` so a completed/failed record is + never clobbered. + """ + now = time.time() + try: + _table().update_item( + Key={PK_ATTR: pk, SK_ATTR: sk}, + UpdateExpression="SET #s = :released, updatedAt = :now", + ConditionExpression="#s = :inflight", + ExpressionAttributeNames={"#s": "status"}, + ExpressionAttributeValues={ + ":released": STATUS_RELEASED, + ":inflight": STATUS_IN_FLIGHT, + ":now": now, + }, + ) + except ClientError as exc: + if exc.response.get("Error", {}).get("Code") == "ConditionalCheckFailedException": + return + raise LedgerError(f"ledger release failed for {pk!r}/{sk!r}: {exc}") from exc + except BotoCoreError as exc: + raise LedgerError(f"ledger release transport error for {pk!r}/{sk!r}: {exc}") from exc + + +def _recorded_result(row: dict[str, Any]) -> Any: + """Extract the recorded result from a completed row (parse inline JSON).""" + if row.get("resultTruncated"): + return { + "status": "success", + "idempotent": True, + "resultTruncated": True, + "resultMarker": row.get("resultMarker"), + "note": "result body exceeded inline cap; faithful replay via S3 is PR2", + } + raw = row.get("result") + if isinstance(raw, str): + try: + return json.loads(raw) + except ValueError: + return raw + return raw + + +def execute_idempotent( + *, + pk: str, + sk: str, + tool_name: str, + mode: str, + run_tool: Callable[[], Any], + now: float | None = None, + wait_kwargs: dict[str, Any] | None = None, +) -> Any: + """Coordinate reserve -> execute -> finalize for one tool call. + + This is the SINGLE atomic seam (as a pure abstraction): reserve, execute, + and finalize live inside one call with no external pre/post window. The + Strands hook wraps a tool so its ``.stream()`` delegates here. + + ``mode == 'bypass'`` skips the ledger entirely (read-only tool) — no row + is written and ``run_tool`` runs directly. Any other mode is + ledger-protected (fail-safe default lives in + ``tool_idempotency.classify_idempotency_mode``). + """ + if mode == MODE_BYPASS: + return run_tool() + + now = time.time() if now is None else now + reservation = reserve(pk, sk, tool_name=tool_name, now=now) + + if reservation.outcome == ReserveOutcome.HIT_COMPLETED: + return _recorded_result(reservation.row or {}) + if reservation.outcome == ReserveOutcome.HIT_FAILED: + row = reservation.row or {} + raise RecordedToolFailure( + f"tool {tool_name!r} previously failed terminally " + f"(errorType={row.get('errorType')!r}, " + f"outcomeIndeterminate={row.get('outcomeIndeterminate')})", + recorded=row, + ) + if reservation.outcome == ReserveOutcome.IN_FLIGHT: + settled = wait_for_terminal(pk, sk, **(wait_kwargs or {})) + if settled is not None and settled.get("status") == STATUS_COMPLETED: + return _recorded_result(settled) + if settled is not None and settled.get("status") == STATUS_FAILED: + raise RecordedToolFailure( + f"tool {tool_name!r} failed terminally on the winning attempt", + recorded=settled, + ) + # Holder still in-flight after the bounded poll — retryable, and we + # NEVER executed the side effect (the load-bearing invariant). + raise RetryableNoExecutionError( + f"tool {tool_name!r} reservation held by a concurrent execution; " + "retry without executing" + ) + + # WON (fresh or reclaimed) — execute under the reservation. + try: + result = run_tool() + except ToolOutcomeError as exc: + if exc.side_effect == "not_sent" and exc.retryable: + release(pk, sk) + raise RetryableNoExecutionError( + f"tool {tool_name!r} failed before sending; reservation released" + ) from exc + if exc.side_effect == "applied": + finalize_failure(pk, sk, error_type=exc.error_type, retryable=False, now=now) + raise RecordedToolFailure( + f"tool {tool_name!r} terminal failure: {exc}" + ) from exc + # Unknown outcome (incl. explicit side_effect='unknown') — fail safe. + finalize_failure( + pk, sk, error_type=exc.error_type, retryable=False, + outcome_indeterminate=True, now=now, + ) + raise OutcomeIndeterminateError( + f"tool {tool_name!r} outcome indeterminate; NOT re-executed" + ) from exc + except Exception as exc: # noqa: BLE001 — any unclassified error == unknown outcome + # Fail-safe: an un-tokened side-effecting call whose outcome we cannot + # determine is NEVER re-executed. Surface, never swallow. + finalize_failure( + pk, sk, error_type=type(exc).__name__, retryable=False, + outcome_indeterminate=True, now=now, + ) + raise OutcomeIndeterminateError( + f"tool {tool_name!r} raised {type(exc).__name__}; outcome " + "indeterminate, NOT re-executed" + ) from exc + + # A tool that returns a status='error' ToolResult is a terminal, known + # failure (the adapter returned rather than raised — the outcome is known). + if isinstance(result, dict) and result.get("status") == "error": + finalize_failure(pk, sk, error_type="tool_error_result", retryable=False, now=now) + return result + + finalize_success(pk, sk, result=result, now=now) + return result diff --git a/arbiter/workerWrapper/__tests__/test_tool_idempotency.py b/arbiter/workerWrapper/__tests__/test_tool_idempotency.py new file mode 100644 index 00000000..01d6415c --- /dev/null +++ b/arbiter/workerWrapper/__tests__/test_tool_idempotency.py @@ -0,0 +1,268 @@ +"""Tests for arbiter/workerWrapper/tool_idempotency.py (PR1). + +Covers canonicalization (incl. the two flagged determinism traps: non-string +dict keys and integral-float/-0.0 collapse), key derivation, org-scoping in +the partition key, bypass classification (fail-safe default), and the +bypass-misflag guard (strict-mode block). +""" +from __future__ import annotations + +import os +import sys + +import pytest +from hypothesis import given, settings, strategies as st + +_PROJECT_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..") +) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from arbiter.workerWrapper.tool_idempotency import ( # noqa: E402 + BypassMisflagError, + CanonicalizationError, + MODE_BYPASS, + MODE_LEDGER, + args_hash, + build_key, + build_partition_key, + canonicalize, + check_bypass_classification, + classify_idempotency_mode, + detect_write_verbs, +) + + +# --------------------------------------------------------------------------- +# Canonicalization — determinism basics +# --------------------------------------------------------------------------- + + +class TestCanonicalizationDeterminism: + def test_key_order_permutation_hashes_equal(self): + a = {"b": 1, "a": 2, "c": {"y": 1, "x": 2}} + b = {"c": {"x": 2, "y": 1}, "a": 2, "b": 1} + assert args_hash(a) == args_hash(b) + + def test_naive_json_dumps_differs_on_shuffled_keys_but_canonical_does_not(self): + # Differential: unsorted json.dumps produces different text on + # reordered keys; our canonicalizer collapses them. + import json + + a = {"b": 1, "a": 2} + b = {"a": 2, "b": 1} + assert json.dumps(a) != json.dumps(b) # RED for naive approach + assert canonicalize(a) == canonicalize(b) # GREEN for canonical + + def test_string_key_distinct_from_int_value(self): + # "1" (string) must NOT collapse into 1 (number). + assert args_hash({"k": "1"}) != args_hash({"k": 1}) + + def test_null_is_not_equal_to_missing(self): + assert args_hash({"a": None}) != args_hash({}) + + def test_nested_lists_stable(self): + assert args_hash({"xs": [1, 2, 3]}) == args_hash({"xs": [1, 2, 3]}) + assert args_hash({"xs": [1, 2, 3]}) != args_hash({"xs": [3, 2, 1]}) + + +# --------------------------------------------------------------------------- +# Flagged trap #1 — non-string dict keys reject deterministically +# --------------------------------------------------------------------------- + + +class TestNonStringKeys: + def test_int_key_rejected(self): + with pytest.raises(CanonicalizationError): + canonicalize({1: "a"}) + + def test_mixed_str_int_keys_rejected(self): + # json.dumps(sort_keys=True) would raise TypeError on this; we reject + # with a clear, deterministic CanonicalizationError instead. + with pytest.raises(CanonicalizationError): + canonicalize({"a": 1, 2: "b"}) + + def test_nested_non_string_key_rejected(self): + with pytest.raises(CanonicalizationError): + canonicalize({"outer": {None: "x"}}) + + def test_rejection_is_deterministic(self): + # Same pathological input always raises (never sometimes-collides). + for _ in range(5): + with pytest.raises(CanonicalizationError): + canonicalize({True: 1}) # bool key is not a str + + +# --------------------------------------------------------------------------- +# Flagged trap #2 — integral-float / -0.0 semantic collapse +# --------------------------------------------------------------------------- + + +class TestNumberNormalization: + def test_integral_float_collapses_to_int(self): + assert args_hash({"n": 2.0}) == args_hash({"n": 2}) + + def test_exponent_integral_float_collapses(self): + assert args_hash({"n": 1e0}) == args_hash({"n": 1}) + + def test_negative_zero_collapses_to_zero(self): + assert args_hash({"n": -0.0}) == args_hash({"n": 0}) + assert args_hash({"n": -0.0}) == args_hash({"n": 0.0}) + + def test_non_integral_float_preserved(self): + assert args_hash({"n": 2.5}) != args_hash({"n": 2}) + assert canonicalize({"n": 2.5}) == canonicalize({"n": 2.5}) + + def test_nan_rejected(self): + with pytest.raises(CanonicalizationError): + canonicalize({"n": float("nan")}) + + def test_inf_rejected(self): + with pytest.raises(CanonicalizationError): + canonicalize({"n": float("inf")}) + with pytest.raises(CanonicalizationError): + canonicalize({"n": float("-inf")}) + + def test_bool_not_collapsed_to_int(self): + # bool is an int subclass; True must stay a boolean, not become 1. + assert args_hash({"b": True}) != args_hash({"b": 1}) + assert canonicalize({"b": True}) == canonicalize({"b": True}) + + +# --------------------------------------------------------------------------- +# Property tests +# --------------------------------------------------------------------------- + +_json_scalars = st.one_of( + st.none(), + st.booleans(), + st.integers(min_value=-(10**12), max_value=10**12), + st.text(max_size=20), +) +_json_values = st.recursive( + _json_scalars, + lambda children: st.one_of( + st.lists(children, max_size=5), + st.dictionaries(st.text(min_size=1, max_size=8), children, max_size=5), + ), + max_leaves=25, +) + + +class TestCanonicalizationProperties: + @settings(max_examples=200, deadline=None) + @given(_json_values) + def test_canonicalize_is_stable(self, value): + # Property: canonicalize is a pure function of value — same input, + # same output, and the same key never yields two different hashes. + assert args_hash(value) == args_hash(value) + + @settings(max_examples=200, deadline=None) + @given(st.dictionaries(st.text(min_size=1, max_size=6), _json_scalars, max_size=6)) + def test_dict_key_reordering_is_hash_invariant(self, d): + import random + + items = list(d.items()) + random.shuffle(items) + shuffled = dict(items) + assert args_hash(d) == args_hash(shuffled) + + +# --------------------------------------------------------------------------- +# Key derivation + org scoping +# --------------------------------------------------------------------------- + + +class TestKeyDerivation: + def test_partition_key_is_org_prefixed(self): + assert build_partition_key("orgA", "exec1") == "orgA#exec1" + + def test_same_call_different_org_yields_different_pk(self): + pk_a, sk_a = build_key("orgA", "exec1", "node1", 0, "createTicket", {"x": 1}) + pk_b, sk_b = build_key("orgB", "exec1", "node1", 0, "createTicket", {"x": 1}) + assert pk_a != pk_b # structural cross-org isolation + assert sk_a == sk_b # same logical call -> same SK + + def test_sort_key_composition(self): + _, sk = build_key("o", "e", "node9", 3, "sendEmail", {"to": "a@b.c"}) + assert sk.startswith("node9#3#sendEmail#") + assert len(sk.rsplit("#", 1)[1]) == 64 # sha256 hex + + def test_same_call_index_different_tool_differs(self): + _, sk_a = build_key("o", "e", "n", 0, "toolA", {"x": 1}) + _, sk_b = build_key("o", "e", "n", 0, "toolB", {"x": 1}) + assert sk_a != sk_b # toolName in SK prevents wrongful absorption + + def test_different_args_differ(self): + _, sk_a = build_key("o", "e", "n", 0, "t", {"x": 1}) + _, sk_b = build_key("o", "e", "n", 0, "t", {"x": 2}) + assert sk_a != sk_b + + +# --------------------------------------------------------------------------- +# Bypass classification (fail-safe default = ledger) +# --------------------------------------------------------------------------- + + +class TestBypassClassification: + def test_absent_flag_defaults_to_ledger(self): + assert classify_idempotency_mode({}) == MODE_LEDGER + assert classify_idempotency_mode(None) == MODE_LEDGER + + def test_malformed_idempotency_block_defaults_to_ledger(self): + assert classify_idempotency_mode({"idempotency": "nope"}) == MODE_LEDGER + assert classify_idempotency_mode({"idempotency": {}}) == MODE_LEDGER + + def test_unrecognized_mode_defaults_to_ledger(self): + assert classify_idempotency_mode({"idempotency": {"mode": "weird"}}) == MODE_LEDGER + + def test_explicit_bypass(self): + assert classify_idempotency_mode({"idempotency": {"mode": "bypass"}}) == MODE_BYPASS + + def test_explicit_ledger(self): + assert classify_idempotency_mode({"idempotency": {"mode": "ledger"}}) == MODE_LEDGER + + +# --------------------------------------------------------------------------- +# Bypass misflag guard (blocks in strict mode) +# --------------------------------------------------------------------------- + + +class TestBypassMisflagGuard: + _WRITING_CODE = "def handler(x):\n ddb.put_item(Item=x)\n return 'ok'\n" + _READONLY_CODE = "def handler(x):\n return ddb.get_item(Key=x)\n" + + def test_detect_write_verbs_finds_put_item(self): + assert "put_item" in detect_write_verbs(self._WRITING_CODE) + + def test_detect_write_verbs_clean_on_readonly(self): + assert detect_write_verbs(self._READONLY_CODE) == [] + + def test_ledger_tool_is_never_scrutinized(self): + # A ledger-classified tool with writing code is fine — the guard only + # scrutinizes bypass claims. + assert check_bypass_classification( + {"idempotency": {"mode": "ledger"}}, self._WRITING_CODE, + enforcement_mode="strict", + ) == [] + + def test_misflagged_bypass_blocks_in_strict(self): + with pytest.raises(BypassMisflagError): + check_bypass_classification( + {"idempotency": {"mode": "bypass"}}, self._WRITING_CODE, + enforcement_mode="strict", + ) + + def test_misflagged_bypass_warns_not_blocks_in_shadow(self): + hits = check_bypass_classification( + {"idempotency": {"mode": "bypass"}}, self._WRITING_CODE, + enforcement_mode="shadow", + ) + assert "put_item" in hits # returned for the caller to WARN/record + + def test_clean_bypass_passes_in_strict(self): + assert check_bypass_classification( + {"idempotency": {"mode": "bypass"}}, self._READONLY_CODE, + enforcement_mode="strict", + ) == [] diff --git a/arbiter/workerWrapper/__tests__/test_tool_idempotency_threading.py b/arbiter/workerWrapper/__tests__/test_tool_idempotency_threading.py new file mode 100644 index 00000000..8a9fbdff --- /dev/null +++ b/arbiter/workerWrapper/__tests__/test_tool_idempotency_threading.py @@ -0,0 +1,117 @@ +"""Tests for tool-call idempotency context threading (PR1). + +Covers the three additive env vars in ``build_subprocess_env`` (back-compat +byte-identity when absent), the server-side orgId resolver in ``index.py`` +(never trusts a payload value), and the hook's back-compat no-op contract. +""" +from __future__ import annotations + +import os +import sys +from unittest.mock import MagicMock + +_PROJECT_ROOT = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "..", "..") +) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from worker_governance import build_subprocess_env # noqa: E402 + + +class TestSubprocessEnvThreading: + def test_idempotency_context_set_when_provided(self): + env = build_subprocess_env( + {}, execution_id="exec1", node_id="node1", org_id="orgA" + ) + assert env["CITADEL_EXECUTION_ID"] == "exec1" + assert env["CITADEL_NODE_ID"] == "node1" + assert env["CITADEL_ORG_ID"] == "orgA" + + def test_absent_context_is_byte_identical_backcompat(self): + base = build_subprocess_env({}, agent_id="a", workflow_id="w") + assert "CITADEL_EXECUTION_ID" not in base + assert "CITADEL_NODE_ID" not in base + assert "CITADEL_ORG_ID" not in base + + def test_empty_org_id_omitted_but_exec_node_kept(self): + # An empty orgId is allowed (executionId is globally unique); it is + # simply not written as an env var, so the hook reads '' by default. + env = build_subprocess_env({}, execution_id="e", node_id="n", org_id="") + assert env["CITADEL_EXECUTION_ID"] == "e" + assert env["CITADEL_NODE_ID"] == "n" + assert "CITADEL_ORG_ID" not in env + + +class TestServerSideOrgResolution: + def _index(self): + import index # resolved to workerWrapper/index by conftest + return index + + def test_org_id_read_from_execution_row(self, monkeypatch): + index = self._index() + monkeypatch.setenv("EXECUTIONS_TABLE", "citadel-executions-test") + table = MagicMock() + table.get_item.return_value = {"Item": {"executionId": "e1", "orgId": "org-trusted"}} + ddb = MagicMock() + ddb.Table.return_value = table + monkeypatch.setattr(index, "_get_dynamodb", lambda: ddb) + assert index._resolve_execution_org_id("e1") == "org-trusted" + + def test_falls_back_to_env_when_row_missing_org(self, monkeypatch): + index = self._index() + monkeypatch.setenv("EXECUTIONS_TABLE", "citadel-executions-test") + monkeypatch.setenv("RELEASE_DEFAULT_ORG_ID", "org-default") + table = MagicMock() + table.get_item.return_value = {"Item": {"executionId": "e1"}} # no orgId + ddb = MagicMock() + ddb.Table.return_value = table + monkeypatch.setattr(index, "_get_dynamodb", lambda: ddb) + assert index._resolve_execution_org_id("e1") == "org-default" + + def test_returns_empty_when_no_table_and_no_env(self, monkeypatch): + index = self._index() + monkeypatch.delenv("EXECUTIONS_TABLE", raising=False) + monkeypatch.delenv("RELEASE_DEFAULT_ORG_ID", raising=False) + assert index._resolve_execution_org_id("e1") == "" + + def test_read_failure_is_non_fatal(self, monkeypatch): + index = self._index() + monkeypatch.setenv("EXECUTIONS_TABLE", "citadel-executions-test") + monkeypatch.delenv("RELEASE_DEFAULT_ORG_ID", raising=False) + ddb = MagicMock() + ddb.Table.side_effect = RuntimeError("ddb down") + monkeypatch.setattr(index, "_get_dynamodb", lambda: ddb) + assert index._resolve_execution_org_id("e1") == "" # never raises + + +class TestHookBackCompat: + def test_hook_disabled_without_execution_node(self): + from tool_idempotency_hook import IdempotencyToolHook + + assert IdempotencyToolHook(org_id="o", execution_id="", node_id="").enabled is False + assert IdempotencyToolHook(org_id="o", execution_id="e", node_id="").enabled is False + + def test_hook_enabled_with_execution_and_node(self): + from tool_idempotency_hook import IdempotencyToolHook + + assert IdempotencyToolHook(org_id="", execution_id="e", node_id="n").enabled is True + + def test_register_hooks_noop_when_disabled(self): + from tool_idempotency_hook import IdempotencyToolHook + + registry = MagicMock() + IdempotencyToolHook(org_id="o", execution_id="", node_id="").register_hooks(registry) + registry.add_callback.assert_not_called() + + def test_mode_resolver_failure_defaults_to_ledger(self): + from tool_idempotency import MODE_LEDGER + from tool_idempotency_hook import IdempotencyToolHook + + def boom(_name): + raise ValueError("resolver blew up") + + hook = IdempotencyToolHook( + org_id="o", execution_id="e", node_id="n", mode_resolver=boom + ) + assert hook._resolve_mode("anyTool") == MODE_LEDGER # fail-safe diff --git a/arbiter/workerWrapper/__tests__/test_workflow_node_configuration.py b/arbiter/workerWrapper/__tests__/test_workflow_node_configuration.py index 673a4d5f..ad0094dc 100644 --- a/arbiter/workerWrapper/__tests__/test_workflow_node_configuration.py +++ b/arbiter/workerWrapper/__tests__/test_workflow_node_configuration.py @@ -169,10 +169,16 @@ def test_empty_configuration_builds_todays_exact_env(self): mock_run, mock_events, agent_cfg = _run_node({}) extra_env = _extract_extra_env(mock_run.call_args) - # Byte-identical to the pre-feature env: governance triplet only. + # Governance triplet PLUS the tool-call idempotency context (PR1): + # the workflow-node path threads executionId/nodeId so the subprocess + # can build ledger keys. orgId is resolved server-side and omitted + # here (no EXECUTIONS_TABLE / RELEASE_DEFAULT_ORG_ID in this test env), + # which is a legal empty-org case (executionId is globally unique). assert extra_env == { 'CITADEL_AGENT_ID': 'agent-A', 'CITADEL_WORKFLOW_ID': 'exec-1', + 'CITADEL_EXECUTION_ID': 'exec-1', + 'CITADEL_NODE_ID': 'n0', } assert agent_cfg['config']['description'] == 'Base agent.' # Node completed normally. diff --git a/arbiter/workerWrapper/agent_runner.py b/arbiter/workerWrapper/agent_runner.py index 648d3d97..ba617eb8 100644 --- a/arbiter/workerWrapper/agent_runner.py +++ b/arbiter/workerWrapper/agent_runner.py @@ -390,6 +390,76 @@ def _governed_init(self, *args, **kwargs): return True +def _install_idempotency_hook(): + """Patch ``strands.Agent.__init__`` to attach an idempotency HookProvider. + + Tool-call idempotency (PR1). Uses the ONLY tool-call extension surface the + pinned ``strands-agents==1.30.0`` actually exposes: the hooks system + (``Agent(hooks=[...])`` + ``BeforeToolCallEvent``). Verified against the + 1.30.0 source — that release has NO ``strands.handlers.tool_handler`` + /``AgentToolHandler`` and ``Agent.__init__`` accepts neither + ``tool_handler`` nor ``**kwargs``; it does accept ``hooks``. We therefore + APPEND an ``IdempotencyToolHook`` to whatever ``hooks`` list the caller + passed (never clobbering caller-supplied hooks). + + No-op unless BOTH ``CITADEL_EXECUTION_ID`` and ``CITADEL_NODE_ID`` are set + (back-compat: an agent run outside the idempotency envelope is unchanged). + ``CITADEL_ORG_ID`` is optional (empty -> shared org prefix; executionId is + still globally unique) and is read ONLY from the trusted subprocess env + that the worker set server-side, never from tool/agent input. + + Graceful degrade (WARN, return False) when strands or the hook module + cannot be imported — a missing idempotency layer must never halt an + otherwise-valid agent. + + Returns True when the patch was installed, False otherwise. + """ + if not (os.environ.get('CITADEL_EXECUTION_ID') and os.environ.get('CITADEL_NODE_ID')): + return False + + try: + import strands # type: ignore[import-not-found] + except ImportError as exc: + sys.stderr.write( + f'[agent_runner] WARN idempotency hook skipped — ' + f'strands unavailable: {exc}\n' + ) + return False + + _here = os.path.dirname(os.path.abspath(__file__)) + if _here not in sys.path: + sys.path.insert(0, _here) + + try: + from tool_idempotency_hook import IdempotencyToolHook + except ImportError as exc: + sys.stderr.write( + f'[agent_runner] WARN idempotency hook skipped — ' + f'tool_idempotency_hook unavailable: {exc}\n' + ) + return False + + original_init = strands.Agent.__init__ + + def _idempotent_init(self, *args, **kwargs): + hook = IdempotencyToolHook( + org_id=os.environ.get('CITADEL_ORG_ID', ''), + execution_id=os.environ.get('CITADEL_EXECUTION_ID', ''), + node_id=os.environ.get('CITADEL_NODE_ID', ''), + ) + existing = kwargs.get('hooks') + if existing is None: + kwargs['hooks'] = [hook] + elif isinstance(existing, list): + kwargs['hooks'] = [*existing, hook] + # If a caller passed a non-list hooks value we leave it untouched — + # strands will validate it; we never overwrite caller intent. + return original_init(self, *args, **kwargs) + + strands.Agent.__init__ = _idempotent_init + return True + + def _install_model_override(): """Patch ``strands.models.BedrockModel.__init__`` to force ``model_id``. @@ -452,6 +522,10 @@ def main(): # subprocess env lacks CITADEL_AGENT_ID (backward compatible). _install_governed_tool_handler() + # Install the tool-call idempotency hook (PR1) via the strands hooks + # system — no-op unless CITADEL_EXECUTION_ID + CITADEL_NODE_ID are set. + _install_idempotency_hook() + # Overrides the model id for operator-selected per-agent overrides; # no-op unless MODEL_OVERRIDE is set in the subprocess environment. _install_model_override() diff --git a/arbiter/workerWrapper/index.py b/arbiter/workerWrapper/index.py index ecb5d1d6..86608396 100644 --- a/arbiter/workerWrapper/index.py +++ b/arbiter/workerWrapper/index.py @@ -813,6 +813,43 @@ def _extract_worker_trace_context(event, message_attributes): pass return extract_carried(event) +def _resolve_execution_org_id(execution_id: str) -> str: + """Resolve an execution's ``orgId`` SERVER-SIDE from the execution row. + + Tool-call idempotency (PR1) org-scoping: ``orgId`` is the ledger PK prefix + and provides structural cross-org isolation, so it MUST come from a + trusted server-side source — the ``EXECUTIONS_TABLE`` row keyed by + ``executionId`` — and NEVER from a subprocess-supplied payload that could + be spoofed to cross orgs. + + Best-effort and non-fatal: returns ``''`` when the table binding is + absent, the row/attribute is missing, or the read fails. An empty orgId is + safe — ``executionId`` is globally unique, so the ledger key stays unique; + the org prefix is defense-in-depth, not the uniqueness guarantee. Falls + back to ``RELEASE_DEFAULT_ORG_ID`` (the same trusted env the release path + uses) before ``''``. Never raises — org resolution must not fail a node. + """ + table_name = os.environ.get('EXECUTIONS_TABLE') + if table_name and execution_id: + try: + resp = _get_dynamodb().Table(table_name).get_item( + Key={'executionId': execution_id} + ) + org_id = (resp.get('Item') or {}).get('orgId') + if isinstance(org_id, str) and org_id: + return org_id + except Exception as exc: # noqa: BLE001 — org resolution is best-effort + print(json.dumps({ + 'level': 'WARN', + 'component': 'WorkerWrapper', + 'action': 'idempotency_org_resolve_failed', + 'executionId': execution_id, + 'error': str(exc), + })) + fallback = os.environ.get('RELEASE_DEFAULT_ORG_ID') + return fallback if isinstance(fallback, str) and fallback else '' + + def _process_workflow_node(event, message_attributes=None): """Run the agent for a dispatched workflow node and emit its result. @@ -890,6 +927,14 @@ def _process_workflow_node(event, message_attributes=None): model_override=model_override, agent_id=msg.agent_id, workflow_id=msg.execution_id, + # Tool-call idempotency (PR1) context. orgId is resolved + # SERVER-SIDE from the execution row (never trusted from the + # dispatch payload); executionId/nodeId come from the validated + # node-dispatch message. When these are threaded, agent_runner + # installs the idempotency hook in the subprocess. + execution_id=msg.execution_id, + node_id=msg.node_id, + org_id=_resolve_execution_org_id(msg.execution_id), ) usage_sink: list = [] diff --git a/arbiter/workerWrapper/tool_idempotency.py b/arbiter/workerWrapper/tool_idempotency.py new file mode 100644 index 00000000..07e437ca --- /dev/null +++ b/arbiter/workerWrapper/tool_idempotency.py @@ -0,0 +1,291 @@ +"""Pure canonical args-hash + idempotency-key derivation (PR1 of 2). + +This module is the deterministic, I/O-free core of tool-call idempotency. It +produces a stable ``argsHash`` from a tool's input and derives the ledger key +``(orgId#executionId, nodeId#callIndex#toolName#argsHash)`` used by +``arbiter/governance/tool_execution_ledger.py``. + +Guarantee scope (read this precisely — the honest framing per the security +consensus, do NOT collapse it to a bare "exactly-once"): + +* The key is **attempt-scoped**: ``callIndex`` is a per-handler-instance + monotonic counter (one handler instance == one subprocess == one node + attempt) and ``argsHash`` derives from the exact tool input the model + produced on *this* attempt. Two byte-identical (post-canonicalization) + re-issues of the same logical call within one attempt collapse to the same + key → exactly-once within an attempt, plus reservation-race safety. +* It does **NOT** provide exactly-once across nondeterministic re-dispatch + (a watchdog re-dispatch runs a fresh LLM body whose calls may reorder or + reword → different keys). Closing that requires the worker + ``dispatchGeneration`` fence, which is **deferred to PR2** and is REQUIRED + for the complete guarantee. Nothing here should be read as the complete + guarantee. + +Why not just ``json.dumps(sort_keys=True)``? Two flagged determinism traps +that raw ``dumps`` gets wrong (both are property-tested): + +1. **Non-string dict keys.** ``json.dumps(..., sort_keys=True)`` raises + ``TypeError`` on mixed-type keys (``'<' not supported between 'str' and + 'int'``) and silently coerces ``int`` keys to strings otherwise — + producing ``{1: 2}`` and ``{"1": 2}`` as the *same* text (a collision). + We reject non-string keys deterministically (raise + :class:`CanonicalizationError`) rather than silently collide; model- + produced tool JSON always has string keys, so this fails loudly only on + the pathological case. +2. **Integral-float / -0.0 collapse.** ``1`` vs ``1.0`` vs ``1e0`` dump to + different text, and ``-0.0`` vs ``0.0`` differ — so semantically-equal + numbers would hash differently. We normalize integral floats to ``int`` + (``2.0 -> 2``, ``-0.0 -> 0``) so they hash equal, and reject non-finite + (``NaN``/``Infinity``). ``null`` is preserved (``{"a": null}`` never + equals ``{}``), unlike the ledger *serializer* which strips ``None`` for + storage — that is a storage concern, not a hash-input concern. +""" + +from __future__ import annotations + +import hashlib +import json +import math +from typing import Any + +__all__ = [ + "CanonicalizationError", + "BypassMisflagError", + "canonicalize", + "args_hash", + "build_partition_key", + "build_sort_key", + "build_key", + "classify_idempotency_mode", + "detect_write_verbs", + "check_bypass_classification", + "MODE_LEDGER", + "MODE_BYPASS", +] + + +class CanonicalizationError(ValueError): + """Raised when a tool input cannot be canonicalized deterministically. + + The two deterministic-failure cases are a non-string dict key and a + non-finite number (``NaN``/``Infinity``). Callers MUST treat this as a + fail-closed condition for a side-effecting tool: a call whose key cannot + be derived cannot be deduplicated, so it must not be executed + unprotected. + """ + + +def _normalize(obj: Any) -> Any: + """Recursively normalize a JSON-ish value to a canonical form. + + * ``dict`` — keys MUST be strings (reject otherwise); values normalized + recursively; ``None`` values preserved (``null`` != missing). + * ``float`` — reject ``NaN``/``Infinity``; integral floats (incl. + ``-0.0``) collapse to ``int`` so ``2.0`` and ``2`` hash equal. + * ``bool`` — preserved as-is (checked before ``int`` since ``bool`` is an + ``int`` subclass; ``True`` must not become ``1``). + * ``int`` / ``str`` / ``None`` — as-is. + * ``list`` / ``tuple`` — normalized element-wise (tuple -> list). + * anything else — coerced to ``str`` deterministically (defensive; model + tool input is plain JSON, so this is a rare fallback). + """ + if obj is None or isinstance(obj, str): + return obj + if isinstance(obj, bool): # bool BEFORE int (bool is a subclass of int) + return obj + if isinstance(obj, int): + return obj + if isinstance(obj, float): + if math.isnan(obj) or math.isinf(obj): + raise CanonicalizationError( + f"non-finite number is not canonicalizable: {obj!r}" + ) + # Integral floats (including -0.0) collapse to int so 2.0 == 2 and + # -0.0 == 0 at the hash layer. is_integer() is True for -0.0. + if obj.is_integer(): + return int(obj) + return obj + if isinstance(obj, dict): + out: dict[str, Any] = {} + for key, value in obj.items(): + if not isinstance(key, str): + raise CanonicalizationError( + f"non-string dict key is not canonicalizable: {key!r} " + f"(type {type(key).__name__}); coerce keys to str upstream" + ) + out[key] = _normalize(value) + return out + if isinstance(obj, (list, tuple)): + return [_normalize(v) for v in obj] + # Defensive coercion for a non-JSON scalar (Decimal, datetime, set, ...). + return str(obj) + + +def canonicalize(tool_input: Any) -> str: + """Return the deterministic canonical JSON string for a tool input. + + Pure and total except for the two deterministic rejections documented on + :class:`CanonicalizationError` (non-string dict key, non-finite number). + After normalization every dict key is a string, so ``sort_keys=True`` + cannot raise; integral floats are ints, so ``1``/``1.0`` serialize + identically. + """ + normalized = _normalize(tool_input) + return json.dumps( + normalized, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + + +def args_hash(tool_input: Any) -> str: + """Return the SHA-256 hex digest of the canonicalized tool input.""" + canonical = canonicalize(tool_input) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def build_partition_key(org_id: str, execution_id: str) -> str: + """Ledger partition key: ``orgId#executionId``. + + ``orgId`` is prefixed for structural cross-org isolation — a key minted + for org A can never collide with or read org B's rows. ``orgId`` MUST be + resolved server-side (execution row / trusted env), never from a + subprocess-supplied payload. ``executionId`` alone is globally unique, so + correctness holds even when ``orgId`` is an empty/sentinel string. + """ + return f"{org_id}#{execution_id}" + + +def build_sort_key(node_id: str, call_index: int, tool_name: str, hash_hex: str) -> str: + """Ledger sort key: ``nodeId#callIndex#toolName#argsHash``. + + ``toolName`` and ``argsHash`` are included even though + ``(executionId, nodeId, callIndex)`` is already unique within one attempt: + they keep the key stable and *verifiable* across attempts — a replayed + call with the same ``callIndex`` but a different tool/args yields a + different key and is correctly treated as a different call, never wrongly + absorbed. Dispatch generation is deliberately NOT in the key (that would + mint a fresh key on every re-dispatch and guarantee duplicates — the + opposite of the goal); cross-dispatch closure is PR2's worker fence. + """ + return f"{node_id}#{call_index}#{tool_name}#{hash_hex}" + + +def build_key( + org_id: str, + execution_id: str, + node_id: str, + call_index: int, + tool_name: str, + tool_input: Any, +) -> tuple[str, str]: + """Derive the ``(partitionKey, sortKey)`` ledger key for a tool call.""" + hash_hex = args_hash(tool_input) + return ( + build_partition_key(org_id, execution_id), + build_sort_key(node_id, call_index, tool_name, hash_hex), + ) + + +# --------------------------------------------------------------------------- +# Bypass classification (read-only tools skip the ledger) +# --------------------------------------------------------------------------- + +MODE_LEDGER = "ledger" +MODE_BYPASS = "bypass" + + +def classify_idempotency_mode(tool_config: Any) -> str: + """Resolve a tool's idempotency mode from its per-tool config. + + **Fail-safe default = ``ledger`` (treat as side-effecting).** The + dangerous direction — a side-effecting tool silently unprotected — + requires an *explicit* ``mode: "bypass"`` flag, never a mere omission, + a malformed config, or an unrecognized value. Anything that is not + exactly the string ``"bypass"`` resolves to ``ledger``. + + Config shape (per-tool item, loaded by ``fabricator/tools_config.py``):: + + {"idempotency": {"mode": "ledger" | "bypass", ...}} + """ + if not isinstance(tool_config, dict): + return MODE_LEDGER + idem = tool_config.get("idempotency") + if not isinstance(idem, dict): + return MODE_LEDGER + mode = idem.get("mode") + if mode == MODE_BYPASS: + return MODE_BYPASS + # Absent / None / malformed / any unrecognized value -> fail-safe ledger. + return MODE_LEDGER + + +class BypassMisflagError(Exception): + """Raised when a ``bypass``-flagged tool looks demonstrably side-effecting. + + Security linkage (consensus condition C1): a ``bypass`` tool writes no + ledger row, so it also never reaches the reserve step — meaning a + mis-flagged side-effecting tool loses BOTH dedupe AND (in PR2) the + dispatch-generation fence at once (double jeopardy). A ``warn`` is too + weak for a control whose failure silently removes protection, so in + ``strict`` enforcement mode this BLOCKS (raises); in other modes the + caller should WARN and record a governance signal, never swallow. + """ + + +# Write-verb heuristic: substrings that strongly imply a mutating operation. +# Deliberately broad (over-blocking a bypass mis-flag is the safe direction); +# matched case-insensitively against the tool's source/binding code. This is +# a heuristic, NOT a soundness boundary — a determined author can hide a write +# behind indirection; the real control is the fail-safe default plus this gate. +_WRITE_VERB_TOKENS = ( + "put_item", "update_item", "delete_item", "batch_write", "transact_write", + "put_object", "delete_object", "putcommand", "updatecommand", "deletecommand", + "createticket", "create_ticket", "create(", "update(", "delete(", "insert(", + "post(", "put(", "patch(", "send_message", "sendmessage", "publish(", + "requests.post", "requests.put", "requests.patch", "requests.delete", + ".create_", ".update_", ".delete_", ".write(", "execute(", "commit(", +) + + +def detect_write_verbs(tool_code: Any) -> list[str]: + """Return the write-verb tokens found in ``tool_code`` (case-insensitive). + + Empty list when ``tool_code`` is not a non-empty string or contains no + recognized mutating token. Pure and total — never raises. + """ + if not isinstance(tool_code, str) or not tool_code: + return [] + lowered = tool_code.lower() + return [tok for tok in _WRITE_VERB_TOKENS if tok in lowered] + + +def check_bypass_classification( + tool_config: Any, + tool_code: Any, + *, + enforcement_mode: str, +) -> list[str]: + """Guard a ``bypass`` classification against a demonstrably-writing tool. + + Returns the list of detected write-verb tokens (empty = clean). For a + ``ledger``-classified tool this is always a clean no-op (empty list) — + the guard only scrutinizes ``bypass`` claims. + + In ``strict`` ``enforcement_mode`` a non-empty detection RAISES + :class:`BypassMisflagError` (blocks activation, consensus condition C1). + In any other mode the detection is returned for the caller to WARN and + record — never silently dropped. + """ + if classify_idempotency_mode(tool_config) != MODE_BYPASS: + return [] + hits = detect_write_verbs(tool_code) + if hits and enforcement_mode == "strict": + raise BypassMisflagError( + "tool is flagged idempotency.mode='bypass' but its code contains " + f"write verbs {hits!r}; blocked in strict enforcement mode " + "(a bypass tool skips the ledger and, in PR2, the dispatch fence)" + ) + return hits diff --git a/arbiter/workerWrapper/tool_idempotency_hook.py b/arbiter/workerWrapper/tool_idempotency_hook.py new file mode 100644 index 00000000..59b0c129 --- /dev/null +++ b/arbiter/workerWrapper/tool_idempotency_hook.py @@ -0,0 +1,276 @@ +"""Strands tool-idempotency hook — the single atomic seam (PR1 of 2). + +Verified against ``strands-agents==1.30.0`` (the version pinned in +``arbiter/workerWrapper/requirements.txt``): there is **no** +``strands.handlers.tool_handler.AgentToolHandler`` and ``Agent.__init__`` +accepts neither ``tool_handler`` nor ``**kwargs``. The supported tool-call +extension surface is the **hooks system**: + +* ``ToolExecutor._stream`` fires ``BeforeToolCallEvent`` (whose + ``selected_tool`` is writable), then runs ``selected_tool.stream(...)``, + then fires ``AfterToolCallEvent``. +* A ``HookProvider`` registers callbacks via ``registry.add_callback``, and is + attached with ``Agent(hooks=[...])``. + +To get reserve -> execute -> finalize with **no pre/post window** (the +requirement, and security condition C2), this hook does NOT split logic across +Before/After. Instead, in ``BeforeToolCallEvent`` it REPLACES ``selected_tool`` +with an :class:`_IdempotentToolWrapper` whose ``stream()`` performs +reserve -> delegate-to-real-tool -> finalize inside one coroutine. The reserve +strictly precedes any adapter call, and there is no seam between reserve and +execute that another actor could slip through. + +The idempotency decision logic itself lives in +``arbiter/governance/tool_execution_ledger.py`` (fully unit-tested via the +synchronous ``execute_idempotent`` coordinator with a stubbed adapter + a +conditional-write fake). This module is the thin async glue that binds that +logic to the real Strands seam; it degrades to a no-op import when +``strands`` is unavailable (dev/CI), mirroring ``governed_tool_handler.py``. + +Scope note (honesty requirement): this delivers exactly-once WITHIN an +attempt + reservation-race safety. Exactly-once across nondeterministic +re-dispatch needs the worker ``dispatchGeneration`` fence, which is DEFERRED +to PR2 and REQUIRED for the complete guarantee. +""" + +from __future__ import annotations + +import logging +import os +import sys +from typing import Any, Callable + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_PROJECT_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from arbiter.governance import tool_execution_ledger as ledger # noqa: E402 +from arbiter.workerWrapper.tool_idempotency import ( # noqa: E402 + MODE_LEDGER, + build_key, + classify_idempotency_mode, +) + +logger = logging.getLogger(__name__) + +try: + from strands.hooks import BeforeToolCallEvent # type: ignore + from strands.types.tools import AgentTool # type: ignore + + _STRANDS_AVAILABLE = True +except ImportError: # pragma: no cover — dev/CI without strands-agents + _STRANDS_AVAILABLE = False + BeforeToolCallEvent = object # type: ignore[assignment,misc] + AgentTool = object # type: ignore[assignment,misc] + + +def _error_result(tool_use_id: str, message: str) -> dict[str, Any]: + """A Strands ToolResult-shaped error dict.""" + return { + "toolUseId": tool_use_id, + "status": "error", + "content": [{"text": message}], + } + + +class _IdempotentToolWrapper(AgentTool): # type: ignore[misc] + """Wraps a selected tool so its ``stream()`` is ledger-protected. + + Reserve -> execute -> finalize happen inside this one ``stream()`` + coroutine, mirroring ``ledger.execute_idempotent``'s branches (that sync + coordinator is the unit-tested authority for the invariant; this async + path applies the identical failure matrix over the same primitives). + """ + + def __init__(self, inner: Any, pk: str, sk: str, tool_name: str, mode: str): + try: + super().__init__() + except TypeError: # pragma: no cover — base signature drift + pass + self._inner = inner + self._pk = pk + self._sk = sk + self._tool_name = tool_name + self._mode = mode + + # --- AgentTool interface delegation -------------------------------------- + @property + def tool_name(self) -> str: # pragma: no cover — thin delegation + return self._inner.tool_name + + @property + def tool_spec(self) -> Any: # pragma: no cover — thin delegation + return self._inner.tool_spec + + @property + def tool_type(self) -> str: # pragma: no cover — thin delegation + return self._inner.tool_type + + def get_display_properties(self) -> dict[str, str]: # pragma: no cover + return self._inner.get_display_properties() + + async def stream(self, tool_use: Any, invocation_state: dict[str, Any], **kwargs: Any): # pragma: no cover — requires strands runtime + from strands.types._events import ToolResultEvent # local import; strands-only + + tool_use_id = str(tool_use.get("toolUseId", "")) if hasattr(tool_use, "get") else "" + + reservation = ledger.reserve(self._pk, self._sk, tool_name=self._tool_name) + + if reservation.outcome == ledger.ReserveOutcome.HIT_COMPLETED: + yield ToolResultEvent(ledger._recorded_result(reservation.row or {})) + return + if reservation.outcome == ledger.ReserveOutcome.HIT_FAILED: + yield ToolResultEvent(_error_result(tool_use_id, "prior terminal failure (idempotent replay)")) + return + if reservation.outcome == ledger.ReserveOutcome.IN_FLIGHT: + settled = ledger.wait_for_terminal(self._pk, self._sk) + if settled is not None and settled.get("status") == ledger.STATUS_COMPLETED: + yield ToolResultEvent(ledger._recorded_result(settled)) + return + # Concurrent loser (or terminal-failed winner): retryable, no + # execution. Surface as an error ToolResult — we NEVER ran the tool. + yield ToolResultEvent(_error_result( + tool_use_id, + "tool execution reserved by a concurrent run; retry (no side effect performed)", + )) + return + + # WON — execute the real tool under the reservation. + last_result: Any = None + try: + async for event in self._inner.stream(tool_use, invocation_state, **kwargs): + if isinstance(event, ToolResultEvent): + last_result = event.tool_result + yield event + except Exception as exc: # noqa: BLE001 — unknown outcome, fail safe + ledger.finalize_failure( + self._pk, self._sk, error_type=type(exc).__name__, + retryable=False, outcome_indeterminate=True, + ) + raise + + if isinstance(last_result, dict) and last_result.get("status") == "error": + ledger.finalize_failure(self._pk, self._sk, error_type="tool_error_result", retryable=False) + elif last_result is not None: + ledger.finalize_success(self._pk, self._sk, result=last_result) + + +class IdempotencyToolHook: + """A Strands ``HookProvider`` that installs ledger-backed idempotency. + + Attach with ``Agent(hooks=[IdempotencyToolHook(...)])``. A no-op unless + ``execution_id`` and ``node_id`` are present (back-compat: an agent run + outside the idempotency envelope behaves byte-identically to today). + + ``org_id`` MUST be resolved server-side (execution row / trusted env) by + the caller — never taken from a subprocess-supplied payload. ``call_index`` + is a per-instance monotonic counter (one hook instance == one agent == one + subprocess == one node attempt), so it is attempt-scoped by construction. + + ``mode_resolver`` maps a tool name to ``'ledger'``/``'bypass'`` (fail-safe + default ``'ledger'`` via ``classify_idempotency_mode``); when absent every + tool is ledger-protected. + """ + + def __init__( + self, + *, + org_id: str, + execution_id: str, + node_id: str, + mode_resolver: Callable[[str], str] | None = None, + ): + self._org_id = org_id or "" + self._execution_id = execution_id or "" + self._node_id = node_id or "" + self._mode_resolver = mode_resolver + self._call_index = 0 + + @property + def enabled(self) -> bool: + return bool(self._execution_id and self._node_id) + + def register_hooks(self, registry: Any, **_kwargs: Any) -> None: + if not _STRANDS_AVAILABLE: # pragma: no cover — dev/CI guard + logger.warning("idempotency hook skipped — strands unavailable") + return + if not self.enabled: + logger.warning( + "idempotency hook skipped — execution/node context absent " + "(back-compat no-op)" + ) + return + registry.add_callback(BeforeToolCallEvent, self._on_before_tool_call) + + def _resolve_mode(self, tool_name: str) -> str: + if self._mode_resolver is None: + return MODE_LEDGER + try: + return classify_idempotency_mode({"idempotency": {"mode": self._mode_resolver(tool_name)}}) + except Exception: # noqa: BLE001 — resolver failure must fail safe + return MODE_LEDGER + + def _on_before_tool_call(self, event: Any) -> None: # pragma: no cover — requires strands runtime + tool_use = getattr(event, "tool_use", None) + selected = getattr(event, "selected_tool", None) + if tool_use is None or selected is None: + return + tool_name = tool_use.get("name", "") if hasattr(tool_use, "get") else "" + tool_input = tool_use.get("input", {}) if hasattr(tool_use, "get") else {} + + # Attempt-scoped monotonic index: increment for EVERY intercepted call + # (ledger and bypass alike) so numbering is stable within the attempt. + call_index = self._call_index + self._call_index += 1 + + mode = self._resolve_mode(tool_name) + if mode != MODE_LEDGER: + return # bypass: leave the real tool in place, write no ledger row + + try: + pk, sk = build_key( + self._org_id, self._execution_id, self._node_id, + call_index, tool_name, tool_input, + ) + except Exception: # noqa: BLE001 — canonicalization failure: fail closed + # A side-effecting call whose key cannot be derived must not run + # unprotected. Replace the tool with one that errors deterministically. + event.selected_tool = _KeyDerivationFailedTool(selected) + return + + event.selected_tool = _IdempotentToolWrapper(selected, pk, sk, tool_name, mode) + + +class _KeyDerivationFailedTool(AgentTool): # type: ignore[misc] + """Replacement tool that refuses to run when the idempotency key cannot be + derived (fail-closed for a side-effecting call).""" + + def __init__(self, inner: Any): + try: + super().__init__() + except TypeError: # pragma: no cover + pass + self._inner = inner + + @property + def tool_name(self) -> str: # pragma: no cover + return self._inner.tool_name + + @property + def tool_spec(self) -> Any: # pragma: no cover + return self._inner.tool_spec + + @property + def tool_type(self) -> str: # pragma: no cover + return self._inner.tool_type + + async def stream(self, tool_use: Any, invocation_state: dict[str, Any], **kwargs: Any): # pragma: no cover + from strands.types._events import ToolResultEvent + + tool_use_id = str(tool_use.get("toolUseId", "")) if hasattr(tool_use, "get") else "" + yield ToolResultEvent(_error_result( + tool_use_id, + "idempotency key could not be derived from tool input; refused " + "(fail-closed, no side effect performed)", + )) diff --git a/arbiter/workerWrapper/worker_governance.py b/arbiter/workerWrapper/worker_governance.py index 732aaa12..2bdeab7c 100644 --- a/arbiter/workerWrapper/worker_governance.py +++ b/arbiter/workerWrapper/worker_governance.py @@ -172,6 +172,9 @@ def build_subprocess_env( workflow_id: str | None = None, denied_tools: list[str] | None = None, eval_run_id: str | None = None, + execution_id: str | None = None, + node_id: str | None = None, + org_id: str | None = None, ) -> dict: """Build the subprocess environment with governance and config overrides. @@ -263,4 +266,22 @@ def build_subprocess_env( if isinstance(eval_run_id, str) and eval_run_id: env['CITADEL_EVAL_RUN_ID'] = eval_run_id + # Tool-call idempotency (PR1): the execution/node/org context threaded to + # the worker subprocess, consumed by + # agent_runner._install_idempotency_hook to build the ledger key + # (orgId#executionId, nodeId#...). All three are additive and optional: + # when any is absent the hook is a back-compat no-op (idempotency + # disabled, pure pre-feature behavior preserved). ``org_id`` MUST be + # resolved server-side by the caller (execution row / trusted env) — this + # function only serializes whatever trusted value it is given, and NEVER + # reads it from a subprocess-supplied payload. An empty org_id is allowed + # (executionId is globally unique, so the key stays unique; the org prefix + # is defense-in-depth cross-org isolation). + if isinstance(execution_id, str) and execution_id: + env['CITADEL_EXECUTION_ID'] = execution_id + if isinstance(node_id, str) and node_id: + env['CITADEL_NODE_ID'] = node_id + if isinstance(org_id, str) and org_id: + env['CITADEL_ORG_ID'] = org_id + return env diff --git a/backend/lib/arbiter-stack.ts b/backend/lib/arbiter-stack.ts index 49806629..51923dff 100644 --- a/backend/lib/arbiter-stack.ts +++ b/backend/lib/arbiter-stack.ts @@ -418,6 +418,28 @@ export class ArbiterStack extends cdk.Stack { }, ); + // Tool-call idempotency ledger (PR1). Org-scoped, TTL'd operational + // dedupe table — NOT an audit artifact (distinct from the 90-day + // governance ledger). PK = orgId#executionId, SK = + // nodeId#callIndex#toolName#argsHash. TTL (attribute `ttl`) is 48h, + // derived server-side at write time by the worker; it exists to bound + // storage, not to retain accountability records. Encrypted at rest with + // an AWS-managed KMS key and PITR on. + const toolExecutionLedgerTable = new dynamodb.Table( + this, + "ToolExecutionLedgerTable", + { + tableName: `citadel-tool-execution-ledger-${props.environment}`, + partitionKey: { name: "pk", type: dynamodb.AttributeType.STRING }, + sortKey: { name: "sk", type: dynamodb.AttributeType.STRING }, + billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, + timeToLiveAttribute: "ttl", + encryption: dynamodb.TableEncryption.AWS_MANAGED, + pointInTimeRecoverySpecification: { pointInTimeRecoveryEnabled: true }, + removalPolicy: cdk.RemovalPolicy.DESTROY, + }, + ); + const workerAgentWrapperLambda = new PythonFunction( this, "WorkerAgentWrapper", @@ -443,6 +465,11 @@ export class ArbiterStack extends cdk.Stack { ...(props.executionsTable && { EXECUTIONS_TABLE: props.executionsTable.tableName, }), + // Tool-call idempotency (PR1): the ledger the worker reserves/ + // finalizes tool executions against. Always wired; the worker's + // idempotency hook is itself gated on per-node execution/node + // context, so a missing key context is a back-compat no-op. + TOOL_EXECUTION_LEDGER_TABLE: toolExecutionLedgerTable.tableName, ...(props.registryId && { REGISTRY_ID: props.registryId }), ...(props.registryId && { REGISTRY_ENABLED: "true" }), }, @@ -518,6 +545,24 @@ export class ArbiterStack extends cdk.Stack { ); } + // Tool-call idempotency (PR1): least-privilege grant on the tool-execution + // ledger. The worker reserves (conditional PutItem), reads recorded + // results (GetItem), and finalizes/releases/reclaims (UpdateItem). It is + // DELIBERATELY NOT grantReadWriteData: no dynamodb:DeleteItem (release is a + // status transition, not a delete) and no dynamodb:Scan/Query (all access + // is by exact key). Scoped to this one table ARN. + workerAgentWrapperLambda.addToRolePolicy( + new iam.PolicyStatement({ + effect: iam.Effect.ALLOW, + actions: [ + "dynamodb:PutItem", + "dynamodb:GetItem", + "dynamodb:UpdateItem", + ], + resources: [toolExecutionLedgerTable.tableArn], + }), + ); + // The worker emits best-effort node-level metrics (NodeDurationMs / // NodeFailure) into the Citadel/Workflows namespace after running each // workflow node. PutMetricData has no resource-level scoping; the call is diff --git a/backend/test/arbiter-stack-step-runner.test.ts b/backend/test/arbiter-stack-step-runner.test.ts index aeaa53c5..5c0df628 100644 --- a/backend/test/arbiter-stack-step-runner.test.ts +++ b/backend/test/arbiter-stack-step-runner.test.ts @@ -447,13 +447,37 @@ describe("ArbiterStack — Step Runner Lambda and EventBridge rules (Task 1.6)", return out; } - test("worker has dynamodb:UpdateItem but NOT Put/Delete/BatchWrite on any table", () => { + test("worker DDB writes stay least-privilege: executions is UpdateItem-only, ledger is Put/Get/Update, no Delete/BatchWrite anywhere", () => { const actions = actionsForRole("WorkerAgentWrapper"); expect(actions.has("dynamodb:UpdateItem")).toBe(true); - // Bare grantWriteData would have added these — it must NOT be used. - expect(actions.has("dynamodb:PutItem")).toBe(false); + // Bare grantWriteData would have added these on the executions table — + // it must NOT be used, on the executions table or anywhere else. expect(actions.has("dynamodb:DeleteItem")).toBe(false); expect(actions.has("dynamodb:BatchWriteItem")).toBe(false); + // PutItem is now present, but ONLY in the tool-execution-ledger grant + // (PR1): the sole statement carrying PutItem must be exactly the + // ledger's Put/Get/Update trio (no Delete/Scan/Query, no FGAC-condition + // executions statement leaking Put). Scope the assertion to that + // statement rather than the whole role. + const policies = template.findResources("AWS::IAM::Policy"); + const putStatements: any[] = []; + for (const p of Object.values(policies) as any[]) { + const roles = p.Properties?.Roles || []; + if ( + !roles.some((r: any) => (r?.Ref || "").includes("WorkerAgentWrapper")) + ) + continue; + for (const s of p.Properties?.PolicyDocument?.Statement || []) { + const acts = Array.isArray(s.Action) ? s.Action : [s.Action]; + if (acts.includes("dynamodb:PutItem")) putStatements.push(acts); + } + } + expect(putStatements.length).toBe(1); + expect(putStatements[0]).toEqual([ + "dynamodb:PutItem", + "dynamodb:GetItem", + "dynamodb:UpdateItem", + ]); }); test("the UpdateItem grant is FGAC-restricted to the nodeResults/executionId attributes", () => { diff --git a/backend/test/arbiter-stack-tool-execution-ledger.test.ts b/backend/test/arbiter-stack-tool-execution-ledger.test.ts new file mode 100644 index 00000000..34d8529f --- /dev/null +++ b/backend/test/arbiter-stack-tool-execution-ledger.test.ts @@ -0,0 +1,167 @@ +import * as cdk from "aws-cdk-lib"; +import { Template, Match } from "aws-cdk-lib/assertions"; +import * as dynamodb from "aws-cdk-lib/aws-dynamodb"; +import * as events from "aws-cdk-lib/aws-events"; +import * as lambda from "aws-cdk-lib/aws-lambda"; +import * as appsync from "aws-cdk-lib/aws-appsync"; +import { Bucket } from "aws-cdk-lib/aws-s3"; +import * as path from "path"; +import { + scaffoldBackendAssetDirs, + scaffoldArbiterStubs, +} from "./helpers/scaffold-stub-assets"; + +scaffoldBackendAssetDirs(["dist/lambda", "src/schema"]); +scaffoldArbiterStubs(); + +import { ArbiterStack } from "../lib/arbiter-stack"; + +// Tool-call idempotency (PR1): the org-scoped TTL'd tool-execution ledger, +// its worker env wiring, and the least-privilege (Put/Get/Update, NO +// Delete/Scan) worker grant scoped to the one table. +describe("ArbiterStack — tool-execution idempotency ledger (PR1)", () => { + let template: Template; + + beforeAll(() => { + const app = new cdk.App({ context: { "aws:cdk:bundling-stacks": [] } }); + const backendStack = new cdk.Stack(app, "MockBackendStack", { + env: { account: "123456789012", region: "us-east-1" }, + }); + + const agentEventBus = new events.EventBus(backendStack, "AgentEventBus", { + eventBusName: "citadel-agents-test", + }); + const mkTable = (id: string, name: string, pk: string) => + new dynamodb.Table(backendStack, id, { + tableName: name, + partitionKey: { name: pk, type: dynamodb.AttributeType.STRING }, + billingMode: dynamodb.BillingMode.PAY_PER_REQUEST, + removalPolicy: cdk.RemovalPolicy.DESTROY, + }); + const agentConfigTable = mkTable( + "AgentConfigTable", + "citadel-agents-test", + "agentId", + ); + const workflowsTable = mkTable( + "WorkflowsTable", + "citadel-workflows-test", + "workflowId", + ); + const executionsTable = mkTable( + "ExecutionsTable", + "citadel-executions-test", + "executionId", + ); + const executionSpecificationsTable = mkTable( + "ExecutionSpecificationsTable", + "citadel-execution-specifications-test", + "specId", + ); + const codeBucket = new Bucket(backendStack, "CodeBucket", { + bucketName: "citadel-code-test", + }); + const fanoutFunction = new lambda.Function(backendStack, "FanoutFunction", { + runtime: lambda.Runtime.NODEJS_24_X, + handler: "workflow-progress-fanout.handler", + code: lambda.Code.fromAsset("dist/lambda"), + timeout: cdk.Duration.seconds(30), + }); + const appSyncApi = new appsync.GraphqlApi(backendStack, "MockApi", { + name: "mock-api", + schema: appsync.SchemaFile.fromAsset( + path.resolve(__dirname, "../src/schema/schema.graphql"), + ), + }); + + const stack = new ArbiterStack(app, "TestArbiterStack", { + environment: "test", + env: { account: "123456789012", region: "us-east-1" }, + agentEventBus, + agentConfigTable, + codeBucket, + workflowsTable, + executionsTable, + fanoutFunction, + appSyncEndpoint: appSyncApi.graphqlUrl, + executionSpecificationsTable, + }); + template = Template.fromStack(stack); + }); + + test("ledger table has org-scoped composite key, TTL, PITR, and SSE", () => { + template.hasResourceProperties("AWS::DynamoDB::Table", { + TableName: "citadel-tool-execution-ledger-test", + KeySchema: Match.arrayWith([ + { AttributeName: "pk", KeyType: "HASH" }, + { AttributeName: "sk", KeyType: "RANGE" }, + ]), + TimeToLiveSpecification: { AttributeName: "ttl", Enabled: true }, + PointInTimeRecoverySpecification: { PointInTimeRecoveryEnabled: true }, + SSESpecification: Match.objectLike({ SSEEnabled: true }), + BillingMode: "PAY_PER_REQUEST", + }); + }); + + test("worker env wires the ledger table name", () => { + template.hasResourceProperties("AWS::Lambda::Function", { + Handler: "index.lambda_handler", + Environment: { + Variables: Match.objectLike({ + TOOL_EXECUTION_LEDGER_TABLE: Match.anyValue(), + }), + }, + }); + }); + + test("worker grant on the ledger is Put/Get/Update only — no Delete/Scan", () => { + // The exact-array match asserts the statement carries precisely these + // three actions (least privilege): a stray DeleteItem/Scan would fail it. + template.hasResourceProperties("AWS::IAM::Policy", { + PolicyDocument: { + Statement: Match.arrayWith([ + Match.objectLike({ + Effect: "Allow", + Action: [ + "dynamodb:PutItem", + "dynamodb:GetItem", + "dynamodb:UpdateItem", + ], + }), + ]), + }, + }); + }); + + test("no policy grants Delete/Scan on the ledger table", () => { + // Defense-in-depth: scan every IAM policy statement and assert none that + // references the ledger table's ARN carries DeleteItem/Scan/Query. + const ledgerTables = template.findResources("AWS::DynamoDB::Table", { + Properties: { TableName: "citadel-tool-execution-ledger-test" }, + }); + const ledgerLogicalId = Object.keys(ledgerTables)[0]; + expect(ledgerLogicalId).toBeDefined(); + + const policies = template.findResources("AWS::IAM::Policy"); + const forbidden = [ + "dynamodb:DeleteItem", + "dynamodb:Scan", + "dynamodb:Query", + ]; + for (const policy of Object.values(policies) as any[]) { + for (const stmt of policy.Properties.PolicyDocument.Statement) { + const actions = Array.isArray(stmt.Action) + ? stmt.Action + : [stmt.Action]; + const refsLedger = JSON.stringify(stmt.Resource ?? "").includes( + ledgerLogicalId, + ); + if (refsLedger) { + for (const f of forbidden) { + expect(actions).not.toContain(f); + } + } + } + } + }); +}); diff --git a/docs/TOOL_IDEMPOTENCY.md b/docs/TOOL_IDEMPOTENCY.md new file mode 100644 index 00000000..5c794d62 --- /dev/null +++ b/docs/TOOL_IDEMPOTENCY.md @@ -0,0 +1,70 @@ +# Tool-Call Idempotency (PR1 of 2) + +Makes a **governed worker tool call** exactly-once *within an attempt* and safe +under a reservation race, using an org-scoped, TTL'd DynamoDB ledger. + +## What PR1 guarantees — stated precisely + +- **Exactly-once execution of a side effect is GUARANTEED for calls that resolve + to the same idempotency key** — SQS/redelivery, same-attempt SDK/Strands + retries, and concurrent split-brain with identical keys. The reservation's + conditional first-write-wins is the mechanism: one caller wins and executes; + every other caller is absorbed (recorded result) or bounced with a + **retryable no-execution error** and never executes. +- **Reservation-race safety**: the concurrent loser bounded-polls for the + winner's result, then returns a retryable error — it never runs the tool. A + dead holder is reclaimed via a conditional CAS. + +## What PR1 does NOT guarantee (and why) + +- It is **NOT** exactly-once across nondeterministic re-dispatch. The key is + attempt-scoped — `(orgId#executionId, nodeId#callIndex#toolName#argsHash)`. + A watchdog re-dispatch runs a fresh LLM body whose tool calls may reorder or + reword, yielding *different* keys the ledger cannot recognize. +- Closing that gap requires a **worker `dispatchGeneration` fence** (a stale + re-dispatched worker refused at reserve time). That fence is **DEFERRED to + PR2** and is **REQUIRED for the complete guarantee**. Do not read PR1 as the + complete exactly-once guarantee. + +Also deferred to PR2: S3 offload of oversized results (PR1 records a +deterministic marker instead), and client-token passthrough to adapters that +support end-to-end dedupe. + +## Key derivation + +`argsHash = sha256(canonicalize(toolInput))`. Canonicalization rejects +non-string dict keys deterministically, collapses integral floats and `-0.0` +(`2.0 == 2`, `-0.0 == 0`), rejects `NaN`/`Infinity`, and preserves `null` +(`{"a": null}` never equals `{}`). Dispatch generation is deliberately NOT in +the key — putting it there would mint a fresh key on every re-dispatch and +guarantee duplicates. + +## Ledger table + +`citadel-tool-execution-ledger-{env}` — PK `orgId#executionId`, SK +`nodeId#callIndex#toolName#argsHash`, TTL attribute `ttl` (48h, derived from the +**server** write time, not a producer clock). It is an **operational dedupe** +table, **NOT an audit artifact** (distinct from the 90-day governance ledger). +`orgId` is resolved server-side from the execution row — never trusted from a +subprocess payload — so the PK prefix gives structural cross-org isolation. + +Worker IAM is least-privilege: `PutItem` / `GetItem` / `UpdateItem` on this one +table only — no `DeleteItem` (release is a status transition), no `Scan`/`Query`. + +## Failure matrix + +| Outcome | Ledger action | Retryable | Re-executed? | +|---|---|---|---| +| Success | `completed` + result | — | No (replay returns recorded result) | +| Terminal (4xx / `applied`) | `failed` | No | No (replay returns recorded failure) | +| Retryable, provably not sent | `released` | Yes | Yes (next attempt re-reserves) | +| Unknown outcome (5xx/timeout after send, un-tokened) | `failed`, `outcomeIndeterminate` | No | **Never** — fail-safe, surfaced | + +## Strands seam + +Verified against `strands-agents==1.30.0`: there is no `AgentToolHandler`; +`Agent.__init__` takes `hooks: list[HookProvider]`. The single atomic seam is a +`BeforeToolCallEvent` hook that replaces `selected_tool` with a wrapper whose +`stream()` runs reserve → execute → finalize in one coroutine (no pre/post +window). The synchronous `execute_idempotent` coordinator is the unit-tested +authority for the invariant.