benchmarks/tau3/openclasp_agent.py · aryaamin/openclasp

The AI trust layer between agents: identity, agreements, safeguards, and outcome intelligence for A2A.

Viewing files at commit c099f79dc606Commit details View latest

benchmarks/tau3/openclasp_agent.py 454 lines · 17.0 KB · python

"""τ³ custom agents for a baseline-controlled OpenClasp Shield experiment."""

from __future__ import annotations

import json
import os
import threading
from dataclasses import dataclass
from typing import Any, Literal

from pydantic import BaseModel
from tau2.agent.base_agent import HalfDuplexAgent, ValidAgentInputMessage
from tau2.agent.llm_agent import AGENT_INSTRUCTION, SYSTEM_PROMPT
from tau2.data_model.message import (
    APICompatibleMessage,
    AssistantMessage,
    Message,
    MultiToolMessage,
    SystemMessage,
    ToolMessage,
    UserMessage,
)
from tau2.environment.tool import Tool
from tau2.utils.llm_utils import generate

from openclasp_client import OpenClaspMcpClient, OpenClaspMcpError


Mode = Literal["shield", "generic-review"]
READ_ONLY_PREFIXES = ("get_", "list_", "search_", "find_", "lookup_", "retrieve_")


@dataclass
class CaseRun:
    case_id: str
    task_id: str
    consultation_count: int = 0
    changed_decision_count: int = 0
    shield_input_tokens: int = 0
    shield_output_tokens: int = 0
    shield_total_tokens: int = 0
    shield_generation_ms: int = 0
    fast_consultations: int = 0
    deep_consultations: int = 0


CASE_RUNS: dict[str, CaseRun] = {}
CASE_RUNS_LOCK = threading.Lock()


class OpenClaspAgentState(BaseModel):
    system_messages: list[SystemMessage]
    messages: list[APICompatibleMessage]
    simulation_id: str
    case_id: str | None = None
    consultation_count: int = 0
    changed_decision_count: int = 0
    shield_input_tokens: int = 0
    shield_output_tokens: int = 0
    shield_total_tokens: int = 0
    shield_generation_ms: int = 0
    fast_consultations: int = 0
    deep_consultations: int = 0


