fix(sandbox): correlate mediation by test case

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(

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."""
 
@@ -27,12 +54,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]:

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)

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,6 +42,13 @@ 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

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,
         )
 

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
@@ -107,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 {})
@@ -155,7 +157,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

assert_ai/integrations/sandbox/mocks/backends.py

@@ -113,33 +113,64 @@ 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")
+        """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
+            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
+            or type(self).matches_state is not ScenarioBackend.matches_state
+        ):
+            return ""
+        return case_id or ""
 
     def reset(self) -> None:
         self._state.clear()
         self._cursor.clear()
 
     def matches_state(self, rule: Mapping[str, Any]) -> bool:
-        """Whether a rule's `when_state` guard holds right now."""
+        """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 "")
         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_for_call(scenario, call.case_id) == str(want)
+
     def resolve(self, rule: Mapping[str, Any], call: MockCall) -> Resolution:
         scenario = str(rule.get("scenario") or "")
         if not scenario:
@@ -159,17 +190,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 = 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
         # 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_for_call(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_for_call(scenario, call.case_id)
 
         is_error = "error" in step
         payload = step.get("error") if is_error else step.get("response")