feat(tools): idempotency keys and execution ledger for tool calls

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",
+        ) == []

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