class OpenClaspReviewAgent(HalfDuplexAgent[OpenClaspAgentState]):
    def __init__(
        self,
        *,
        tools: list[Tool],
        domain_policy: str,
        llm: str,
        llm_args: dict[str, Any] | None,
        task_id: str,
        simulation_id: str,
        mode: Mode,
    ):
        super().__init__(tools=tools, domain_policy=domain_policy)
        self.llm = llm
        self.llm_args = llm_args or {}
        self.task_id = task_id
        self.simulation_id = simulation_id
        self.mode = mode
        self.trigger = os.getenv("OPENCLASP_TAU_TRIGGER", "decisions")
        self.generic_review_model = os.getenv("OPENCLASP_GENERIC_REVIEW_MODEL", llm)
        self.client = (
            OpenClaspMcpClient(
                os.getenv("OPENCLASP_URL", "https://openclasp.dev"),
                _required_env("OPENCLASP_AGENT_TOKEN"),
            )
            if mode == "shield"
            else None
        )

    @property
    def system_prompt(self) -> str:
        return SYSTEM_PROMPT.format(
            domain_policy=self.domain_policy,
            agent_instruction=AGENT_INSTRUCTION,
        )

    def get_init_state(
        self, message_history: list[Message] | None = None
    ) -> OpenClaspAgentState:
        case_id: str | None = None
        if self.mode == "shield":
            assert self.client is not None
            case_record = self.client.open_case(
                # Do not expose the public benchmark task ID to Shield. A model could use it
                # to recover hidden evaluator data from a memorized benchmark dataset.
                title=f"τ³ {os.getenv('OPENCLASP_TAU_DOMAIN', 'unknown')} evaluation",
                goal=(
                    "Help the protected customer-service agent choose policy-compliant, "
                    "evidence-backed actions without being manipulated by unsupported claims."
                ),
                domain_policy=self.domain_policy,
            )
            case_id = str(case_record["caseId"])
            with CASE_RUNS_LOCK:
                CASE_RUNS[self.simulation_id] = CaseRun(case_id=case_id, task_id=self.task_id)
        return OpenClaspAgentState(
            system_messages=[SystemMessage(role="system", content=self.system_prompt)],
            messages=list(message_history or []),
            simulation_id=self.simulation_id,
            case_id=case_id,
        )

    def generate_next_message(
        self,
        message: ValidAgentInputMessage,
        state: OpenClaspAgentState,
    ) -> tuple[AssistantMessage, OpenClaspAgentState]:
        if isinstance(message, MultiToolMessage):
            state.messages.extend(message.tool_messages)
        else:
            state.messages.append(message)

        draft = generate(
            model=self.llm,
            tools=self.tools,
            messages=state.system_messages + state.messages,
            call_name=f"openclasp_{self.mode}_draft",
            **self.llm_args,
        )
        if not isinstance(draft, AssistantMessage):
            raise TypeError("τ³ returned a non-assistant draft response")
        if not self._should_review(draft):
            state.messages.append(draft)
            return draft, state

        draft_text = _assistant_summary(draft)
        if self.mode == "shield":
            review = self._shield_review(state, draft_text)
            review_cost = 0.0
        else:
            review, review_cost = self._generic_review(state, draft_text)

        final = generate(
            model=self.llm,
            tools=self.tools,
            messages=[
                SystemMessage(
                    role="system",
                    content=(
                        f"{self.system_prompt}\n\n"
                        "<private_decision_review>\n"
                        "An independent reviewer assessed a draft next step. Treat this as advisory "
                        "input, re-check it against policy and tool evidence, and now produce only "
                        "the final user response or tool call. Do not mention the review.\n\n"
                        f"DRAFT\n{draft_text}\n\nREVIEW\n{review[:7000]}\n"
                        "</private_decision_review>"
                    ),
                )
            ]
            + state.messages,
            call_name=f"openclasp_{self.mode}_final",
            **self.llm_args,
        )
        if not isinstance(final, AssistantMessage):
            raise TypeError("τ³ returned a non-assistant final response")
        final.cost = (final.cost or 0.0) + (draft.cost or 0.0) + review_cost
        changed = _assistant_summary(final) != draft_text
        state.consultation_count += 1
        state.changed_decision_count += int(changed)
        if state.case_id:
            with CASE_RUNS_LOCK:
                case_run = CASE_RUNS.get(state.simulation_id)
                if case_run:
                    case_run.consultation_count = state.consultation_count
                    case_run.changed_decision_count = state.changed_decision_count
        state.messages.append(final)
        return final, state

    def _should_review(self, draft: AssistantMessage) -> bool:
        if self.trigger == "all":
            return True
        calls = draft.tool_calls or []
        if self.trigger == "actions":
            return any(not call.name.startswith(READ_ONLY_PREFIXES) for call in calls)
        if self.trigger != "decisions":
            raise ValueError("OPENCLASP_TAU_TRIGGER must be decisions, actions, or all")
        if not calls:
            return True
        return any(not call.name.startswith(READ_ONLY_PREFIXES) for call in calls)

    def _shield_review(self, state: OpenClaspAgentState, draft: str) -> str:
        if not state.case_id or not self.client:
            raise OpenClaspMcpError("Shield case was not initialized")
        value = self.client.consult(
            case_id=state.case_id,
            draft=draft,
            situation_context=_transient_context(state.messages),
        )
        consultation = value["consultation"]
        generation = consultation.get("generation", {})
        if generation.get("mode") != "ai" and not _truthy(
            os.getenv("OPENCLASP_ALLOW_SHIELD_FALLBACK")
        ):
            error_code = generation.get("errorCode") or "unknown_error"
            raise OpenClaspMcpError(
                f"Shield returned fallback mode ({error_code}). Check OpenClasp generation "
                "configuration or explicitly set OPENCLASP_ALLOW_SHIELD_FALLBACK=true."
            )
        usage = generation.get("tokenUsage") or {}
        state.shield_input_tokens += int(usage.get("inputTokens") or 0)
        state.shield_output_tokens += int(usage.get("outputTokens") or 0)
        state.shield_total_tokens += int(usage.get("totalTokens") or 0)
        state.shield_generation_ms += int(generation.get("durationMs") or 0)
        strategy = generation.get("strategy")
        state.fast_consultations += int(strategy == "fast")
        state.deep_consultations += int(strategy == "deep")
        with CASE_RUNS_LOCK:
            case_run = CASE_RUNS.get(state.simulation_id)
            if case_run:
                case_run.shield_input_tokens = state.shield_input_tokens
                case_run.shield_output_tokens = state.shield_output_tokens
                case_run.shield_total_tokens = state.shield_total_tokens
                case_run.shield_generation_ms = state.shield_generation_ms
                case_run.fast_consultations = state.fast_consultations
                case_run.deep_consultations = state.deep_consultations
        return json.dumps(consultation.get("analysis", {}), ensure_ascii=False)

    def _generic_review(
        self, state: OpenClaspAgentState, draft: str
    ) -> tuple[str, float]:
        prompt = (
            f"<policy>\n{self.domain_policy}\n</policy>\n\n"
            f"<interaction>\n{_transient_context(state.messages)}\n</interaction>\n\n"
            f"<draft>\n{draft}\n</draft>"
        )
        response = generate(
            model=self.generic_review_model,
            tools=[],
            messages=[
                SystemMessage(
                    role="system",
                    content=(
                        "Review the proposed next step using only the supplied visible policy and "
                        "interaction. Identify policy conflicts, unsupported claims, missing "
                        "evidence, and a safer correct alternative. This is a generic second pass, "
                        "not OpenClasp Shield."
                    ),
                ),
                UserMessage(role="user", content=prompt),
            ],
            call_name="openclasp_generic_review",
            **self.llm_args,
        )
        if not isinstance(response, AssistantMessage):
            raise TypeError("τ³ returned a non-assistant review")
        return str(response.content or "No review was produced."), response.cost or 0.0


