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 6df07fcdd05cCommit details View latest

benchmarks/tau3/openclasp_agent.py 399 lines · 15.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


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


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