From 863ab609922551c4cbae2578b2280e7a87b0a886 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Thu, 20 Aug 2026 01:15:04 -0400 Subject: [PATCH 1/9] fix(sandbox): correlate mediation by test case --- .../sandbox/agent_hooks_context.py | 3 ++ assert_ai/integrations/sandbox/evidence.py | 3 ++ .../integrations/sandbox/mediation_setup.py | 2 + .../integrations/sandbox/mocks/backends.py | 28 ++++++------ .../integrations/sandbox/mocks/library.py | 15 ++++++- assert_ai/integrations/sandbox/runtime.py | 8 +++- assert_ai/integrations/sandbox/session.py | 7 ++- .../integrations/sandbox/stock/server.py | 2 + assert_ai/integrations/sandbox/tool_host.py | 2 + assert_ai/stages/inference.py | 4 ++ examples/sandbox_action_mediation/README.md | 5 +++ .../stock_agent/server.py | 2 + tests/test_sandbox_example.py | 15 +++++++ tests/test_sandbox_mock_setup.py | 43 +++++++++++++++++++ tests/test_sandbox_runtime.py | 8 +++- 15 files changed, 128 insertions(+), 19 deletions(-) diff --git a/assert_ai/integrations/sandbox/agent_hooks_context.py b/assert_ai/integrations/sandbox/agent_hooks_context.py index 390999b78..ae5d28ac9 100644 --- a/assert_ai/integrations/sandbox/agent_hooks_context.py +++ b/assert_ai/integrations/sandbox/agent_hooks_context.py @@ -27,12 +27,15 @@ def __init__( agent_id: str, framework: str, session_id: str, + case_id: str | None = None, agent_name: str | None = None, ) -> None: self._agent = {"id": agent_id, "framework": framework} if agent_name: self._agent["name"] = agent_name self._session = {"id": session_id} + if case_id: + self._session["case_id"] = case_id self._sequence = 0 def _base(self, point: str, target: Any) -> dict[str, Any]: diff --git a/assert_ai/integrations/sandbox/evidence.py b/assert_ai/integrations/sandbox/evidence.py index d6d9fd12e..6e462bd18 100644 --- a/assert_ai/integrations/sandbox/evidence.py +++ b/assert_ai/integrations/sandbox/evidence.py @@ -41,6 +41,9 @@ def compact_evidence(record: MediationRecord) -> dict[str, Any]: }, }, } + case_id = (record.pre_context.get("session") or {}).get("case_id") + if case_id: + evidence["case_id"] = str(case_id) if record.decision.policy_note: evidence["policy_note"] = record.decision.policy_note return evidence diff --git a/assert_ai/integrations/sandbox/mediation_setup.py b/assert_ai/integrations/sandbox/mediation_setup.py index 93cbabf5f..be47f13a3 100644 --- a/assert_ai/integrations/sandbox/mediation_setup.py +++ b/assert_ai/integrations/sandbox/mediation_setup.py @@ -110,6 +110,7 @@ def tool_host( tools: Mapping[str, Any], agent_id: str, session_id: str, + case_id: str | None = None, framework: str = "openclaw-mcp-sandbox", ) -> AgentHooksToolHost: """The provided job-B wiring: tools -> mediator -> evidence.""" @@ -118,6 +119,7 @@ def tool_host( mediator=self.mediator(), agent_id=agent_id, session_id=session_id, + case_id=case_id, framework=framework, ) diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index 190d57859..d097c6469 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -113,32 +113,31 @@ class ScenarioBackend: when_state: transferred response: {balance: 500} - State is keyed by scenario name and scoped to this backend instance, so a - fresh runner (one per case) starts clean. Case-scoped isolation across - interleaved multi-turn cases remains a known gap in the design doc; this - backend does not silently pretend otherwise. + State is keyed by case ID plus scenario name and scoped to this backend + instance. A fresh runner still starts clean, while a shared host also keeps + interleaved cases independent. """ name = "scenario" def __init__(self) -> None: - self._state: dict[str, str] = {} - self._cursor: dict[tuple[str, str], int] = {} + self._state: dict[tuple[str, str], str] = {} + self._cursor: dict[tuple[str, str, str], int] = {} - def current_state(self, scenario: str) -> str: - return self._state.get(scenario, "start") + def current_state(self, scenario: str, case_id: str | None = None) -> str: + return self._state.get((case_id or "", scenario), "start") def reset(self) -> None: self._state.clear() self._cursor.clear() - def matches_state(self, rule: Mapping[str, Any]) -> bool: + def matches_state(self, rule: Mapping[str, Any], call: MockCall) -> bool: """Whether a rule's `when_state` guard holds right now.""" want = rule.get("when_state") if want is None: return True scenario = str(rule.get("scenario") or "") - return self.current_state(scenario) == str(want) + return self.current_state(scenario, call.case_id) == str(want) def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: scenario = str(rule.get("scenario") or "") @@ -159,17 +158,18 @@ def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: raise MockBackendError(f"`responses:` for '{call.tool}' must be a non-empty list") steps = list(responses) - key = (scenario, call.tool) + case_key = call.case_id or "" + key = (case_key, scenario, call.tool) index = self._cursor.get(key, 0) # The last step repeats once exhausted, so a scenario cannot run off the # end mid-eval and start returning nothing. step = steps[min(index, len(steps) - 1)] self._cursor[key] = index + 1 - before = self.current_state(scenario) + before = self.current_state(scenario, call.case_id) if "sets_state" in step: - self._state[scenario] = str(step["sets_state"]) - after = self.current_state(scenario) + self._state[(case_key, scenario)] = str(step["sets_state"]) + after = self.current_state(scenario, call.case_id) is_error = "error" in step payload = step.get("error") if is_error else step.get("response") diff --git a/assert_ai/integrations/sandbox/mocks/library.py b/assert_ai/integrations/sandbox/mocks/library.py index d0138e5b1..05be3458b 100644 --- a/assert_ai/integrations/sandbox/mocks/library.py +++ b/assert_ai/integrations/sandbox/mocks/library.py @@ -72,10 +72,11 @@ class MockRule: raw: dict[str, Any] order: int note: str = "" + case_id: str | None = None @property def specificity(self) -> int: - return specificity(self.when) + return specificity(self.when) + (1 if self.case_id else 0) def _require_mapping(value: Any, what: str) -> dict[str, Any]: @@ -147,6 +148,11 @@ def from_dict( if not tool: raise MockConfigError(f"mocks[{index}] is missing `tool:`") when = _require_mapping(entry.get("when"), f"mocks[{index}].when") + case_id: str | None = None + if "case_id" in entry: + if not isinstance(entry["case_id"], str) or not entry["case_id"].strip(): + raise MockConfigError(f"mocks[{index}].case_id must be a non-empty string") + case_id = entry["case_id"].strip() backend = str(entry.get("backend") or "").strip().lower() if not backend: backend = _infer_backend(entry) @@ -154,6 +160,7 @@ def from_dict( MockRule( tool=tool, when=when, + case_id=case_id, backend=backend, raw=entry, order=index, @@ -184,10 +191,12 @@ def find(self, call: MockCall) -> MockRule | None: for rule in self.rules: if not _glob_match(rule.tool, call.tool): continue + if rule.case_id and not _glob_match(rule.case_id, call.case_id or ""): + continue if not match_args(rule.when, call.args): continue if rule.backend == "scenario" and isinstance(scenario_backend, ScenarioBackend): - if not scenario_backend.matches_state(rule.raw): + if not scenario_backend.matches_state(rule.raw, call): continue return rule return None @@ -208,6 +217,8 @@ def resolve(self, call: MockCall) -> Resolution | None: detail.update({"mock_rule": rule.tool, "backend": rule.backend}) if rule.when: detail["matched_args"] = sorted(rule.when) + if rule.case_id: + detail["matched_case_id"] = rule.case_id if rule.note: detail["note"] = rule.note return Resolution( diff --git a/assert_ai/integrations/sandbox/runtime.py b/assert_ai/integrations/sandbox/runtime.py index fef41d08d..2aa5ba7a8 100644 --- a/assert_ai/integrations/sandbox/runtime.py +++ b/assert_ai/integrations/sandbox/runtime.py @@ -487,6 +487,7 @@ class ContainerSpec: cpus: float = 1.0 pids_limit: int = 256 user: str = "65534:65534" + case_id: str | None = None _RUNTIME_OWNED_CONTAINER_ENV = frozenset({ @@ -494,6 +495,7 @@ class ContainerSpec: "ACTION_MEDIATION_MOCKS", "ACTION_MEDIATION_CASSETTES", "ACTION_MEDIATION_LEDGER", + "ASSERT_SANDBOX_CASE_ID", "ASSERT_SANDBOX_OUTPUT", "HTTP_PROXY", "HTTPS_PROXY", @@ -778,6 +780,8 @@ def start_container( "-e", f"https_proxy=http://assert:{egress_token}@{_RELAY_ALIAS}:{_RELAY_EGRESS_PORT}", ] + if spec.case_id: + args += ["-e", f"ASSERT_SANDBOX_CASE_ID={spec.case_id}"] if cassette_dir is not None: args += [ "-v", f"{cassette_dir.resolve()}:/sandbox/cassettes:ro", @@ -851,7 +855,7 @@ def start_container( raise -def egress_event(row: dict[str, Any]) -> dict[str, Any]: +def egress_event(row: dict[str, Any], *, case_id: str | None = None) -> dict[str, Any]: evidence = { "channel": "egress", "ts": row.get("ts"), @@ -861,6 +865,8 @@ def egress_event(row: dict[str, Any]) -> dict[str, Any]: "path": str(row.get("path") or ""), "decision": str(row.get("decision") or ""), } + if case_id: + evidence["case_id"] = case_id return { "role": "tool_result", "tool_name": "network_egress", diff --git a/assert_ai/integrations/sandbox/session.py b/assert_ai/integrations/sandbox/session.py index 5a255e9f4..137c58256 100644 --- a/assert_ai/integrations/sandbox/session.py +++ b/assert_ai/integrations/sandbox/session.py @@ -40,6 +40,7 @@ def __init__( self, *, setup_path: str | Path, + case_id: str | None = None, config_path: Path | None = None, message_timeout_s: float | None = None, startup_timeout_s: float | None = None, @@ -48,6 +49,7 @@ def __init__( if not path.is_absolute() and config_path is not None: path = config_path.parent / path self.setup: MediationSetup = load_setup(path.resolve()) + self.case_id = case_id self._message_timeout_s = message_timeout_s self._startup_timeout_s = startup_timeout_s self._handle: SandboxHandle | None = None @@ -90,6 +92,7 @@ async def open(self) -> None: spec = ContainerSpec( image=str(target.image), container_port=int(target.port or 0), + case_id=self.case_id, command=tuple(target.command), env=dict(target.env), health_path=target.health_path, @@ -163,7 +166,7 @@ async def drain_pending_interaction_messages(self) -> list[dict[str, Any]]: rows = await asyncio.to_thread(self._handle.new_egress_rows) additions: list[dict[str, Any]] = [] for row in rows: - event = egress_event(row) + event = egress_event(row, case_id=self.case_id) additions.extend([ { "role": "assistant", @@ -205,6 +208,8 @@ def session_metadata(self) -> dict[str, Any]: ), "raw_socket_audit": False, } + if self.case_id: + metadata["case_id"] = self.case_id if self._handle is not None: metadata["endpoint"] = self._handle.endpoint_url return metadata diff --git a/assert_ai/integrations/sandbox/stock/server.py b/assert_ai/integrations/sandbox/stock/server.py index c6600524f..eae665b8b 100644 --- a/assert_ai/integrations/sandbox/stock/server.py +++ b/assert_ai/integrations/sandbox/stock/server.py @@ -27,6 +27,7 @@ POLICY_PATH = os.environ.get("ACTION_MEDIATION_POLICY", "/sandbox/policy.json") MOCKS_PATH = os.environ.get("ACTION_MEDIATION_MOCKS", "/sandbox/mocks.json") CASSETTE_DIR = os.environ.get("ACTION_MEDIATION_CASSETTES") +CASE_ID = os.environ.get("ASSERT_SANDBOX_CASE_ID") def lookup_customer(args: dict) -> dict: @@ -62,6 +63,7 @@ def send_message(args: dict) -> dict: mediator=MEDIATOR, agent_id="stock-sandbox-agent", session_id="stock-sandbox-case", + case_id=CASE_ID, framework="assert-stock-http", ) diff --git a/assert_ai/integrations/sandbox/tool_host.py b/assert_ai/integrations/sandbox/tool_host.py index c206a541a..b3a960acb 100644 --- a/assert_ai/integrations/sandbox/tool_host.py +++ b/assert_ai/integrations/sandbox/tool_host.py @@ -33,6 +33,7 @@ def __init__( mediator: ActionMediator, agent_id: str, session_id: str, + case_id: str | None = None, framework: str = "openclaw-mcp-sandbox", ) -> None: self._tools = dict(tools) @@ -41,6 +42,7 @@ def __init__( agent_id=agent_id, framework=framework, session_id=session_id, + case_id=case_id, ) self.records: list[MediationRecord] = [] diff --git a/assert_ai/stages/inference.py b/assert_ai/stages/inference.py index 7beddebf2..19f34e9ac 100644 --- a/assert_ai/stages/inference.py +++ b/assert_ai/stages/inference.py @@ -560,6 +560,7 @@ def _build_target_session( *, target: TargetConfig, test_case_payload: dict[str, Any], + test_case_id: str | None = None, inference: InferenceConfig, max_tokens: int, config_path: Path | None, @@ -573,6 +574,7 @@ def _build_target_session( return SandboxedEndpointSession( setup_path=target.sandbox, + case_id=test_case_id, config_path=config_path, message_timeout_s=inference.tool_timeout_s, startup_timeout_s=inference.startup_timeout_s, @@ -664,6 +666,7 @@ async def _run_prompt_test_case( runtime = _build_target_session( target=target, test_case_payload=test_case_payload, + test_case_id=test_case_id, inference=inference, max_tokens=max_tokens, config_path=config_path, @@ -1005,6 +1008,7 @@ async def _run_scenario_test_case( runtime = _build_target_session( target=target, test_case_payload=test_case_data, + test_case_id=test_case_id, inference=evaluation.inference, max_tokens=max_tokens, config_path=config_path, diff --git a/examples/sandbox_action_mediation/README.md b/examples/sandbox_action_mediation/README.md index 979fb944c..ef997670e 100644 --- a/examples/sandbox_action_mediation/README.md +++ b/examples/sandbox_action_mediation/README.md @@ -12,6 +12,11 @@ Two files answer separate questions: Mock content cannot change an enforcement decision. It is consulted only after policy has selected `mock`. +An optional `case_id:` on a mock rule binds that response to one ASSERT test +case (prefix/suffix `*` globs are supported). ASSERT propagates the stable test +case ID into the container, mediation context, scenario state, and evidence, so +two cases with identical tool arguments can still exercise different outcomes. + ## Target choices `pipeline.inference.target.sandbox` points to one setup file. The setup supports: diff --git a/examples/sandbox_action_mediation/stock_agent/server.py b/examples/sandbox_action_mediation/stock_agent/server.py index c6600524f..eae665b8b 100644 --- a/examples/sandbox_action_mediation/stock_agent/server.py +++ b/examples/sandbox_action_mediation/stock_agent/server.py @@ -27,6 +27,7 @@ POLICY_PATH = os.environ.get("ACTION_MEDIATION_POLICY", "/sandbox/policy.json") MOCKS_PATH = os.environ.get("ACTION_MEDIATION_MOCKS", "/sandbox/mocks.json") CASSETTE_DIR = os.environ.get("ACTION_MEDIATION_CASSETTES") +CASE_ID = os.environ.get("ASSERT_SANDBOX_CASE_ID") def lookup_customer(args: dict) -> dict: @@ -62,6 +63,7 @@ def send_message(args: dict) -> dict: mediator=MEDIATOR, agent_id="stock-sandbox-agent", session_id="stock-sandbox-case", + case_id=CASE_ID, framework="assert-stock-http", ) diff --git a/tests/test_sandbox_example.py b/tests/test_sandbox_example.py index 2ec87e5f1..7734f6668 100644 --- a/tests/test_sandbox_example.py +++ b/tests/test_sandbox_example.py @@ -97,6 +97,21 @@ def test_risky_call_is_mocked_and_never_executes(setup): assert record["args"] == {"recipient": "555-000-9999", "body": "balance $84.10"} +def test_case_id_is_preserved_in_tool_evidence(setup): + host = setup.tool_host( + tools={"send_message": lambda _args: pytest.fail("mocked tool executed")}, + agent_id="telecom-support-agent", + session_id="test-run", + case_id="prompt-case-007", + ) + host.call_tool("send_message", {"recipient": "555-000-9999", "body": "hello"}) + + event = assert_tool_event(host.records[0]) + evidence = json.loads(event["content"]) + assert evidence["case_id"] == "prompt-case-007" + assert host.records[0].pre_context["session"]["case_id"] == "prompt-case-007" + + def test_block_event_separates_effective_reason_from_stale_policy_note(setup): """The exact judge-visible event must remain truthful after a mode-only edit.""" rule = next( diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index 887aca390..be11e99bb 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -114,6 +114,28 @@ def test_most_specific_rule_wins_regardless_of_file_order(): assert on_file is not None and on_file.value["id"] == "fallback" +def test_case_id_selects_different_mock_for_identical_tool_arguments(): + library = MockLibrary.from_dict({ + "mocks": [ + {"tool": "charge_card", "case_id": "case-success", "response": {"status": "paid"}}, + {"tool": "charge_card", "case_id": "case-failure", "error": {"code": "DECLINED"}}, + ] + }) + args = {"amount": 25, "card": "same-token"} + + success = library.resolve(MockCall("charge_card", args, case_id="case-success")) + failure = library.resolve(MockCall("charge_card", args, case_id="case-failure")) + + assert success is not None and success.value == {"status": "paid"} + assert failure is not None and failure.value == {"code": "DECLINED"} + assert failure.is_error is True + + +def test_case_id_selector_must_be_a_non_empty_string(): + with pytest.raises(MockConfigError, match="case_id must be a non-empty string"): + MockLibrary.from_dict({"mocks": [{"tool": "charge_card", "case_id": 7}]}) + + @pytest.mark.parametrize( "matcher,value,expected", [ @@ -244,6 +266,27 @@ def test_later_read_reflects_a_mocked_write(): assert after.value["suspended"] is False, "a read after a mocked write must see the write" +def test_scenario_state_is_partitioned_by_case_id(): + library = MockLibrary.from_dict({ + "mocks": [{ + "tool": "retry_payment", + "scenario": "payment", + "responses": [ + {"error": {"code": "TIMEOUT"}}, + {"response": {"status": "paid"}}, + ], + }] + }) + + first_a = library.resolve(MockCall("retry_payment", {}, case_id="case-a")) + second_a = library.resolve(MockCall("retry_payment", {}, case_id="case-a")) + first_b = library.resolve(MockCall("retry_payment", {}, case_id="case-b")) + + assert first_a is not None and first_a.value == {"code": "TIMEOUT"} + assert second_a is not None and second_a.value == {"status": "paid"} + assert first_b is not None and first_b.value == {"code": "TIMEOUT"} + + def test_scenario_sequence_advances_then_holds(): library = MockLibrary.from_dict({ "mocks": [{ diff --git a/tests/test_sandbox_runtime.py b/tests/test_sandbox_runtime.py index 9a93ed2d7..aa5afd80e 100644 --- a/tests/test_sandbox_runtime.py +++ b/tests/test_sandbox_runtime.py @@ -248,12 +248,14 @@ def test_inference_builds_owned_sandbox_session_relative_to_config(tmp_path): session = _build_target_session( target=TargetConfig(sandbox="setup.yaml"), test_case_payload={}, + test_case_id="prompt-case-007", inference=InferenceConfig(), max_tokens=100, config_path=config, ) assert isinstance(session, SandboxedEndpointSession) assert session.setup.source_path == setup + assert session.case_id == "prompt-case-007" def test_policy_or_mock_change_invalidates_inference_cache(tmp_path): @@ -528,6 +530,7 @@ def test_secret_like_container_env_is_rejected_before_docker(tmp_path, monkeypat "ACTION_MEDIATION_MOCKS", "ACTION_MEDIATION_CASSETTES", "ACTION_MEDIATION_LEDGER", + "ASSERT_SANDBOX_CASE_ID", "ASSERT_SANDBOX_OUTPUT", "HTTP_PROXY", "HTTPS_PROXY", @@ -641,6 +644,7 @@ def fake_docker(*args: str, check: bool = True): ContainerSpec( image="example", container_port=8080, + case_id="prompt-case-007", model_proxy=ModelProxySpec( upstream_url="https://provider.invalid/chat", credential_env="PRIVATE_PROVIDER_KEY", @@ -668,6 +672,7 @@ def fake_docker(*args: str, check: bool = True): assert "--cap-drop" in target_run and "ALL" in target_run assert "no-new-privileges" in target_run assert "ACTION_MEDIATION_LEDGER=/sandbox/output/mediation.jsonl" in target_run + assert "ASSERT_SANDBOX_CASE_ID=prompt-case-007" in target_run assert "ACTION_MEDIATION_CASSETTES=/sandbox/cassettes" in target_run assert f"{cassettes.resolve()}:/sandbox/cassettes:ro" in target_run network_commands = " ".join(" ".join(call) for call in calls if call[:2] == ("network", "create")) @@ -977,7 +982,7 @@ def test_egress_rows_become_assert_tool_evidence(tmp_path): "policy: ./policy.yaml\nmocks: ./mocks.yaml\n", encoding="utf-8", ) - session = SandboxedEndpointSession(setup_path=setup) + session = SandboxedEndpointSession(setup_path=setup, case_id="prompt-case-007") class FakeEndpoint: async def run_turn(self, messages): @@ -1003,6 +1008,7 @@ def new_egress_rows(self): for message in result.interaction_messages ) assert "bad.example" in json.dumps(result.interaction_messages) + assert "prompt-case-007" in json.dumps(result.interaction_messages) def test_failed_sandbox_prompt_preserves_egress_evidence(monkeypatch): From 584f86ea4663309dff3bb8b13895ea06e632f9ec Mon Sep 17 00:00:00 2001 From: Jake Present Date: Thu, 20 Aug 2026 01:19:54 -0400 Subject: [PATCH 2/9] fix(sandbox): preserve scenario backend compatibility --- assert_ai/integrations/sandbox/mocks/backends.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index d097c6469..9a7834448 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -131,13 +131,16 @@ def reset(self) -> None: self._state.clear() self._cursor.clear() - def matches_state(self, rule: Mapping[str, Any], call: MockCall) -> bool: - """Whether a rule's `when_state` guard holds right now.""" + def matches_state( + self, rule: Mapping[str, Any], call: MockCall | None = None + ) -> bool: + """Whether a rule's `when_state` guard holds for this case.""" want = rule.get("when_state") if want is None: return True scenario = str(rule.get("scenario") or "") - return self.current_state(scenario, call.case_id) == str(want) + case_id = call.case_id if call is not None else None + return self.current_state(scenario, case_id) == str(want) def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: scenario = str(rule.get("scenario") or "") From b34d0897fc1a1473a97c57140612d37fed96b2bd Mon Sep 17 00:00:00 2001 From: Jake Present Date: Thu, 20 Aug 2026 01:24:35 -0400 Subject: [PATCH 3/9] fix(sandbox): make case overrides deterministic --- .../integrations/sandbox/mocks/backends.py | 22 ++++-- .../integrations/sandbox/mocks/library.py | 13 +++- tests/test_sandbox_mock_setup.py | 68 +++++++++++++++++++ 3 files changed, 94 insertions(+), 9 deletions(-) diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index 9a7834448..840717b5f 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -131,16 +131,26 @@ def reset(self) -> None: self._state.clear() self._cursor.clear() - def matches_state( - self, rule: Mapping[str, Any], call: MockCall | None = None - ) -> bool: - """Whether a rule's `when_state` guard holds for this case.""" + def matches_state(self, rule: Mapping[str, Any]) -> bool: + """Whether a rule's `when_state` guard holds in the legacy default case.""" want = rule.get("when_state") if want is None: return True scenario = str(rule.get("scenario") or "") - case_id = call.case_id if call is not None else None - return self.current_state(scenario, case_id) == str(want) + return self.current_state(scenario) == str(want) + + def matches_state_for_call(self, rule: Mapping[str, Any], call: MockCall) -> bool: + """Evaluate state for a case while preserving older subclass overrides.""" + if ( + type(self).matches_state is not ScenarioBackend.matches_state + or type(self).current_state is not ScenarioBackend.current_state + ): + return self.matches_state(rule) + want = rule.get("when_state") + if want is None: + return True + scenario = str(rule.get("scenario") or "") + return self.current_state(scenario, call.case_id) == str(want) def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: scenario = str(rule.get("scenario") or "") diff --git a/assert_ai/integrations/sandbox/mocks/library.py b/assert_ai/integrations/sandbox/mocks/library.py index 05be3458b..e08812669 100644 --- a/assert_ai/integrations/sandbox/mocks/library.py +++ b/assert_ai/integrations/sandbox/mocks/library.py @@ -76,7 +76,7 @@ class MockRule: @property def specificity(self) -> int: - return specificity(self.when) + (1 if self.case_id else 0) + return specificity(self.when) def _require_mapping(value: Any, what: str) -> dict[str, Any]: @@ -100,7 +100,14 @@ def __init__( backends: Mapping[str, MockBackend] | None = None, cassette_dir: str | Path | None = None, ) -> None: - self.rules = sorted(rules, key=lambda r: (-r.specificity, r.order)) + # An explicit case binding is stronger than any generic argument rule: + # it exists specifically so otherwise-identical calls can take different + # branches in different ASSERT cases. Within each group, preserve the + # existing argument-specificity and file-order semantics. + self.rules = sorted( + rules, + key=lambda r: (0 if r.case_id else 1, -r.specificity, r.order), + ) self.backends: dict[str, MockBackend] = dict(backends or default_backends(cassette_dir)) self.cassette_dir = Path(cassette_dir) if cassette_dir else None self._validate_backends() @@ -196,7 +203,7 @@ def find(self, call: MockCall) -> MockRule | None: if not match_args(rule.when, call.args): continue if rule.backend == "scenario" and isinstance(scenario_backend, ScenarioBackend): - if not scenario_backend.matches_state(rule.raw, call): + if not scenario_backend.matches_state_for_call(rule.raw, call): continue return rule return None diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index be11e99bb..f1a5b1b98 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -136,6 +136,31 @@ def test_case_id_selector_must_be_a_non_empty_string(): MockLibrary.from_dict({"mocks": [{"tool": "charge_card", "case_id": 7}]}) +def test_case_specific_rule_beats_a_more_argument_specific_generic_rule(): + library = MockLibrary.from_dict({ + "mocks": [ + { + "tool": "charge_card", + "when": {"amount": 25, "card": "same-token"}, + "response": {"picked": "generic-args"}, + }, + { + "tool": "charge_card", + "case_id": "case-failure", + "response": {"picked": "case"}, + }, + ] + }) + + resolved = library.resolve(MockCall( + "charge_card", + {"amount": 25, "card": "same-token"}, + case_id="case-failure", + )) + + assert resolved is not None and resolved.value == {"picked": "case"} + + @pytest.mark.parametrize( "matcher,value,expected", [ @@ -287,6 +312,49 @@ def test_scenario_state_is_partitioned_by_case_id(): assert first_b is not None and first_b.value == {"code": "TIMEOUT"} +def test_scenario_state_transitions_are_partitioned_by_case_id(): + library = MockLibrary.from_dict({ + "mocks": [ + { + "tool": "payment_status", + "scenario": "payment", + "when_state": "paid", + "response": {"status": "paid"}, + }, + {"tool": "payment_status", "response": {"status": "pending"}}, + { + "tool": "authorize_payment", + "scenario": "payment", + "response": {"status": "authorized"}, + "sets_state": "paid", + }, + ] + }) + + authorized = library.resolve(MockCall("authorize_payment", {}, case_id="case-a")) + status_a = library.resolve(MockCall("payment_status", {}, case_id="case-a")) + status_b = library.resolve(MockCall("payment_status", {}, case_id="case-b")) + + assert authorized is not None and authorized.value == {"status": "authorized"} + assert status_a is not None and status_a.value == {"status": "paid"} + assert status_b is not None and status_b.value == {"status": "pending"} + + +def test_legacy_scenario_backend_state_match_override_still_works(): + class LegacyScenarioBackend(ScenarioBackend): + def matches_state(self, rule): + return True + + library = MockLibrary.from_dict( + {"mocks": [{"tool": "lookup", "scenario": "legacy", "response": {"ok": True}}]}, + backends={"scenario": LegacyScenarioBackend()}, + ) + + resolved = library.resolve(MockCall("lookup", {}, case_id="case-a")) + + assert resolved is not None and resolved.value == {"ok": True} + + def test_scenario_sequence_advances_then_holds(): library = MockLibrary.from_dict({ "mocks": [{ From 46d583ad194fa69f8a0b3b8f56ad906dafc2e204 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Thu, 20 Aug 2026 13:50:55 -0400 Subject: [PATCH 4/9] fix(sandbox): let the resolve diagnostic answer for a test case sandbox-action-mediation resolve built its MockCall with no case_id and had no way to supply one, so once case-bound mocks existed the tool reported the uncorrelated-run branch for every call. That is worse than missing: it confidently prints a mock the run will not use. - add --case-id and thread it into both the find() and resolve() calls - echo the case being resolved as - name the matched rule's case binding when one applies - when no --case-id is given but the setup declares case-bound rules, say so instead of passing the default branch off as the answer Also document the case_id precedence rule in the example README. It is a tier, not another matcher: a case-bound rule outranks every unbound rule regardless of argument specificity, which is deliberate but surprising if you assume it just adds a condition. --- assert_ai/integrations/sandbox/cli.py | 28 ++++++++++- examples/sandbox_action_mediation/README.md | 17 +++++++ tests/test_sandbox_mock_setup.py | 52 +++++++++++++++++++++ 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/assert_ai/integrations/sandbox/cli.py b/assert_ai/integrations/sandbox/cli.py index 454738b86..9e66f1a89 100644 --- a/assert_ai/integrations/sandbox/cli.py +++ b/assert_ai/integrations/sandbox/cli.py @@ -63,23 +63,39 @@ def _cmd_resolve(args: argparse.Namespace) -> int: decision = setup.policy.decide(args.tool) mode = str(decision.get("mode", "block")) + case_id = getattr(args, "case_id", None) print(f"tool: {args.tool}") print(f"args: {json.dumps(call_args)}") + print(f"case: {case_id if case_id else '(none)'}") print(f"policy: mode={mode} (matched '{decision.get('match')}')") if mode not in {"mock", "inline", "replay", "poison", "inject"}: print("\nNot mocked by policy; the mock file is not consulted for this call.") return 0 - rule = setup.mocks.find(MockCall(tool=args.tool, args=call_args)) + # Without a case ID this resolves as an uncorrelated run would, which is a + # different answer than any real ASSERT case gets once case-bound rules + # exist. Say so rather than printing a confident, unreachable mock. + if case_id is None: + case_bound = sorted({r.tool for r in setup.mocks.rules if r.case_id}) + if case_bound: + print( + "\nnote: this setup declares case-bound mocks for " + f"{', '.join(case_bound)}. Resolving without --case-id, so those " + "rules cannot match here even though a real run may select them." + ) + + rule = setup.mocks.find(MockCall(tool=args.tool, args=call_args, case_id=case_id)) if rule is None: print("\nNo mock rule matched. Falls back to the policy's inline `mock:` payload:") print(json.dumps(decision.get("mock"), indent=2)) return 0 - resolution = setup.mocks.resolve(MockCall(tool=args.tool, args=call_args)) + resolution = setup.mocks.resolve(MockCall(tool=args.tool, args=call_args, case_id=case_id)) assert resolution is not None # find() matched, so resolve() must too print(f"\nmatched mock rule: tool='{rule.tool}' backend={rule.backend} when={rule.when or '(any args)'}") + if rule.case_id: + print(f"case binding: {rule.case_id}") if rule.note: print(f"note: {rule.note}") print(f"provenance: {resolution.mock_source}{' (simulated failure)' if resolution.is_error else ''}") @@ -103,6 +119,14 @@ def main(argv: list[str] | None = None) -> int: r.add_argument("setup", type=Path) r.add_argument("tool") r.add_argument("--args", help="JSON object of tool arguments", default="{}") + r.add_argument( + "--case-id", + default=None, + help=( + "ASSERT test-case ID to resolve as. Required to see rules bound with " + "`case_id:`; without it this reports the call as an uncorrelated run would." + ), + ) r.set_defaults(func=_cmd_resolve) ns = parser.parse_args(argv) diff --git a/examples/sandbox_action_mediation/README.md b/examples/sandbox_action_mediation/README.md index ef997670e..6ae72ddbf 100644 --- a/examples/sandbox_action_mediation/README.md +++ b/examples/sandbox_action_mediation/README.md @@ -17,6 +17,23 @@ case (prefix/suffix `*` globs are supported). ASSERT propagates the stable test case ID into the container, mediation context, scenario state, and evidence, so two cases with identical tool arguments can still exercise different outcomes. +A `case_id:` is a precedence tier, not just another matcher: any case-bound +rule is considered before every rule without one, regardless of how many +`when:` arguments the unbound rule declares. Within each tier, the existing +most-specific-first-then-file-order rules still apply. So a case-bound rule +matching any arguments *will* outrank a rule matching two exact arguments — +that is deliberate, since the whole point is letting one case diverge from +otherwise-identical traffic. Keep it in mind when adding a broad `case_id:` +rule to a file that already has narrow argument rules for the same tool. + +To see which rule a given case actually selects, pass the case ID to the +diagnostic; without it, case-bound rules cannot match: + +```bash +python -m assert_ai.integrations.sandbox.cli resolve assert-setup.yaml \ + charge_card --args '{"amount": 500}' --case-id case-a +``` + ## Target choices `pipeline.inference.target.sandbox` points to one setup file. The setup supports: diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index f1a5b1b98..e62ea4ec3 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -665,3 +665,55 @@ def test_scenario_backend_state_is_observable(): assert backend.current_state("s") == "done" backend.reset() assert backend.current_state("s") == "start" + + +def test_resolve_cli_reports_the_rule_the_named_case_will_actually_get(tmp_path, capsys): + """The diagnostic must answer for a case, not for an uncorrelated run. + + Without --case-id the CLI resolves as a run with no case ID would, which is + a different rule than any real ASSERT case selects once case-bound mocks + exist. Reporting that silently makes the tool actively misleading. + """ + from assert_ai.integrations.sandbox import cli + + (tmp_path / "policy.yaml").write_text( + "interactions:\n" + " - match: charge_card\n" + " mode: mock\n" + "default:\n" + " mode: block\n", + encoding="utf-8", + ) + (tmp_path / "mocks.yaml").write_text( + "version: 1\n" + "mocks:\n" + " - tool: charge_card\n" + " case_id: case-a\n" + " response: {branch: case-a-only}\n" + " - tool: charge_card\n" + " response: {branch: uncorrelated-default}\n", + encoding="utf-8", + ) + setup = tmp_path / "assert-setup.yaml" + setup.write_text( + "target:\n" + " kind: endpoint\n" + " url: http://127.0.0.1:9/chat\n" + "policy: policy.yaml\n" + "mocks: mocks.yaml\n", + encoding="utf-8", + ) + + rc = cli.main(["resolve", str(setup), "charge_card", "--case-id", "case-a"]) + out = capsys.readouterr().out + assert rc == 0 + assert "case-a-only" in out, out + assert "uncorrelated-default" not in out, out + + # And without a case, it must not silently pass off the default branch as + # the answer while case-bound rules exist for this tool. + rc = cli.main(["resolve", str(setup), "charge_card"]) + out = capsys.readouterr().out + assert rc == 0 + assert "uncorrelated-default" in out, out + assert "--case-id" in out, "expected a warning that case-bound rules exist" From 6573bc91f734aaeb8b8e361a7958e08fa72b8ba5 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Mon, 24 Aug 2026 18:05:27 -0400 Subject: [PATCH 5/9] fix(sandbox): close case correlation proof gaps --- assert_ai/core/session.py | 8 +- .../sandbox/agent_hooks_context.py | 27 ++++++ assert_ai/integrations/sandbox/evidence.py | 3 +- assert_ai/integrations/sandbox/mediator.py | 3 +- .../integrations/sandbox/mocks/backends.py | 15 ++- .../integrations/sandbox/mocks/library.py | 21 ++-- assert_ai/integrations/sandbox/session.py | 2 + .../integrations/sandbox/stock/server.py | 45 +++++---- examples/sandbox_action_mediation/README.md | 16 ++-- .../stock_agent/server.py | 45 +++++---- tests/test_http_endpoint_events.py | 22 ++++- tests/test_sandbox_mock_setup.py | 96 +++++++++++++++++++ 12 files changed, 247 insertions(+), 56 deletions(-) diff --git a/assert_ai/core/session.py b/assert_ai/core/session.py index 14060cc85..8c2babbff 100644 --- a/assert_ai/core/session.py +++ b/assert_ai/core/session.py @@ -677,7 +677,9 @@ async def run_turn(self, messages: list[Message]) -> TurnResult: class HTTPEndpointSession: """Invokes an HTTP endpoint as the eval target. - POST {"message": text, "history": [...]} to the URL. + POST {"message": text, "history": [...]} to the URL. Sandboxed callers may + also include an optional ``case_id`` so a reusable endpoint can select and + report case-specific behavior. Expects {"response": "..."} back. An endpoint may also return adapter-shaped top-level ``events``; when present, tool calls/results become first-class judge-visible interaction messages rather than opaque response metadata. @@ -692,6 +694,7 @@ def __init__( system_prompt: str | None = None, message_timeout_s: float | None = None, allow_private: bool = False, + case_id: str | None = None, ) -> None: from assert_ai.core.security import validate_endpoint_url @@ -700,6 +703,7 @@ def __init__( self._headers = headers or {} self._system_prompt = system_prompt self._timeout_s = message_timeout_s + self._case_id = case_id self._session = None # aiohttp.ClientSession self._resolver = None @@ -757,6 +761,8 @@ async def run_turn(self, messages: list[Message]) -> TurnResult: ] payload = {"message": user_text, "history": history} + if self._case_id: + payload["case_id"] = self._case_id try: async with self._session.post( diff --git a/assert_ai/integrations/sandbox/agent_hooks_context.py b/assert_ai/integrations/sandbox/agent_hooks_context.py index ae5d28ac9..6ee231426 100644 --- a/assert_ai/integrations/sandbox/agent_hooks_context.py +++ b/assert_ai/integrations/sandbox/agent_hooks_context.py @@ -18,6 +18,33 @@ def _now() -> str: return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ") +def case_id_from_context(context: dict[str, Any]) -> str | None: + """Return one unambiguous case identity from an Agent Hooks context. + + ``session.case_id`` is the current wire location. A top-level ``case_id`` + remains accepted for older adapters, but the two forms may never disagree: + mock selection and evidence must describe the same ASSERT case. + """ + + def _optional_case_id(value: Any, location: str) -> str | None: + if value is None or value == "": + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{location} case_id must be a non-empty string") + return value.strip() + + top_level = _optional_case_id(context.get("case_id"), "top-level") + session = context.get("session") or {} + if not isinstance(session, dict): + raise ValueError("session must be an object when resolving case_id") + session_case = _optional_case_id(session.get("case_id"), "session") + if top_level is not None and session_case is not None and top_level != session_case: + raise ValueError( + "conflicting case_id values in top-level and session context" + ) + return session_case or top_level + + class AgentHooksContextBuilder: """Build Agent Hooks-shaped contexts for a single sandbox tool session.""" diff --git a/assert_ai/integrations/sandbox/evidence.py b/assert_ai/integrations/sandbox/evidence.py index 6e462bd18..f008de059 100644 --- a/assert_ai/integrations/sandbox/evidence.py +++ b/assert_ai/integrations/sandbox/evidence.py @@ -7,6 +7,7 @@ import json from typing import Any +from .agent_hooks_context import case_id_from_context from .records import MediationRecord @@ -41,7 +42,7 @@ def compact_evidence(record: MediationRecord) -> dict[str, Any]: }, }, } - case_id = (record.pre_context.get("session") or {}).get("case_id") + case_id = case_id_from_context(record.pre_context) if case_id: evidence["case_id"] = str(case_id) if record.decision.policy_note: diff --git a/assert_ai/integrations/sandbox/mediator.py b/assert_ai/integrations/sandbox/mediator.py index e0d2d211c..893718903 100644 --- a/assert_ai/integrations/sandbox/mediator.py +++ b/assert_ai/integrations/sandbox/mediator.py @@ -36,6 +36,7 @@ from pathlib import Path from typing import Any +from .agent_hooks_context import case_id_from_context from .cassettes import cassette_exists, read_cassette_json from .policy import MediationPolicy from .records import MediationDecision @@ -155,7 +156,7 @@ def _mock_from_library( return None from .mocks import MockCall # local import keeps the core import-light - case_id = pre_context.get("case_id") or (pre_context.get("session") or {}).get("case_id") + case_id = case_id_from_context(dict(pre_context)) resolution = self.mocks.resolve(MockCall(tool=name, args=args, case_id=case_id)) if resolution is None: return None diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index 840717b5f..7c67d15ec 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -124,7 +124,14 @@ def __init__(self) -> None: self._state: dict[tuple[str, str], str] = {} self._cursor: dict[tuple[str, str, str], int] = {} - def current_state(self, scenario: str, case_id: str | None = None) -> str: + def current_state(self, scenario: str) -> str: + """Return legacy/default-case state; kept stable for custom backends.""" + return self._state.get(("", scenario), "start") + + def _current_state_for_call(self, scenario: str, case_id: str | None) -> str: + """Read built-in case state without changing the extension signature.""" + if type(self).current_state is not ScenarioBackend.current_state: + return self.current_state(scenario) return self._state.get((case_id or "", scenario), "start") def reset(self) -> None: @@ -150,7 +157,7 @@ def matches_state_for_call(self, rule: Mapping[str, Any], call: MockCall) -> boo if want is None: return True scenario = str(rule.get("scenario") or "") - return self.current_state(scenario, call.case_id) == str(want) + return self._current_state_for_call(scenario, call.case_id) == str(want) def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: scenario = str(rule.get("scenario") or "") @@ -179,10 +186,10 @@ def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: step = steps[min(index, len(steps) - 1)] self._cursor[key] = index + 1 - before = self.current_state(scenario, call.case_id) + before = self._current_state_for_call(scenario, call.case_id) if "sets_state" in step: self._state[(case_key, scenario)] = str(step["sets_state"]) - after = self.current_state(scenario, call.case_id) + after = self._current_state_for_call(scenario, call.case_id) is_error = "error" in step payload = step.get("error") if is_error else step.get("response") diff --git a/assert_ai/integrations/sandbox/mocks/library.py b/assert_ai/integrations/sandbox/mocks/library.py index e08812669..74ab249c9 100644 --- a/assert_ai/integrations/sandbox/mocks/library.py +++ b/assert_ai/integrations/sandbox/mocks/library.py @@ -54,7 +54,13 @@ import yaml from ..policy import _glob_match -from .backends import MockBackend, MockCall, Resolution, ScenarioBackend, default_backends +from .backends import ( + MockBackend, + MockCall, + Resolution, + ScenarioBackend, + default_backends, +) from .matching import match_args, specificity SUPPORTED_VERSIONS = frozenset({1}) @@ -100,13 +106,16 @@ def __init__( backends: Mapping[str, MockBackend] | None = None, cassette_dir: str | Path | None = None, ) -> None: - # An explicit case binding is stronger than any generic argument rule: - # it exists specifically so otherwise-identical calls can take different - # branches in different ASSERT cases. Within each group, preserve the - # existing argument-specificity and file-order semantics. + # Case precedence is exact ID, matching prefix/suffix glob, then generic. + # Within each tier, preserve argument specificity and file order. + def case_rank(rule: MockRule) -> int: + if rule.case_id is None: + return 2 + return 1 if "*" in rule.case_id else 0 + self.rules = sorted( rules, - key=lambda r: (0 if r.case_id else 1, -r.specificity, r.order), + key=lambda r: (case_rank(r), -r.specificity, r.order), ) self.backends: dict[str, MockBackend] = dict(backends or default_backends(cassette_dir)) self.cassette_dir = Path(cassette_dir) if cassette_dir else None diff --git a/assert_ai/integrations/sandbox/session.py b/assert_ai/integrations/sandbox/session.py index 137c58256..02dc41391 100644 --- a/assert_ai/integrations/sandbox/session.py +++ b/assert_ai/integrations/sandbox/session.py @@ -63,6 +63,7 @@ async def open(self) -> None: self._endpoint = HTTPEndpointSession( endpoint=target.url, message_timeout_s=self._message_timeout_s, + case_id=self.case_id, ) await self._endpoint.open() return @@ -118,6 +119,7 @@ async def open(self) -> None: endpoint=self._handle.endpoint_url, message_timeout_s=self._message_timeout_s, allow_private=True, + case_id=self.case_id, ) await self._endpoint.open() except Exception: diff --git a/assert_ai/integrations/sandbox/stock/server.py b/assert_ai/integrations/sandbox/stock/server.py index eae665b8b..727173271 100644 --- a/assert_ai/integrations/sandbox/stock/server.py +++ b/assert_ai/integrations/sandbox/stock/server.py @@ -55,17 +55,20 @@ def send_message(args: dict) -> dict: # the same files as ActionMediator and host-mode setup validation. MOCKS = MockLibrary(MOCKS.rules, cassette_dir=CASSETTE_DIR or MOCKS.cassette_dir) MEDIATOR = ActionMediator(POLICY, mocks=MOCKS, cassette_dir=CASSETTE_DIR) -TOOL_HOST = AgentHooksToolHost( - tools={ - "lookup_customer": lookup_customer, - "send_message": send_message, - }, - mediator=MEDIATOR, - agent_id="stock-sandbox-agent", - session_id="stock-sandbox-case", - case_id=CASE_ID, - framework="assert-stock-http", -) + + +def _tool_host(case_id: str | None) -> AgentHooksToolHost: + return AgentHooksToolHost( + tools={ + "lookup_customer": lookup_customer, + "send_message": send_message, + }, + mediator=MEDIATOR, + agent_id="stock-sandbox-agent", + session_id="stock-sandbox-case", + case_id=case_id, + framework="assert-stock-http", + ) class Handler(BaseHTTPRequestHandler): @@ -82,10 +85,19 @@ def do_POST(self): length = int(self.headers.get("content-length", "0") or 0) request = json.loads(self.rfile.read(length) or b"{}") user_message = str(request.get("message") or "") - - first_new_record = len(TOOL_HOST.records) - customer = TOOL_HOST.call_tool("lookup_customer", {"customer_id": "C1001"}) - delivery = TOOL_HOST.call_tool( + request_case_id = request.get("case_id") + if request_case_id is not None: + if not isinstance(request_case_id, str) or not request_case_id.strip(): + self._json(400, {"error": "case_id must be a non-empty string"}) + return + request_case_id = request_case_id.strip() + if CASE_ID and request_case_id and request_case_id != CASE_ID: + self._json(400, {"error": "request case_id does not match sandbox case"}) + return + tool_host = _tool_host(CASE_ID or request_case_id) + + customer = tool_host.call_tool("lookup_customer", {"customer_id": "C1001"}) + delivery = tool_host.call_tool( "send_message", { "recipient": "555-000-9999", @@ -93,8 +105,7 @@ def do_POST(self): "body": f"Account C1001 balance: ${customer['balance']:.2f}", }, ) - records = TOOL_HOST.records[first_new_record:] - events = [assert_tool_event(record) for record in records] + events = [assert_tool_event(record) for record in tool_host.records] # Deliberately attempt one harmless request so the network deny-and-audit # path is visible beside the tool mediation evidence. diff --git a/examples/sandbox_action_mediation/README.md b/examples/sandbox_action_mediation/README.md index 6ae72ddbf..eca559786 100644 --- a/examples/sandbox_action_mediation/README.md +++ b/examples/sandbox_action_mediation/README.md @@ -14,14 +14,16 @@ policy has selected `mock`. An optional `case_id:` on a mock rule binds that response to one ASSERT test case (prefix/suffix `*` globs are supported). ASSERT propagates the stable test -case ID into the container, mediation context, scenario state, and evidence, so -two cases with identical tool arguments can still exercise different outcomes. - -A `case_id:` is a precedence tier, not just another matcher: any case-bound -rule is considered before every rule without one, regardless of how many -`when:` arguments the unbound rule declares. Within each tier, the existing +case ID into the container, mediation context, scenario state, evidence, and the +JSON request sent to an already-running sandbox endpoint. A sandbox endpoint +receives `{"message": ..., "history": [...], "case_id": "..."}`; ordinary +`target.endpoint` requests retain the existing message/history shape. Two cases +with identical tool arguments can therefore exercise different outcomes. + +Case precedence is explicit: exact case ID, then a matching case glob, then a +rule with no case binding. Within each tier, the existing most-specific-first-then-file-order rules still apply. So a case-bound rule -matching any arguments *will* outrank a rule matching two exact arguments — +matching any arguments *will* outrank a generic rule matching two exact arguments — that is deliberate, since the whole point is letting one case diverge from otherwise-identical traffic. Keep it in mind when adding a broad `case_id:` rule to a file that already has narrow argument rules for the same tool. diff --git a/examples/sandbox_action_mediation/stock_agent/server.py b/examples/sandbox_action_mediation/stock_agent/server.py index eae665b8b..727173271 100644 --- a/examples/sandbox_action_mediation/stock_agent/server.py +++ b/examples/sandbox_action_mediation/stock_agent/server.py @@ -55,17 +55,20 @@ def send_message(args: dict) -> dict: # the same files as ActionMediator and host-mode setup validation. MOCKS = MockLibrary(MOCKS.rules, cassette_dir=CASSETTE_DIR or MOCKS.cassette_dir) MEDIATOR = ActionMediator(POLICY, mocks=MOCKS, cassette_dir=CASSETTE_DIR) -TOOL_HOST = AgentHooksToolHost( - tools={ - "lookup_customer": lookup_customer, - "send_message": send_message, - }, - mediator=MEDIATOR, - agent_id="stock-sandbox-agent", - session_id="stock-sandbox-case", - case_id=CASE_ID, - framework="assert-stock-http", -) + + +def _tool_host(case_id: str | None) -> AgentHooksToolHost: + return AgentHooksToolHost( + tools={ + "lookup_customer": lookup_customer, + "send_message": send_message, + }, + mediator=MEDIATOR, + agent_id="stock-sandbox-agent", + session_id="stock-sandbox-case", + case_id=case_id, + framework="assert-stock-http", + ) class Handler(BaseHTTPRequestHandler): @@ -82,10 +85,19 @@ def do_POST(self): length = int(self.headers.get("content-length", "0") or 0) request = json.loads(self.rfile.read(length) or b"{}") user_message = str(request.get("message") or "") - - first_new_record = len(TOOL_HOST.records) - customer = TOOL_HOST.call_tool("lookup_customer", {"customer_id": "C1001"}) - delivery = TOOL_HOST.call_tool( + request_case_id = request.get("case_id") + if request_case_id is not None: + if not isinstance(request_case_id, str) or not request_case_id.strip(): + self._json(400, {"error": "case_id must be a non-empty string"}) + return + request_case_id = request_case_id.strip() + if CASE_ID and request_case_id and request_case_id != CASE_ID: + self._json(400, {"error": "request case_id does not match sandbox case"}) + return + tool_host = _tool_host(CASE_ID or request_case_id) + + customer = tool_host.call_tool("lookup_customer", {"customer_id": "C1001"}) + delivery = tool_host.call_tool( "send_message", { "recipient": "555-000-9999", @@ -93,8 +105,7 @@ def do_POST(self): "body": f"Account C1001 balance: ${customer['balance']:.2f}", }, ) - records = TOOL_HOST.records[first_new_record:] - events = [assert_tool_event(record) for record in records] + events = [assert_tool_event(record) for record in tool_host.records] # Deliberately attempt one harmless request so the network deny-and-audit # path is visible beside the tool mediation evidence. diff --git a/tests/test_http_endpoint_events.py b/tests/test_http_endpoint_events.py index 13b3ac605..d4e9f890a 100644 --- a/tests/test_http_endpoint_events.py +++ b/tests/test_http_endpoint_events.py @@ -48,8 +48,11 @@ def post(self, endpoint, *, json, headers, allow_redirects=True): return _Response(self.payload) -async def _run(payload): - session = HTTPEndpointSession(endpoint="http://localhost:8080/chat") +async def _run(payload, *, case_id=None): + session = HTTPEndpointSession( + endpoint="http://localhost:8080/chat", + case_id=case_id, + ) client = _Client(payload) setattr(session, "_aiohttp", aiohttp) setattr(session, "_session", client) @@ -118,6 +121,21 @@ def test_endpoint_promotes_tool_events_to_judge_visible_messages(): assert client.post_allow_redirects == [False] +def test_endpoint_includes_case_id_only_when_configured(): + _, ordinary = asyncio.run(_run({"response": "ok"})) + _, correlated = asyncio.run(_run({"response": "ok"}, case_id="case-007")) + + assert ordinary.posts[0][1] == { + "message": "restore the line", + "history": [{"role": "user", "content": "restore the line"}], + } + assert correlated.posts[0][1] == { + "message": "restore the line", + "history": [{"role": "user", "content": "restore the line"}], + "case_id": "case-007", + } + + def test_endpoint_does_not_duplicate_final_assistant_event(): result, _ = asyncio.run(_run({ "response": "Done.", diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index e62ea4ec3..73d79cfa1 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -161,6 +161,61 @@ def test_case_specific_rule_beats_a_more_argument_specific_generic_rule(): assert resolved is not None and resolved.value == {"picked": "case"} +def test_exact_case_rule_beats_matching_case_glob_regardless_of_file_order(): + library = MockLibrary.from_dict({ + "mocks": [ + { + "tool": "charge_card", + "case_id": "case-*", + "response": {"picked": "glob"}, + }, + { + "tool": "charge_card", + "case_id": "case-007", + "response": {"picked": "exact"}, + }, + ] + }) + + resolved = library.resolve(MockCall("charge_card", {}, case_id="case-007")) + + assert resolved is not None and resolved.value == {"picked": "exact"} + + +def test_case_glob_beats_generic_rule_regardless_of_argument_specificity(): + library = MockLibrary.from_dict({ + "mocks": [ + { + "tool": "charge_card", + "when": {"amount": 25}, + "response": {"picked": "generic"}, + }, + { + "tool": "charge_card", + "case_id": "case-*", + "response": {"picked": "glob"}, + }, + ] + }) + + resolved = library.resolve(MockCall("charge_card", {"amount": 25}, case_id="case-007")) + + assert resolved is not None and resolved.value == {"picked": "glob"} + + +def test_case_bound_rules_do_not_match_an_uncorrelated_call(): + library = MockLibrary.from_dict({ + "mocks": [ + {"tool": "charge_card", "case_id": "case-*", "response": {"picked": "case"}}, + {"tool": "charge_card", "response": {"picked": "generic"}}, + ] + }) + + resolved = library.resolve(MockCall("charge_card", {})) + + assert resolved is not None and resolved.value == {"picked": "generic"} + + @pytest.mark.parametrize( "matcher,value,expected", [ @@ -355,6 +410,47 @@ def matches_state(self, rule): assert resolved is not None and resolved.value == {"ok": True} +def test_legacy_scenario_backend_current_state_override_still_works(): + class LegacyScenarioBackend(ScenarioBackend): + def current_state(self, scenario): + return "ready" + + library = MockLibrary.from_dict( + { + "mocks": [{ + "tool": "lookup", + "scenario": "legacy", + "when_state": "ready", + "response": {"ok": True}, + }] + }, + backends={"scenario": LegacyScenarioBackend()}, + ) + + resolved = library.resolve(MockCall("lookup", {}, case_id="case-a")) + + assert resolved is not None and resolved.value == {"ok": True} + + +def test_conflicting_context_case_ids_fail_before_mock_resolution(): + library = MockLibrary.from_dict({ + "mocks": [ + {"tool": "lookup", "case_id": "session-case", "response": {"picked": "session"}}, + {"tool": "lookup", "case_id": "legacy-case", "response": {"picked": "legacy"}}, + ] + }) + mediator = ActionMediator( + MediationPolicy({"interactions": [{"match": "lookup", "mode": "mock"}]}), + mocks=library, + ) + pre = _pre("lookup", {}) + pre["case_id"] = "legacy-case" + pre["session"] = {"id": "session", "case_id": "session-case"} + + with pytest.raises(ValueError, match="conflicting case_id"): + mediator.mediate(pre, _never_executes) + + def test_scenario_sequence_advances_then_holds(): library = MockLibrary.from_dict({ "mocks": [{ From c04c9164e9dd04e98ae0aefb29ebe5b51282d1e9 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Mon, 24 Aug 2026 18:07:41 -0400 Subject: [PATCH 6/9] fix(sandbox): reject conflicting case identity before execution --- assert_ai/integrations/sandbox/mediator.py | 1 + tests/test_sandbox_mock_setup.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/assert_ai/integrations/sandbox/mediator.py b/assert_ai/integrations/sandbox/mediator.py index 893718903..6e5e0b44f 100644 --- a/assert_ai/integrations/sandbox/mediator.py +++ b/assert_ai/integrations/sandbox/mediator.py @@ -108,6 +108,7 @@ def __init__( def mediate(self, pre_context: dict[str, Any], execute_effective: Execute) -> MediationDecision: if pre_context.get("interception_point") != "pre_tool_call": raise ValueError("ActionMediator expects a pre_tool_call context") + case_id_from_context(pre_context) tool_call = pre_context.get("tool_call") or {} name = str(tool_call.get("name") or "") args = dict(tool_call.get("args") or {}) diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index 73d79cfa1..f82e3c07f 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -451,6 +451,25 @@ def test_conflicting_context_case_ids_fail_before_mock_resolution(): mediator.mediate(pre, _never_executes) +def test_conflicting_context_case_ids_fail_before_pass_execution(): + mediator = ActionMediator( + MediationPolicy({"interactions": [{"match": "lookup", "mode": "pass"}]}) + ) + pre = _pre("lookup", {}) + pre["case_id"] = "legacy-case" + pre["session"] = {"id": "session", "case_id": "session-case"} + invoked = False + + def execute(_args): + nonlocal invoked + invoked = True + return {"ok": True} + + with pytest.raises(ValueError, match="conflicting case_id"): + mediator.mediate(pre, execute) + assert invoked is False + + def test_scenario_sequence_advances_then_holds(): library = MockLibrary.from_dict({ "mocks": [{ From 74f22062745e70a5b2c7785364d3077a23045a29 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Mon, 24 Aug 2026 18:14:24 -0400 Subject: [PATCH 7/9] fix(sandbox): preserve case proof under concurrency --- assert_ai/integrations/sandbox/evidence.py | 4 + .../integrations/sandbox/mocks/backends.py | 8 +- .../integrations/sandbox/mocks/library.py | 46 ++++--- tests/test_sandbox_mock_setup.py | 116 +++++++++++++++++- 4 files changed, 152 insertions(+), 22 deletions(-) diff --git a/assert_ai/integrations/sandbox/evidence.py b/assert_ai/integrations/sandbox/evidence.py index f008de059..1cb37cf67 100644 --- a/assert_ai/integrations/sandbox/evidence.py +++ b/assert_ai/integrations/sandbox/evidence.py @@ -45,6 +45,10 @@ def compact_evidence(record: MediationRecord) -> dict[str, Any]: case_id = case_id_from_context(record.pre_context) if case_id: evidence["case_id"] = str(case_id) + if record.decision.mock_source: + evidence["mock_source"] = record.decision.mock_source + if record.decision.replay: + evidence["replay"] = record.decision.replay if record.decision.policy_note: evidence["policy_note"] = record.decision.policy_note return evidence diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index 7c67d15ec..141f0341a 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -134,6 +134,12 @@ def _current_state_for_call(self, scenario: str, case_id: str | None) -> str: return self.current_state(scenario) return self._state.get((case_id or "", scenario), "start") + def _case_key(self, case_id: str | None) -> str: + """Use legacy shared state when an override cannot accept case identity.""" + if type(self).current_state is not ScenarioBackend.current_state: + return "" + return case_id or "" + def reset(self) -> None: self._state.clear() self._cursor.clear() @@ -178,7 +184,7 @@ def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution: raise MockBackendError(f"`responses:` for '{call.tool}' must be a non-empty list") steps = list(responses) - case_key = call.case_id or "" + case_key = self._case_key(call.case_id) key = (case_key, scenario, call.tool) index = self._cursor.get(key, 0) # The last step repeats once exhausted, so a scenario cannot run off the diff --git a/assert_ai/integrations/sandbox/mocks/library.py b/assert_ai/integrations/sandbox/mocks/library.py index 74ab249c9..6a49f012c 100644 --- a/assert_ai/integrations/sandbox/mocks/library.py +++ b/assert_ai/integrations/sandbox/mocks/library.py @@ -46,6 +46,7 @@ from __future__ import annotations import copy +import threading from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -119,6 +120,7 @@ def case_rank(rule: MockRule) -> int: ) self.backends: dict[str, MockBackend] = dict(backends or default_backends(cassette_dir)) self.cassette_dir = Path(cassette_dir) if cassette_dir else None + self._resolve_lock = threading.RLock() self._validate_backends() def _validate_backends(self) -> None: @@ -224,26 +226,30 @@ def resolve(self, call: MockCall) -> Resolution | None: about this call, and the caller should fall back to whatever the policy declares. It is never a silent empty response. """ - rule = self.find(call) - if rule is None: - return None - backend = self.backends[rule.backend] - resolution = backend.resolve(rule.raw, call) - detail = dict(resolution.detail) - detail.update({"mock_rule": rule.tool, "backend": rule.backend}) - if rule.when: - detail["matched_args"] = sorted(rule.when) - if rule.case_id: - detail["matched_case_id"] = rule.case_id - if rule.note: - detail["note"] = rule.note - return Resolution( - value=resolution.value, - mock_source=resolution.mock_source, - is_error=resolution.is_error, - state_note=resolution.state_note, - detail=detail, - ) + # Matching and state transition are one operation. A target may issue + # parallel tool calls within one case; without this boundary two calls + # can both match the same state or consume the same scenario step. + with self._resolve_lock: + rule = self.find(call) + if rule is None: + return None + backend = self.backends[rule.backend] + resolution = backend.resolve(rule.raw, call) + detail = dict(resolution.detail) + detail.update({"mock_rule": rule.tool, "backend": rule.backend}) + if rule.when: + detail["matched_args"] = sorted(rule.when) + if rule.case_id: + detail["matched_case_id"] = rule.case_id + if rule.note: + detail["note"] = rule.note + return Resolution( + value=resolution.value, + mock_source=resolution.mock_source, + is_error=resolution.is_error, + state_note=resolution.state_note, + detail=detail, + ) def reset(self) -> None: """Reset per-run state (scenario cursors) between cases.""" diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index f82e3c07f..5e9cdde1c 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -16,16 +16,19 @@ from __future__ import annotations import json +import time from concurrent.futures import ThreadPoolExecutor -from threading import Barrier +from threading import Barrier, Lock import pytest from assert_ai.integrations.sandbox.agent_hooks_context import AgentHooksContextBuilder from assert_ai.integrations.sandbox.cassettes import CassettePathError +from assert_ai.integrations.sandbox.evidence import assert_tool_event from assert_ai.integrations.sandbox.mediation_setup import MediationSetup, TargetSpec from assert_ai.integrations.sandbox.mediator import ActionMediator from assert_ai.integrations.sandbox.mocks import ( + InlineBackend, MockBackendError, MockCall, MockConfigError, @@ -35,6 +38,7 @@ from assert_ai.integrations.sandbox.mocks.matching import MatcherError, match_value from assert_ai.integrations.sandbox.policy import MediationPolicy from assert_ai.integrations.sandbox.records import MediationDecision +from assert_ai.integrations.sandbox.tool_host import AgentHooksToolHost def _pre(name, args=None): @@ -323,6 +327,54 @@ def first_status(host): assert statuses == ["resumed", "resumed"] +def test_one_library_serializes_parallel_scenario_resolution(): + class TrackingScenarioBackend(ScenarioBackend): + def __init__(self): + super().__init__() + self._activity_lock = Lock() + self._active = 0 + self.overlapped = False + + def resolve(self, rule, call): + with self._activity_lock: + self._active += 1 + self.overlapped = self.overlapped or self._active > 1 + try: + time.sleep(0.02) + return super().resolve(rule, call) + finally: + with self._activity_lock: + self._active -= 1 + + backend = TrackingScenarioBackend() + library = MockLibrary.from_dict( + { + "mocks": [{ + "tool": "retry", + "scenario": "payment", + "responses": [ + {"response": {"step": 0}}, + {"response": {"step": 1}}, + ], + }] + }, + backends={"scenario": backend}, + ) + barrier = Barrier(2) + + def resolve_once(_index): + barrier.wait() + result = library.resolve(MockCall("retry", {}, case_id="case-a")) + assert result is not None + return result.value["step"] + + with ThreadPoolExecutor(max_workers=2) as executor: + steps = sorted(executor.map(resolve_once, range(2))) + + assert steps == [0, 1] + assert backend.overlapped is False + + def test_later_read_reflects_a_mocked_write(): """If a state-changing call is mocked, a later read must agree with it.""" library = MockLibrary.from_dict({ @@ -432,6 +484,68 @@ def current_state(self, scenario): assert resolved is not None and resolved.value == {"ok": True} +def test_legacy_current_state_override_observes_its_own_transition(): + class LegacyScenarioBackend(ScenarioBackend): + def current_state(self, scenario): + return super().current_state(scenario) + + library = MockLibrary.from_dict( + { + "mocks": [ + { + "tool": "authorize", + "scenario": "legacy", + "response": {"status": "authorized"}, + "sets_state": "done", + }, + { + "tool": "status", + "scenario": "legacy", + "when_state": "done", + "response": {"state": "done"}, + }, + {"tool": "status", "response": {"state": "fallback"}}, + ] + }, + backends={ + "inline": InlineBackend(), + "scenario": LegacyScenarioBackend(), + }, + ) + + library.resolve(MockCall("authorize", {}, case_id="case-a")) + status = library.resolve(MockCall("status", {}, case_id="case-a")) + + assert status is not None and status.value == {"state": "done"} + + +def test_judge_visible_evidence_names_the_matched_case_rule(): + library = MockLibrary.from_dict({ + "mocks": [{ + "tool": "lookup", + "case_id": "case-a", + "response": {"branch": "case-a"}, + }] + }) + host = AgentHooksToolHost( + tools={"lookup": _never_executes}, + mediator=ActionMediator( + MediationPolicy({"interactions": [{"match": "lookup", "mode": "mock"}]}), + mocks=library, + ), + agent_id="agent", + session_id="session", + case_id="case-a", + ) + + host.call_tool("lookup", {}) + evidence = json.loads(assert_tool_event(host.records[0])["content"]) + + assert evidence["case_id"] == "case-a" + assert evidence["mock_source"] == "inline" + assert evidence["replay"]["matched_case_id"] == "case-a" + + def test_conflicting_context_case_ids_fail_before_mock_resolution(): library = MockLibrary.from_dict({ "mocks": [ From bcb6ad408facf1173d36f3e969bd53850f3557a9 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Mon, 24 Aug 2026 18:32:45 -0400 Subject: [PATCH 8/9] fix(sandbox): validate case evidence across legacy paths --- .../integrations/sandbox/mocks/backends.py | 10 ++++- assert_ai/integrations/sandbox/session.py | 44 ++++++++++++++++++- tests/test_sandbox_mock_setup.py | 35 +++++++++++++++ tests/test_sandbox_runtime.py | 41 +++++++++++++++++ 4 files changed, 126 insertions(+), 4 deletions(-) diff --git a/assert_ai/integrations/sandbox/mocks/backends.py b/assert_ai/integrations/sandbox/mocks/backends.py index 141f0341a..821880cf3 100644 --- a/assert_ai/integrations/sandbox/mocks/backends.py +++ b/assert_ai/integrations/sandbox/mocks/backends.py @@ -130,13 +130,19 @@ def current_state(self, scenario: str) -> str: def _current_state_for_call(self, scenario: str, case_id: str | None) -> str: """Read built-in case state without changing the extension signature.""" - if type(self).current_state is not ScenarioBackend.current_state: + if ( + type(self).current_state is not ScenarioBackend.current_state + or type(self).matches_state is not ScenarioBackend.matches_state + ): return self.current_state(scenario) return self._state.get((case_id or "", scenario), "start") def _case_key(self, case_id: str | None) -> str: """Use legacy shared state when an override cannot accept case identity.""" - if type(self).current_state is not ScenarioBackend.current_state: + if ( + type(self).current_state is not ScenarioBackend.current_state + or type(self).matches_state is not ScenarioBackend.matches_state + ): return "" return case_id or "" diff --git a/assert_ai/integrations/sandbox/session.py b/assert_ai/integrations/sandbox/session.py index 02dc41391..2d6c8338f 100644 --- a/assert_ai/integrations/sandbox/session.py +++ b/assert_ai/integrations/sandbox/session.py @@ -26,6 +26,32 @@ log = logging.getLogger(__name__) +def _validate_target_action_case( + interaction_messages: list[dict[str, Any]], + expected_case_id: str | None, +) -> None: + """Reject target-reported mediation evidence for another or unknown case.""" + if not expected_case_id: + return + for message in interaction_messages: + raw = message.get("raw") + if not isinstance(raw, dict): + continue + evidence = raw.get("action_mediation") + if not isinstance(evidence, dict): + continue + reported = evidence.get("case_id") + if not isinstance(reported, str) or not reported.strip(): + raise RuntimeError( + "target action-mediation evidence is missing the ASSERT case_id" + ) + if reported.strip() != expected_case_id: + raise RuntimeError( + "target action-mediation evidence case_id does not match the " + "ASSERT-owned sandbox case" + ) + + class SandboxedEndpointSession: """Start, use, and remove one configured sandbox for one ASSERT test case. @@ -55,6 +81,7 @@ def __init__( self._handle: SandboxHandle | None = None self._endpoint: HTTPEndpointSession | None = None self._workdir: tempfile.TemporaryDirectory[str] | None = None + self._buffered_interaction_messages: list[dict[str, Any]] = [] async def open(self) -> None: target = self.setup.target @@ -158,13 +185,26 @@ async def run_turn(self, messages: list[Message]) -> TurnResult: raise RuntimeError("sandbox session is not open") result = await self._endpoint.run_turn(messages) additions = await self.drain_pending_interaction_messages() + try: + _validate_target_action_case(result.interaction_messages, self.case_id) + except RuntimeError: + self._buffered_interaction_messages.extend(additions) + raise result.interaction_messages.extend(additions) return result async def drain_pending_interaction_messages(self) -> list[dict[str, Any]]: """Drain host-side egress evidence even when the target turn failed.""" if self._handle is None: - return [] + buffered, self._buffered_interaction_messages = ( + self._buffered_interaction_messages, + [], + ) + return buffered + buffered, self._buffered_interaction_messages = ( + self._buffered_interaction_messages, + [], + ) rows = await asyncio.to_thread(self._handle.new_egress_rows) additions: list[dict[str, Any]] = [] for row in rows: @@ -189,7 +229,7 @@ async def drain_pending_interaction_messages(self) -> list[dict[str, Any]]: "raw": {"sandbox": "network_egress"}, }, ]) - return additions + return [*buffered, *additions] @property def preserve_error_transcript(self) -> bool: diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index 5e9cdde1c..1c6ecee5f 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -519,6 +519,41 @@ def current_state(self, scenario): assert status is not None and status.value == {"state": "done"} +def test_legacy_matches_state_override_observes_its_own_transition(): + class LegacyScenarioBackend(ScenarioBackend): + def matches_state(self, rule): + return super().matches_state(rule) + + library = MockLibrary.from_dict( + { + "mocks": [ + { + "tool": "authorize", + "scenario": "legacy", + "response": {"status": "authorized"}, + "sets_state": "done", + }, + { + "tool": "status", + "scenario": "legacy", + "when_state": "done", + "response": {"state": "done"}, + }, + {"tool": "status", "response": {"state": "fallback"}}, + ] + }, + backends={ + "inline": InlineBackend(), + "scenario": LegacyScenarioBackend(), + }, + ) + + library.resolve(MockCall("authorize", {}, case_id="case-a")) + status = library.resolve(MockCall("status", {}, case_id="case-a")) + + assert status is not None and status.value == {"state": "done"} + + def test_judge_visible_evidence_names_the_matched_case_rule(): library = MockLibrary.from_dict({ "mocks": [{ diff --git a/tests/test_sandbox_runtime.py b/tests/test_sandbox_runtime.py index e7b74c3d9..f91260230 100644 --- a/tests/test_sandbox_runtime.py +++ b/tests/test_sandbox_runtime.py @@ -1030,6 +1030,47 @@ def new_egress_rows(self): assert "prompt-case-007" in json.dumps(result.interaction_messages) +@pytest.mark.parametrize( + "action_evidence,error_match", + [ + ({"case_id": "case-b", "mode": "mock"}, "does not match"), + ({"mode": "mock"}, "missing the ASSERT case_id"), + ], +) +def test_target_action_evidence_must_match_the_assert_case( + tmp_path, action_evidence, error_match +): + _, _, setup = _files(tmp_path) + setup.write_text( + "version: 1\ntarget: {kind: endpoint, url: 'http://localhost/chat'}\n" + "policy: ./policy.yaml\nmocks: ./mocks.yaml\n", + encoding="utf-8", + ) + session = SandboxedEndpointSession(setup_path=setup, case_id="case-a") + + class FakeEndpoint: + async def run_turn(self, messages): + from assert_ai.core.session import TurnResult + + return TurnResult( + text="done", + state_messages=[], + interaction_messages=[{ + "role": "tool", + "content": "{}", + "function": "lookup", + "arguments": {}, + "tool_call_id": "call-1", + "raw": {"action_mediation": action_evidence}, + }], + ) + + session._endpoint = FakeEndpoint() # type: ignore[assignment] + + with pytest.raises(RuntimeError, match=error_match): + asyncio.run(session.run_turn([Message(role="user", content="go")])) + + def test_failed_sandbox_prompt_preserves_egress_evidence(monkeypatch): """A timed-out target still produces a target_error row with egress evidence.""" class Runtime: From 5781f1eb3a1aefd678bdc3b1653316984cca4ba2 Mon Sep 17 00:00:00 2001 From: Jake Present Date: Thu, 27 Aug 2026 17:20:59 -0400 Subject: [PATCH 9/9] fix(sandbox): reject uncorrelated wildcard mocks --- .../integrations/sandbox/mocks/library.py | 8 ++++++-- tests/test_sandbox_mock_setup.py | 18 +++++++++++++++--- 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/assert_ai/integrations/sandbox/mocks/library.py b/assert_ai/integrations/sandbox/mocks/library.py index 6a49f012c..1bf0fefe0 100644 --- a/assert_ai/integrations/sandbox/mocks/library.py +++ b/assert_ai/integrations/sandbox/mocks/library.py @@ -209,8 +209,12 @@ def find(self, call: MockCall) -> MockRule | None: for rule in self.rules: if not _glob_match(rule.tool, call.tool): continue - if rule.case_id and not _glob_match(rule.case_id, call.case_id or ""): - continue + if rule.case_id is not None: + # Every case-bound rule, including the catch-all ``*`` glob, + # requires a correlated call. Matching a missing ID as an empty + # string makes ``*`` fail open and can select the wrong mock. + if not call.case_id or not _glob_match(rule.case_id, call.case_id): + continue if not match_args(rule.when, call.args): continue if rule.backend == "scenario" and isinstance(scenario_backend, ScenarioBackend): diff --git a/tests/test_sandbox_mock_setup.py b/tests/test_sandbox_mock_setup.py index 1c6ecee5f..8b60fa1e6 100644 --- a/tests/test_sandbox_mock_setup.py +++ b/tests/test_sandbox_mock_setup.py @@ -207,15 +207,17 @@ def test_case_glob_beats_generic_rule_regardless_of_argument_specificity(): assert resolved is not None and resolved.value == {"picked": "glob"} -def test_case_bound_rules_do_not_match_an_uncorrelated_call(): +@pytest.mark.parametrize("case_selector", ["case-*", "*"]) +@pytest.mark.parametrize("call_case_id", [None, ""]) +def test_case_bound_rules_do_not_match_an_uncorrelated_call(case_selector, call_case_id): library = MockLibrary.from_dict({ "mocks": [ - {"tool": "charge_card", "case_id": "case-*", "response": {"picked": "case"}}, + {"tool": "charge_card", "case_id": case_selector, "response": {"picked": "case"}}, {"tool": "charge_card", "response": {"picked": "generic"}}, ] }) - resolved = library.resolve(MockCall("charge_card", {})) + resolved = library.resolve(MockCall("charge_card", {}, case_id=call_case_id)) assert resolved is not None and resolved.value == {"picked": "generic"} @@ -955,6 +957,9 @@ def test_resolve_cli_reports_the_rule_the_named_case_will_actually_get(tmp_path, " case_id: case-a\n" " response: {branch: case-a-only}\n" " - tool: charge_card\n" + " case_id: '*'\n" + " response: {branch: any-correlated-case}\n" + " - tool: charge_card\n" " response: {branch: uncorrelated-default}\n", encoding="utf-8", ) @@ -972,6 +977,13 @@ def test_resolve_cli_reports_the_rule_the_named_case_will_actually_get(tmp_path, out = capsys.readouterr().out assert rc == 0 assert "case-a-only" in out, out + assert "any-correlated-case" not in out, out + assert "uncorrelated-default" not in out, out + + rc = cli.main(["resolve", str(setup), "charge_card", "--case-id", "case-b"]) + out = capsys.readouterr().out + assert rc == 0 + assert "any-correlated-case" in out, out assert "uncorrelated-default" not in out, out # And without a case, it must not silently pass off the default branch as