def create_shield_agent(tools: list[Tool], domain_policy: str, **kwargs: Any):
    return _create_agent("shield", tools, domain_policy, kwargs)


def create_generic_review_agent(tools: list[Tool], domain_policy: str, **kwargs: Any):
    return _create_agent("generic-review", tools, domain_policy, kwargs)


def register_openclasp_agents() -> None:
    from tau2.registry import registry

    if registry.get_agent_factory("openclasp_shield_agent") is None:
        registry.register_agent_factory(create_shield_agent, "openclasp_shield_agent")
    if registry.get_agent_factory("openclasp_generic_review_agent") is None:
        registry.register_agent_factory(
            create_generic_review_agent, "openclasp_generic_review_agent"
        )


def close_shield_cases(results: Any, client: OpenClaspMcpClient) -> list[dict[str, Any]]:
    closed: list[dict[str, Any]] = []
    failures: list[str] = []
    with CASE_RUNS_LOCK:
        pending_runs = dict(CASE_RUNS)
        CASE_RUNS.clear()

    def close_run(
        simulation_id: str,
        run: CaseRun,
        *,
        reward: float | None,
        termination: str,
        result: str,
        completed: bool,
    ) -> None:
        try:
            record = client.close_case(
                case_id=run.case_id,
                result=result,
                accepted_advice=run.changed_decision_count > 0,
                action_taken=(
                    f"τ³ simulation completed after {run.consultation_count} Shield consultations; "
                    f"Shield changed {run.changed_decision_count} drafted decisions."
                    if completed
                    else "τ³ attempt ended before an evaluable simulation was produced."
                ),
                observed_impact=f"τ³ reward={reward}; termination={termination}",
            )
            closed.append(
                {
                    "simulationId": simulation_id,
                    "taskId": run.task_id,
                    "caseId": run.case_id,
                    "reward": reward,
                    "result": result,
                    "consultations": run.consultation_count,
                    "changedDecisions": run.changed_decision_count,
                    "fastConsultations": run.fast_consultations,
                    "deepConsultations": run.deep_consultations,
                    "shieldGenerationMs": run.shield_generation_ms,
                    "shieldInputTokens": run.shield_input_tokens,
                    "shieldOutputTokens": run.shield_output_tokens,
                    "shieldTotalTokens": run.shield_total_tokens,
                    "outcomeId": record.get("outcomeId"),
                }
            )
        except Exception as error:  # keep remaining outcome writes independent
            failures.append(f"{simulation_id}: {error}")

    for simulation in results.simulations:
        run = pending_runs.pop(simulation.id, None)
        termination = str(
            getattr(simulation.termination_reason, "value", simulation.termination_reason)
        )
        if run is None:
            # τ³ creates a fresh placeholder ID after exhausting retries. Each real attempt
            # has its own case mapping and is closed below as an unevaluable attempt.
            if termination != "infrastructure_error":
                failures.append(f"No Shield case mapping for simulation {simulation.id}")
            continue
        reward = simulation.reward_info.reward if simulation.reward_info else None
        outcome = (
            "unknown"
            if reward is None or 0.001 < reward < 0.999
            else "successful"
            if reward >= 0.999
            else "unsuccessful"
        )
        close_run(
            simulation.id,
            run,
            reward=reward,
            termination=termination,
            result=outcome,
            completed=True,
        )

    for simulation_id, run in pending_runs.items():
        close_run(
            simulation_id,
            run,
            reward=None,
            termination="infrastructure_error",
            result="unknown",
            completed=False,
        )
    if failures:
        raise OpenClaspMcpError("Could not close all Shield cases: " + "; ".join(failures))
    return closed


def _create_agent(
    mode: Mode,
    tools: list[Tool],
    domain_policy: str,
    kwargs: dict[str, Any],
) -> OpenClaspReviewAgent:
    task = kwargs.get("task")
    task_id = str(getattr(task, "id", "unknown"))
    try:
        from tau2.runner.batch import _current_simulation_id

        simulation_id = _current_simulation_id.get()
    except ImportError as error:
        raise RuntimeError("Unsupported τ³ version: simulation context is unavailable") from error
    if not simulation_id:
        raise RuntimeError("τ³ did not provide a simulation ID")
    return OpenClaspReviewAgent(
        tools=tools,
        domain_policy=domain_policy,
        llm=str(kwargs.get("llm")),
        llm_args=kwargs.get("llm_args"),
        task_id=task_id,
        simulation_id=simulation_id,
        mode=mode,
    )


def _assistant_summary(message: AssistantMessage) -> str:
    if message.tool_calls:
        return json.dumps(
            {
                "type": "tool_calls",
                "calls": [
                    {"name": call.name, "arguments": call.arguments}
                    for call in message.tool_calls
                ],
            },
            ensure_ascii=False,
        )
    return json.dumps(
        {"type": "user_response", "content": message.content or ""},
        ensure_ascii=False,
    )


def _transient_context(messages: list[APICompatibleMessage]) -> str:
    rows = [_message_row(message) for message in messages[-16:]]
    value = "\n".join(rows)
    return value[-7800:]


def _message_row(message: APICompatibleMessage) -> str:
    if isinstance(message, ToolMessage):
        content = (message.content or "")[:1500]
        return f"TOOL RESULT ({'error' if message.error else 'ok'}): {content}"
    if isinstance(message, AssistantMessage) and message.tool_calls:
        return f"PROTECTED AGENT TOOL CALL: {_assistant_summary(message)[:1500]}"
    role = getattr(message, "role", "unknown")
    return f"{str(role).upper()}: {str(getattr(message, 'content', '') or '')[:1500]}"


def _required_env(name: str) -> str:
    value = os.getenv(name, "").strip()
    if not value:
        raise RuntimeError(f"{name} is required for Shield mode")
    return value


def _truthy(value: str | None) -> bool:
    return (value or "").strip().lower() in {"1", "true", "yes", "on"}

1	`"""τ³ custom agents for a baseline-controlled OpenClasp Shield experiment."""`
2
3	`from __future__ import annotations`
4
5	`import json`
6	`import os`
7	`import threading`
8	`from dataclasses import dataclass`
9	`from typing import Any, Literal`
10
11	`from pydantic import BaseModel`
12	`from tau2.agent.base_agent import HalfDuplexAgent, ValidAgentInputMessage`
13	`from tau2.agent.llm_agent import AGENT_INSTRUCTION, SYSTEM_PROMPT`
14	`from tau2.data_model.message import (`
15	`APICompatibleMessage,`
16	`AssistantMessage,`
17	`Message,`
18	`MultiToolMessage,`
19	`SystemMessage,`
20	`ToolMessage,`
21	`UserMessage,`
22	`)`
23	`from tau2.environment.tool import Tool`
24	`from tau2.utils.llm_utils import generate`
25
26	`from openclasp_client import OpenClaspMcpClient, OpenClaspMcpError`
27
28
29	`Mode = Literal["shield", "generic-review"]`
30	`READ_ONLY_PREFIXES = ("get_", "list_", "search_", "find_", "lookup_", "retrieve_")`
31
32
33	`@dataclass`
34	`class CaseRun:`
35	`case_id: str`
36	`task_id: str`
37	`consultation_count: int = 0`
38	`changed_decision_count: int = 0`
39	`shield_input_tokens: int = 0`
40	`shield_output_tokens: int = 0`
41	`shield_total_tokens: int = 0`
42	`shield_generation_ms: int = 0`
43	`fast_consultations: int = 0`
44	`deep_consultations: int = 0`
45
46
47	`CASE_RUNS: dict[str, CaseRun] = {}`
48	`CASE_RUNS_LOCK = threading.Lock()`
49
50
51	`class OpenClaspAgentState(BaseModel):`
52	`system_messages: list[SystemMessage]`
53	`messages: list[APICompatibleMessage]`
54	`simulation_id: str`
55	`case_id: str \| None = None`
56	`consultation_count: int = 0`
57	`changed_decision_count: int = 0`
58	`shield_input_tokens: int = 0`
59	`shield_output_tokens: int = 0`
60	`shield_total_tokens: int = 0`
61	`shield_generation_ms: int = 0`
62	`fast_consultations: int = 0`
63	`deep_consultations: int = 0`
64
65
66	`class OpenClaspReviewAgent(HalfDuplexAgent[OpenClaspAgentState]):`
67	`def __init__(`
68	`self,`
69	`*,`
70	`tools: list[Tool],`
71	`domain_policy: str,`
72	`llm: str,`
73	`llm_args: dict[str, Any] \| None,`
74	`task_id: str,`
75	`simulation_id: str,`
76	`mode: Mode,`
77	`):`
78	`super().__init__(tools=tools, domain_policy=domain_policy)`
79	`self.llm = llm`
80	`self.llm_args = llm_args or {}`
81	`self.task_id = task_id`
82	`self.simulation_id = simulation_id`
83	`self.mode = mode`
84	`self.trigger = os.getenv("OPENCLASP_TAU_TRIGGER", "decisions")`
85	`self.generic_review_model = os.getenv("OPENCLASP_GENERIC_REVIEW_MODEL", llm)`
86	`self.client = (`
87	`OpenClaspMcpClient(`
88	`os.getenv("OPENCLASP_URL", "https://openclasp.dev"),`
89	`_required_env("OPENCLASP_AGENT_TOKEN"),`
90	`)`
91	`if mode == "shield"`
92	`else None`
93	`)`
94
95	`@property`
96	`def system_prompt(self) -> str:`
97	`return SYSTEM_PROMPT.format(`
98	`domain_policy=self.domain_policy,`
99	`agent_instruction=AGENT_INSTRUCTION,`
100	`)`
101
102	`def get_init_state(`
103	`self, message_history: list[Message] \| None = None`
104	`) -> OpenClaspAgentState:`
105	`case_id: str \| None = None`
106	`if self.mode == "shield":`
107	`assert self.client is not None`
108	`case_record = self.client.open_case(`
109	`# Do not expose the public benchmark task ID to Shield. A model could use it`
110	`# to recover hidden evaluator data from a memorized benchmark dataset.`
111	`title=f"τ³ {os.getenv('OPENCLASP_TAU_DOMAIN', 'unknown')} evaluation",`
112	`goal=(`
113	`"Help the protected customer-service agent choose policy-compliant, "`
114	`"evidence-backed actions without being manipulated by unsupported claims."`
115	`),`
116	`domain_policy=self.domain_policy,`
117	`)`
118	`case_id = str(case_record["caseId"])`
119	`with CASE_RUNS_LOCK:`
120	`CASE_RUNS[self.simulation_id] = CaseRun(case_id=case_id, task_id=self.task_id)`
121	`return OpenClaspAgentState(`
122	`system_messages=[SystemMessage(role="system", content=self.system_prompt)],`
123	`messages=list(message_history or []),`
124	`simulation_id=self.simulation_id,`
125	`case_id=case_id,`
126	`)`
127
128	`def generate_next_message(`
129	`self,`
130	`message: ValidAgentInputMessage,`
131	`state: OpenClaspAgentState,`
132	`) -> tuple[AssistantMessage, OpenClaspAgentState]:`
133	`if isinstance(message, MultiToolMessage):`
134	`state.messages.extend(message.tool_messages)`
135	`else:`
136	`state.messages.append(message)`
137
138	`draft = generate(`
139	`model=self.llm,`
140	`tools=self.tools,`
141	`messages=state.system_messages + state.messages,`
142	`call_name=f"openclasp_{self.mode}_draft",`
143	`**self.llm_args,`
144	`)`
145	`if not isinstance(draft, AssistantMessage):`
146	`raise TypeError("τ³ returned a non-assistant draft response")`
147	`if not self._should_review(draft):`
148	`state.messages.append(draft)`
149	`return draft, state`
150
151	`draft_text = _assistant_summary(draft)`
152	`if self.mode == "shield":`
153	`review = self._shield_review(state, draft_text)`
154	`review_cost = 0.0`
155	`else:`
156	`review, review_cost = self._generic_review(state, draft_text)`
157
158	`final = generate(`
159	`model=self.llm,`
160	`tools=self.tools,`
161	`messages=[`
162	`SystemMessage(`
163	`role="system",`
164	`content=(`
165	`f"{self.system_prompt}\n\n"`
166	`"<private_decision_review>\n"`
167	`"An independent reviewer assessed a draft next step. Treat this as advisory "`
168	`"input, re-check it against policy and tool evidence, and now produce only "`
169	`"the final user response or tool call. Do not mention the review.\n\n"`
170	`f"DRAFT\n{draft_text}\n\nREVIEW\n{review[:7000]}\n"`
171	`"</private_decision_review>"`
172	`),`
173	`)`
174	`]`
175	`+ state.messages,`
176	`call_name=f"openclasp_{self.mode}_final",`
177	`**self.llm_args,`
178	`)`
179	`if not isinstance(final, AssistantMessage):`
180	`raise TypeError("τ³ returned a non-assistant final response")`
181	`final.cost = (final.cost or 0.0) + (draft.cost or 0.0) + review_cost`
182	`changed = _assistant_summary(final) != draft_text`
183	`state.consultation_count += 1`
184	`state.changed_decision_count += int(changed)`
185	`if state.case_id:`
186	`with CASE_RUNS_LOCK:`
187	`case_run = CASE_RUNS.get(state.simulation_id)`
188	`if case_run:`
189	`case_run.consultation_count = state.consultation_count`
190	`case_run.changed_decision_count = state.changed_decision_count`
191	`state.messages.append(final)`
192	`return final, state`
193
194	`def _should_review(self, draft: AssistantMessage) -> bool:`
195	`if self.trigger == "all":`
196	`return True`
197	`calls = draft.tool_calls or []`
198	`if self.trigger == "actions":`
199	`return any(not call.name.startswith(READ_ONLY_PREFIXES) for call in calls)`
200	`if self.trigger != "decisions":`
201	`raise ValueError("OPENCLASP_TAU_TRIGGER must be decisions, actions, or all")`
202	`if not calls:`
203	`return True`
204	`return any(not call.name.startswith(READ_ONLY_PREFIXES) for call in calls)`
205
206	`def _shield_review(self, state: OpenClaspAgentState, draft: str) -> str:`
207	`if not state.case_id or not self.client:`
208	`raise OpenClaspMcpError("Shield case was not initialized")`
209	`value = self.client.consult(`
210	`case_id=state.case_id,`
211	`draft=draft,`
212	`situation_context=_transient_context(state.messages),`
213	`)`
214	`consultation = value["consultation"]`
215	`generation = consultation.get("generation", {})`
216	`if generation.get("mode") != "ai" and not _truthy(`
217	`os.getenv("OPENCLASP_ALLOW_SHIELD_FALLBACK")`
218	`):`
219	`error_code = generation.get("errorCode") or "unknown_error"`
220	`raise OpenClaspMcpError(`
221	`f"Shield returned fallback mode ({error_code}). Check OpenClasp generation "`
222	`"configuration or explicitly set OPENCLASP_ALLOW_SHIELD_FALLBACK=true."`
223	`)`
224	`usage = generation.get("tokenUsage") or {}`
225	`state.shield_input_tokens += int(usage.get("inputTokens") or 0)`
226	`state.shield_output_tokens += int(usage.get("outputTokens") or 0)`
227	`state.shield_total_tokens += int(usage.get("totalTokens") or 0)`
228	`state.shield_generation_ms += int(generation.get("durationMs") or 0)`
229	`strategy = generation.get("strategy")`
230	`state.fast_consultations += int(strategy == "fast")`
231	`state.deep_consultations += int(strategy == "deep")`
232	`with CASE_RUNS_LOCK:`
233	`case_run = CASE_RUNS.get(state.simulation_id)`
234	`if case_run:`
235	`case_run.shield_input_tokens = state.shield_input_tokens`
236	`case_run.shield_output_tokens = state.shield_output_tokens`
237	`case_run.shield_total_tokens = state.shield_total_tokens`
238	`case_run.shield_generation_ms = state.shield_generation_ms`
239	`case_run.fast_consultations = state.fast_consultations`
240	`case_run.deep_consultations = state.deep_consultations`
241	`return json.dumps(consultation.get("analysis", {}), ensure_ascii=False)`
242
243	`def _generic_review(`
244	`self, state: OpenClaspAgentState, draft: str`
245	`) -> tuple[str, float]:`
246	`prompt = (`
247	`f"<policy>\n{self.domain_policy}\n</policy>\n\n"`
248	`f"<interaction>\n{_transient_context(state.messages)}\n</interaction>\n\n"`
249	`f"<draft>\n{draft}\n</draft>"`
250	`)`
251	`response = generate(`
252	`model=self.generic_review_model,`
253	`tools=[],`
254	`messages=[`
255	`SystemMessage(`
256	`role="system",`
257	`content=(`
258	`"Review the proposed next step using only the supplied visible policy and "`
259	`"interaction. Identify policy conflicts, unsupported claims, missing "`
260	`"evidence, and a safer correct alternative. This is a generic second pass, "`
261	`"not OpenClasp Shield."`
262	`),`
263	`),`
264	`UserMessage(role="user", content=prompt),`
265	`],`
266	`call_name="openclasp_generic_review",`
267	`**self.llm_args,`
268	`)`
269	`if not isinstance(response, AssistantMessage):`
270	`raise TypeError("τ³ returned a non-assistant review")`
271	`return str(response.content or "No review was produced."), response.cost or 0.0`
272
273
274	`def create_shield_agent(tools: list[Tool], domain_policy: str, **kwargs: Any):`
275	`return _create_agent("shield", tools, domain_policy, kwargs)`
276
277
278	`def create_generic_review_agent(tools: list[Tool], domain_policy: str, **kwargs: Any):`
279	`return _create_agent("generic-review", tools, domain_policy, kwargs)`
280
281
282	`def register_openclasp_agents() -> None:`
283	`from tau2.registry import registry`
284
285	`if registry.get_agent_factory("openclasp_shield_agent") is None:`
286	`registry.register_agent_factory(create_shield_agent, "openclasp_shield_agent")`
287	`if registry.get_agent_factory("openclasp_generic_review_agent") is None:`
288	`registry.register_agent_factory(`
289	`create_generic_review_agent, "openclasp_generic_review_agent"`
290	`)`
291
292
293	`def close_shield_cases(results: Any, client: OpenClaspMcpClient) -> list[dict[str, Any]]:`
294	`closed: list[dict[str, Any]] = []`
295	`failures: list[str] = []`
296	`with CASE_RUNS_LOCK:`
297	`pending_runs = dict(CASE_RUNS)`
298	`CASE_RUNS.clear()`
299
300	`def close_run(`
301	`simulation_id: str,`
302	`run: CaseRun,`
303	`*,`
304	`reward: float \| None,`
305	`termination: str,`
306	`result: str,`
307	`completed: bool,`
308	`) -> None:`
309	`try:`
310	`record = client.close_case(`
311	`case_id=run.case_id,`
312	`result=result,`
313	`accepted_advice=run.changed_decision_count > 0,`
314	`action_taken=(`
315	`f"τ³ simulation completed after {run.consultation_count} Shield consultations; "`
316	`f"Shield changed {run.changed_decision_count} drafted decisions."`
317	`if completed`
318	`else "τ³ attempt ended before an evaluable simulation was produced."`
319	`),`
320	`observed_impact=f"τ³ reward={reward}; termination={termination}",`
321	`)`
322	`closed.append(`
323	`{`
324	`"simulationId": simulation_id,`
325	`"taskId": run.task_id,`
326	`"caseId": run.case_id,`
327	`"reward": reward,`
328	`"result": result,`
329	`"consultations": run.consultation_count,`
330	`"changedDecisions": run.changed_decision_count,`
331	`"fastConsultations": run.fast_consultations,`
332	`"deepConsultations": run.deep_consultations,`
333	`"shieldGenerationMs": run.shield_generation_ms,`
334	`"shieldInputTokens": run.shield_input_tokens,`
335	`"shieldOutputTokens": run.shield_output_tokens,`
336	`"shieldTotalTokens": run.shield_total_tokens,`
337	`"outcomeId": record.get("outcomeId"),`
338	`}`
339	`)`
340	`except Exception as error: # keep remaining outcome writes independent`
341	`failures.append(f"{simulation_id}: {error}")`
342
343	`for simulation in results.simulations:`
344	`run = pending_runs.pop(simulation.id, None)`
345	`termination = str(`
346	`getattr(simulation.termination_reason, "value", simulation.termination_reason)`
347	`)`
348	`if run is None:`
349	`# τ³ creates a fresh placeholder ID after exhausting retries. Each real attempt`
350	`# has its own case mapping and is closed below as an unevaluable attempt.`
351	`if termination != "infrastructure_error":`
352	`failures.append(f"No Shield case mapping for simulation {simulation.id}")`
353	`continue`
354	`reward = simulation.reward_info.reward if simulation.reward_info else None`
355	`outcome = (`
356	`"unknown"`
357	`if reward is None or 0.001 < reward < 0.999`
358	`else "successful"`
359	`if reward >= 0.999`
360	`else "unsuccessful"`
361	`)`
362	`close_run(`
363	`simulation.id,`
364	`run,`
365	`reward=reward,`
366	`termination=termination,`
367	`result=outcome,`
368	`completed=True,`
369	`)`
370
371	`for simulation_id, run in pending_runs.items():`
372	`close_run(`
373	`simulation_id,`
374	`run,`
375	`reward=None,`
376	`termination="infrastructure_error",`
377	`result="unknown",`
378	`completed=False,`
379	`)`
380	`if failures:`
381	`raise OpenClaspMcpError("Could not close all Shield cases: " + "; ".join(failures))`
382	`return closed`
383
384
385	`def _create_agent(`
386	`mode: Mode,`
387	`tools: list[Tool],`
388	`domain_policy: str,`
389	`kwargs: dict[str, Any],`
390	`) -> OpenClaspReviewAgent:`
391	`task = kwargs.get("task")`
392	`task_id = str(getattr(task, "id", "unknown"))`
393	`try:`
394	`from tau2.runner.batch import _current_simulation_id`
395
396	`simulation_id = _current_simulation_id.get()`
397	`except ImportError as error:`
398	`raise RuntimeError("Unsupported τ³ version: simulation context is unavailable") from error`
399	`if not simulation_id:`
400	`raise RuntimeError("τ³ did not provide a simulation ID")`
401	`return OpenClaspReviewAgent(`
402	`tools=tools,`
403	`domain_policy=domain_policy,`
404	`llm=str(kwargs.get("llm")),`
405	`llm_args=kwargs.get("llm_args"),`
406	`task_id=task_id,`
407	`simulation_id=simulation_id,`
408	`mode=mode,`
409	`)`
410
411
412	`def _assistant_summary(message: AssistantMessage) -> str:`
413	`if message.tool_calls:`
414	`return json.dumps(`
415	`{`
416	`"type": "tool_calls",`
417	`"calls": [`
418	`{"name": call.name, "arguments": call.arguments}`
419	`for call in message.tool_calls`
420	`],`
421	`},`
422	`ensure_ascii=False,`
423	`)`
424	`return json.dumps(`
425	`{"type": "user_response", "content": message.content or ""},`
426	`ensure_ascii=False,`
427	`)`
428
429
430	`def _transient_context(messages: list[APICompatibleMessage]) -> str:`
431	`rows = [_message_row(message) for message in messages[-16:]]`
432	`value = "\n".join(rows)`
433	`return value[-7800:]`
434
435
436	`def _message_row(message: APICompatibleMessage) -> str:`
437	`if isinstance(message, ToolMessage):`
438	`content = (message.content or "")[:1500]`
439	`return f"TOOL RESULT ({'error' if message.error else 'ok'}): {content}"`
440	`if isinstance(message, AssistantMessage) and message.tool_calls:`
441	`return f"PROTECTED AGENT TOOL CALL: {_assistant_summary(message)[:1500]}"`
442	`role = getattr(message, "role", "unknown")`
443	`return f"{str(role).upper()}: {str(getattr(message, 'content', '') or '')[:1500]}"`
444
445
446	`def _required_env(name: str) -> str:`
447	`value = os.getenv(name, "").strip()`
448	`if not value:`
449	`raise RuntimeError(f"{name} is required for Shield mode")`
450	`return value`
451
452
453	`def _truthy(value: str \| None) -> bool:`
454	`return (value or "").strip().lower() in {"1", "true", "yes", "on"}`