From 24c53854b2b7304004509668cf1aa49c7e5d040b Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Fri, 13 Nov 2026 17:07:43 -0800 Subject: [PATCH 1/3] feat(rubrics_verifier): opt-in judge timeout retries, a judge turn cap, and override points for staging and judge config - judge_timeout_retries (unset: a timeout fails at once, as before): an agent judge call that runs out of time is tried again; the timed-out attempt's trajectory is fetched and logged under judge_output_discrepancies (kind judge_timeout), or carried on JudgeTimeoutError. Running out of time is the poll's own deadline, or a failed task whose error code says so (`.timeout`); only a code that doesn't tell falls back to the message's "timed out after Ns". - The Claude Code agent fails a turn that ran out of its timeout with error code claude_code.timeout (was cli_error), as other harnesses name theirs. - The Claude Code agent's config takes peer_tools (default true): false leaves the SDK's peer_agents MCP server out of a turn, for an agent that must not message the task's other agents. - judge_max_turns (unset: the judge's own): sent as the judge's max_turns. - stage_trajectory: a subclass may stage the run in the judge's sandbox itself (StagedTrajectory); None, the default, loads the trajectory as before. - judge_agent_config: extra judge settings, {} by default. - Run.firings_read: the firings of every log that could be read. Co-Authored-By: Claude Opus 5.5 --- .../agents/agents/claude-code/a2a_server.py | 10 +- .../agents/claude-code/claude_bridge.py | 4 +- .../task_steps/verifiers/rubrics_verifier.py | 127 +++++++++++++- .../task_steps/verifiers/run_context.py | 5 + .../agents/claude_code/test_card_contract.py | 2 +- .../claude_code/test_claude_bridge_stdin.py | 2 +- .../claude_code/test_message_contract.py | 29 +++- ...st_rubrics_verifier_judge_timeout_retry.py | 164 ++++++++++++++++++ 8 files changed, 332 insertions(+), 11 deletions(-) create mode 100644 tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py diff --git a/src/agent_env/agents/agents/claude-code/a2a_server.py b/src/agent_env/agents/agents/claude-code/a2a_server.py index 7dc27625e..38f4775fd 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -57,6 +57,7 @@ enable, extension, ) +from agentenv_protocol.a2a_agent.peer_agents import PEER_MCP_SERVER_NAME from agentenv_protocol.transfers import Uploaded, download, upload from starlette.applications import Starlette from starlette.exceptions import HTTPException @@ -104,6 +105,8 @@ INPUTS_DIR = WORKSPACE_DIR / "inputs" _SKIPPABLE_WORKSPACE_ERRORS = (FileNotFoundError, PermissionError, NotADirectoryError, IsADirectoryError) +# The error code of a turn that ran out of its timeout_seconds, as other harnesses name theirs (``.timeout``). +TIMEOUT_CODE = "claude_code.timeout" # What agent-env keeps of a failed task's error_message. _ERROR_MESSAGE_LIMIT = 1000 # How much of a changing workspace file one read stages. @@ -137,6 +140,8 @@ class ClaudeCodeConfig(AgentConfig): max_thinking_tokens: int | None = None # Which of the CLI's own tools it may use. With any selection, its MCP servers are only the ones agent-env gives it. builtin_tools: ClaudeCodeTools | None = None + # Whether a turn gets the SDK's peer_agents MCP tools, to message the task's other agents. + peer_tools: bool = True def _rewrite_session_ids(jsonl_text: str, session_id: str) -> tuple[list[str], int]: @@ -481,7 +486,8 @@ async def _turn(self, request: TaskRequest[ClaudeCodeConfig], cfg: ClaudeCodeCon resume=resume, model=cfg.model, allowed_tools=cfg.allowed_tools, - mcp_servers=dict(request.mcp_servers), + mcp_servers={name: server for name, server in request.mcp_servers.items() + if cfg.peer_tools or name != PEER_MCP_SERVER_NAME}, system_prompt=cfg.system_prompt, output_format=cfg.output_format, timeout_s=cfg.timeout_seconds, @@ -508,7 +514,7 @@ async def _turn(self, request: TaskRequest[ClaudeCodeConfig], cfg: ClaudeCodeCon data["cli_attempts"] = result.attempts builder = TaskResult.builder() if result.error: - builder.failed("cli_error", result.error[:_ERROR_MESSAGE_LIMIT]) + builder.failed(TIMEOUT_CODE if result.timed_out else "cli_error", result.error[:_ERROR_MESSAGE_LIMIT]) else: builder.succeeded() return ( diff --git a/src/agent_env/agents/agents/claude-code/claude_bridge.py b/src/agent_env/agents/agents/claude-code/claude_bridge.py index e8fcfc651..d7070bb71 100644 --- a/src/agent_env/agents/agents/claude-code/claude_bridge.py +++ b/src/agent_env/agents/agents/claude-code/claude_bridge.py @@ -148,6 +148,8 @@ class ClaudeResult: structured_output: dict | None = None # How many `claude -p` runs the turn took; more than one means it was resumed. attempts: int = 1 + # The turn ran out of its timeout_s, rather than failing on its own. + timed_out: bool = False def _attribution_headers(project_id: str | None, task_id: str | None) -> dict[str, str]: @@ -745,7 +747,7 @@ def on_line(line: str) -> None: _kill_cli(proc) await _reap(proc) return ClaudeResult(response="", tool_call_count=tool_count, session_id=session_id, - error=f"claude CLI timed out after {timeout_s}s") + error=f"claude CLI timed out after {timeout_s}s", timed_out=True) except Exception: # Reading failed, as when an event_callback raises: the CLI mustn't keep running. _kill_cli(proc) diff --git a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index a7796c22c..4f3038852 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -317,6 +317,44 @@ def for_screenshot( per_criterion=per_criterion) +@dataclass(frozen=True) +class StagedTrajectory: + """A run a subclass staged in the agent judge's sandbox itself (``stage_trajectory``): the path the judge prompt + names as the trajectory, a note added to every judge prompt on how to read it, and what the verification entry + keeps of the staging (``trajectory_staging``).""" + + path: str + prompt_note: str + record: dict = field(default_factory=dict) + + +class JudgeTimeoutError(TimeoutError): + """An agent judge ran out of ``judge_timeout_seconds``; ``trajectory_object_url`` is what it had done by then, when + its trajectory could be fetched.""" + + def __init__(self, message: str, trajectory_object_url: Optional[str]): + super().__init__(message) + self.trajectory_object_url = trajectory_object_url + + +# A harness that ran out of its own time limit fails its task with an error code that says so (``.timeout``). +# One whose code doesn't say, or an older image, is read from its message, which harnesses phrase "timed out after Ns"; +# that phrase only, since a model API call that timed out ("The operation timed out") is not the judge running out. +_TIMEOUT_CODE = re.compile(r"(^|\.)timeout$") +_UNSPECIFIC_CODES = frozenset({None, "", "cli_error"}) +_TIMED_OUT_MESSAGE = re.compile(r"\btimed out after \d+") + + +def _judge_ran_out_of_time(tr: Any) -> bool: + """Whether a failed judge task failed by running out of its time limit: its error code first, its message only + when the code doesn't tell.""" + if tr.error_code and _TIMEOUT_CODE.search(tr.error_code): + return True + if tr.error_code not in _UNSPECIFIC_CODES: + return False + return bool(_TIMED_OUT_MESSAGE.search(" ".join(filter(None, [tr.error_message, tr.response_text])))) + + def _image_count(blocks: list[dict]) -> int: """Images among judge message blocks (per-criterion frames interleave a ``FRAME:`` text block before each).""" return sum(1 for b in blocks if b.get("type") == "image_url") @@ -463,6 +501,8 @@ def __init__( run_context_for_judge: Optional[list[str]] = None, judge_batch_size: Optional[int] = None, judge_batch_concurrency: Optional[int] = None, + judge_timeout_retries: Optional[int] = None, + judge_max_turns: Optional[int] = None, depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True, fail_task_when_not_passing: bool = False, @@ -481,6 +521,15 @@ def __init__( # judge_batch_concurrency at once), and retry a batch whose rows read as bulk verdicts. self.judge_batch_size = _validated_batch_option("judge_batch_size", judge_batch_size) self.judge_batch_concurrency = _validated_batch_option("judge_batch_concurrency", judge_batch_concurrency) + # An agent judge call (a batch's, with batches) that runs out of judge_timeout_seconds is tried again this many + # times, and the timed-out attempt's trajectory is kept; None keeps the first timeout fatal, as it always was. + if judge_timeout_retries is not None and (isinstance(judge_timeout_retries, bool) + or not isinstance(judge_timeout_retries, int) + or judge_timeout_retries < 0): + raise ValueError(f"judge_timeout_retries must be a non-negative integer, got {judge_timeout_retries!r}") + self.judge_timeout_retries = judge_timeout_retries + # The agent judge's turn cap (its agent config's max_turns); None leaves the judge's own. + self.judge_max_turns = _validated_batch_option("judge_max_turns", judge_max_turns) # Grade a multi-turn run as it stood after this (1-based) turn: the judge sees turns 1..k only, # so a per-turn env snapshot can be scored on what the agent had done by then. None grades all. self.grade_up_to_turn = _validated_turn(grade_up_to_turn) @@ -608,6 +657,10 @@ def to_dict(self) -> dict: base["judge_batch_size"] = self.judge_batch_size if self.judge_batch_concurrency is not None: base["judge_batch_concurrency"] = self.judge_batch_concurrency + if self.judge_timeout_retries is not None: + base["judge_timeout_retries"] = self.judge_timeout_retries + if self.judge_max_turns is not None: + base["judge_max_turns"] = self.judge_max_turns if self.fail_task_when_not_passing: base["fail_task_when_not_passing"] = True return base @@ -640,6 +693,8 @@ def from_dict(cls, data: dict) -> RubricsVerifierTaskStep: run_context_for_judge=data.get("run_context_for_judge"), judge_batch_size=data.get("judge_batch_size"), judge_batch_concurrency=data.get("judge_batch_concurrency"), + judge_timeout_retries=data.get("judge_timeout_retries"), + judge_max_turns=data.get("judge_max_turns"), fail_task_when_not_passing=data.get("fail_task_when_not_passing", False), ) @@ -648,6 +703,19 @@ async def judge_prompt_sections(self, context: TaskStepContext, prompt_response: per run (never for one skipped) with the prompt's stored response. None here.""" return [] + async def stage_trajectory(self, context: TaskStepContext, prompt_response: PromptResponse, + run_context: Optional[Run], sandbox: Any, + container: Optional[str]) -> Optional[StagedTrajectory]: + """An override point for an agent judge with use_trajectory: stage the run in the judge's ``sandbox`` (into + ``container`` when the judge runs in one) and say where. None here: the prompt's trajectory files are loaded + as they are.""" + return None + + def judge_agent_config(self) -> dict: + """Settings an override adds to what the agent judge is configured with (its card's agent-config extension + takes only the ones it supports). None here.""" + return {} + def _resolved_criteria(self, context: TaskStepContext) -> list[dict]: """Criteria for this run: a per-run ``step_params`` override replaces the stored list. @@ -767,6 +835,7 @@ async def _grade(self, context: TaskStepContext) -> TaskStepContext: ) auto_deployed_judge = None + staged: Optional[StagedTrajectory] = None judge_agent = None judge_a2a_url: Optional[str] = None judge_sandbox_id: Optional[str] = None @@ -821,10 +890,11 @@ async def _grade(self, context: TaskStepContext) -> TaskStepContext: judge_sandbox_type = auto_deployed_judge.sandbox_type if auto_deployed_judge else (agent.sandbox_type if self.agent_name else None) provider = build_sandbox_provider(judge_sandbox_type) if judge_sandbox_type else get_agent_sandbox_provider() sandbox = await provider.get_sandbox(judge_sandbox_id) - if multiturn: + staged = await self.stage_trajectory(context, stored_response, run_context, sandbox, judge_container) + if staged is None and multiturn: await self._load_per_turn_trajectories( sandbox, per_turn_uris, multiturn_dir, filter_to_apply, judge_container) - else: + elif staged is None: await sandbox.write_file_from_object(compacted.container_object_url, self.TRAJECTORY_CONTAINER_PATH, **into_container(judge_container)) if compacted.tool_result_files: @@ -862,7 +932,9 @@ async def _grade(self, context: TaskStepContext) -> TaskStepContext: f"verifier '{self.verifier_id}' uses an agent judge — raw trajectory " "passed through unchanged." ) - if self.use_trajectory: + if staged is not None: + trajectory_path = staged.path + elif self.use_trajectory: trajectory_path = multiturn_dir if multiturn else self.TRAJECTORY_CONTAINER_PATH else: image_blocks = compacted.image_blocks @@ -890,7 +962,9 @@ def eval_prompt_for(batch: list[dict]) -> str: run_context=run_context, extra_sections=extra_sections, ) - if self.use_trajectory and multiturn: + if staged is not None: + prompt += staged.prompt_note + elif self.use_trajectory and multiturn: prompt += self._MULTITURN_DIR_NOTE.format(n=len(per_turn_uris)) else: prompt = self._build_eval_prompt( @@ -991,6 +1065,8 @@ async def grade(batch: list[dict]) -> tuple[list[dict], int, list[dict], Optiona judge_call_fields["judge_batch_trajectory_object_urls"] = batch_trajectory_urls if grade_up_to_turn is not None: judge_call_fields["graded_up_to_turn"] = grade_up_to_turn + if staged is not None: + judge_call_fields["trajectory_staging"] = staged.record self._write_entries(context, verification_results, judge_call_fields) return context @@ -1285,7 +1361,9 @@ async def _configure_judge_a2a( "max_thinking_tokens": judge_max_thinking_tokens, "output_format": self._spec().output_format, "timeout_seconds": self.judge_timeout_seconds, + "max_turns": self.judge_max_turns, "task_id": context.metadata.get("task_id"), + **self.judge_agent_config(), } desired_config = {k: v for k, v in desired_config.items() if v is not None} config_ext = A2AAgent.find_extension(judge_agent_card, A2AAgent.EXT_AGENT_CONFIG) @@ -1293,6 +1371,9 @@ async def _configure_judge_a2a( if config_ext and desired_config: supported = (config_ext.get("params") or {}).get("methods", {}).get("set", {}).get("request", {}).get("supported", []) config_payload = {k: v for k, v in desired_config.items() if k in supported} + if unsupported := sorted(set(self.judge_agent_config()) - set(supported)): + logger.warning("Verifier '%s': the judge agent doesn't take %s; it runs without them", + self.verifier_id, unsupported) if config_payload: endpoint = (config_ext.get("params") or {}).get("endpoint", "/ext/agent-config") await protocol.post_agent_config(judge_a2a_url + endpoint, config_payload) @@ -1330,9 +1411,18 @@ async def _invoke_judge_a2a( # A sibling batch failed: stop this judge too, or it keeps working on a reused judge sandbox. await asyncio.shield(protocol.cancel_a2a_task(judge_a2a_url, task_id)) raise + except protocol.AgentUnreachableError: + raise + except TimeoutError as e: + if self.judge_timeout_retries is None: + raise + raise await self._judge_timed_out(str(e), judge_a2a_url, judge_agent_card, task_id, judge_agent) from e state = result["status"]["state"] status_msg = (result.get("status") or {}).get("message") or {} tr = protocol.TerminalResponse.from_message(status_msg) + if state == "failed" and self.judge_timeout_retries is not None and _judge_ran_out_of_time(tr): + raise await self._judge_timed_out(agent_error_text(tr.error_message, tr.response_text) or "timed out", + judge_a2a_url, judge_agent_card, task_id, judge_agent) if state == "failed": sandbox_id = getattr(judge_agent, "sandbox_id", None) logger.error( @@ -1366,6 +1456,32 @@ async def _invoke_judge_a2a( ) return {"response": tr.response_text, "trajectory_object_url": judge_trajectory_url} + async def _judge_timed_out(self, why: str, judge_a2a_url: str, judge_agent_card: dict, task_id: str, + judge_agent: Any) -> JudgeTimeoutError: + """The error for a judge call that ran out of time, carrying the judge's trajectory so far.""" + url = await self._fetch_judge_trajectory(judge_a2a_url=judge_a2a_url, judge_agent_card=judge_agent_card, + a2a_server_task_id=task_id, + sandbox_type=getattr(judge_agent, "sandbox_type", None)) + logger.warning("Verifier '%s': judge task %s ran out of its %ss (%s); its trajectory so far: %s", + self.verifier_id, task_id, self.judge_timeout_seconds, why, url) + return JudgeTimeoutError(f"Judge A2A task {task_id} ran out of its {self.judge_timeout_seconds}s: {why} " + f"(its trajectory so far: {url})", url) + + async def _invoke_judge_retrying_timeouts(self, *, discrepancy_log: list[dict], **kwargs: Any) -> dict: + """``_invoke_judge_a2a``, tried again up to ``judge_timeout_retries`` times when it runs out of time; each + timeout is logged in ``discrepancy_log`` with the trajectory the judge left.""" + retries = self.judge_timeout_retries or 0 + for attempt in range(1, retries + 2): + try: + return await self._invoke_judge_a2a(**kwargs) + except JudgeTimeoutError as e: + discrepancy_log.append({"kind": "judge_timeout", "attempt": attempt, + "timeout_seconds": self.judge_timeout_seconds, + "judge_trajectory_object_url": e.trajectory_object_url}) + if attempt > retries: + raise + raise RuntimeError("Judge timeout retry loop exited unexpectedly") + async def _fetch_judge_trajectory( self, *, @@ -1474,7 +1590,8 @@ async def _run_judge_with_output_retries( for attempt in range(1, self.DEFAULT_MAX_RETRIES + 1): if use_agent_judge: - data = await self._invoke_judge_a2a( + data = await self._invoke_judge_retrying_timeouts( + discrepancy_log=discrepancy_log, eval_prompt=prompt, judge_a2a_url=judge_a2a_url, judge_agent_card=judge_agent_card, diff --git a/src/agent_env/task_step/task_steps/verifiers/run_context.py b/src/agent_env/task_step/task_steps/verifiers/run_context.py index 7c976e172..a37624338 100644 --- a/src/agent_env/task_step/task_steps/verifiers/run_context.py +++ b/src/agent_env/task_step/task_steps/verifiers/run_context.py @@ -105,6 +105,11 @@ def firings(self, trigger_id: Optional[str] = None, *, env: Optional[str] = None return tuple(f for f in self._firings if (trigger_id is None or f.trigger_id == trigger_id) and (source is None or f.source == source)) + def firings_read(self) -> tuple[Firing, ...]: + """Every firing of the logs that could be read, oldest first, whichever couldn't (``unavailable`` names + those).""" + return self._firings + def firing(self, trigger_id: str, *, env: Optional[str] = None) -> Optional[Firing]: """The first firing of ``trigger_id``, or None if it never fired.""" return next(iter(self.firings(trigger_id, env=env)), None) diff --git a/tst/unit/agents/claude_code/test_card_contract.py b/tst/unit/agents/claude_code/test_card_contract.py index 0ace58b4f..b270f971b 100644 --- a/tst/unit/agents/claude_code/test_card_contract.py +++ b/tst/unit/agents/claude_code/test_card_contract.py @@ -68,7 +68,7 @@ def test_the_config_takes_every_field_prompt_agent_sends_that_the_cli_has(app): config = _extensions(app)["urn:agentenv:agent-config/v1"]["params"]["methods"] assert config["set"]["request"]["supported"] == [ "allowed_tools", "builtin_tools", "description", "effort", "max_thinking_tokens", "max_turns", "model", "name", - "output_format", "project_id", "role", "system_prompt", "task_id", "timeout_seconds", + "output_format", "peer_tools", "project_id", "role", "system_prompt", "task_id", "timeout_seconds", ] assert set(config["set"]["request"]["schema"]["properties"]) == set(config["set"]["request"]["supported"]) assert config["get"] == {"method": "GET", "response": {"required": ["config"]}} diff --git a/tst/unit/agents/claude_code/test_claude_bridge_stdin.py b/tst/unit/agents/claude_code/test_claude_bridge_stdin.py index 02a775bce..bafcad87d 100644 --- a/tst/unit/agents/claude_code/test_claude_bridge_stdin.py +++ b/tst/unit/agents/claude_code/test_claude_bridge_stdin.py @@ -111,7 +111,7 @@ def test_timeout_kills_process_and_returns_error(spawn_capture_hang): claude_bridge.run_claude(prompt="p", session_id="sess-1", resume=False, timeout_s=0) ) assert result.error is not None - assert "timed out" in result.error + assert "timed out" in result.error and result.timed_out is True assert result.response == "" assert spawn_capture_hang["proc"].killed is True diff --git a/tst/unit/agents/claude_code/test_message_contract.py b/tst/unit/agents/claude_code/test_message_contract.py index 22f8e434a..b0708ce59 100644 --- a/tst/unit/agents/claude_code/test_message_contract.py +++ b/tst/unit/agents/claude_code/test_message_contract.py @@ -10,7 +10,7 @@ import httpx import pytest from agentenv_protocol.a2a_agent import AGENT_CONFIG_V1, FilePart -from agentenv_protocol.a2a_agent.peer_agents import PEER_MCP_PATH +from agentenv_protocol.a2a_agent.peer_agents import PEER_MCP_PATH, PEER_MCP_SERVER_NAME from tst.unit.agents.claude_code.support import send, terminal from tst.unit.agents.claude_code.agent import a2a_server, claude_bridge, kit @@ -85,6 +85,17 @@ def test_a_cli_error_is_a_failed_task_with_the_capped_message(calls, client): ] +def test_a_turn_that_ran_out_of_its_time_says_so_in_its_error_code(calls, client): + calls.result_for = lambda kwargs: claude_bridge.ClaudeResult( + response="", tool_call_count=7, session_id=kwargs["session_id"], error="claude CLI timed out after 60s", + timed_out=True, + ) + + state, _, data = terminal(send(client, "hello")) + + assert state == "failed" and data[1]["error_code"] == a2a_server.TIMEOUT_CODE == "claude_code.timeout" + + def test_a_cli_error_keeps_the_response_text_when_there_is_one(calls, client): calls.result_for = lambda kwargs: claude_bridge.ClaudeResult( response="partial answer", tool_call_count=0, session_id=kwargs["session_id"], error="boom", @@ -332,3 +343,19 @@ def test_a_raw_store_uri_is_refused_saying_how_a_stored_object_arrives(calls, cl assert state == "failed" and text.startswith("unsupported URI scheme 's3' in s3://bucket/a.png") assert "agent-env sends an object its stores own as an https one" in text and calls == [] + + +def test_peer_tools_off_leaves_the_sdk_s_peer_agents_server_out_of_the_turn(monkeypatch, client): + sent: list[dict] = [] + + async def run_claude(**kwargs): + sent.append(kwargs["mcp_servers"]) + return claude_bridge.ClaudeResult(response="{}", tool_call_count=0, session_id=kwargs["session_id"]) + + monkeypatch.setattr(a2a_server, "run_claude", run_claude) + for peer_tools in (True, False): + r = client.post("/ext/agent-config", json={"peer_tools": peer_tools}) + assert r.status_code == 200, r.text + send(client, "grade it") + + assert PEER_MCP_SERVER_NAME in sent[0] and PEER_MCP_SERVER_NAME not in sent[1] diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py new file mode 100644 index 000000000..fc81517c8 --- /dev/null +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py @@ -0,0 +1,164 @@ +"""judge_timeout_retries: an agent judge call that runs out of judge_timeout_seconds is tried again, and the timed-out +attempt's trajectory is kept. Without it, a timeout fails the call at once, as it always has.""" + +from __future__ import annotations + +import json +import re +from types import SimpleNamespace + +import pytest + +from agent_env.a2a_agent import A2AAgent, protocol +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_steps.verifiers.rubrics_verifier import JudgeTimeoutError, RubricsVerifierTaskStep + +CRITERIA = [{"id": "c1", "title": "asked first"}] +ANSWER = json.dumps({"results": [{"id": "c1", "score": 1.0, "justification": "wake 2 call 3 asked"}]}) + + +def _verifier(**kw) -> RubricsVerifierTaskStep: + return RubricsVerifierTaskStep(id="verify", version=1, criteria=CRITERIA, prompt_id="p", use_trajectory=False, + verifier_id="vt", judge_timeout_seconds=60, **kw) + + +def _judge(monkeypatch, verifier, outcomes: list) -> list[str]: + """The judge's calls: each outcome is "poll-timeout", "failed::", or a completed answer.""" + sent: list[str] = [] + fetched: list[str] = [] + + async def send(url, **kwargs): + sent.append(kwargs["parts"][0]["text"]) + return f"task-{len(sent)}", None + + async def poll(url, task_id, **kwargs): + outcome = outcomes.pop(0) + if outcome == "poll-timeout": + raise TimeoutError(f"A2A task {task_id} did not finish in 60s") + if outcome.startswith("failed:"): + code, _, message = outcome[7:].partition(":") + return {"status": {"state": "failed", "message": {"parts": [ + {"kind": "text", "text": ""}, + {"kind": "data", "data": {"error_message": message, **({"error_code": code} if code else {})}}]}}} + return {"status": {"state": "completed", "message": {"parts": [{"kind": "text", "text": outcome}]}}} + + async def fetch(*, a2a_server_task_id, **kwargs): + fetched.append(a2a_server_task_id) + return f"mem://judge/{a2a_server_task_id}.json" + + monkeypatch.setattr(protocol, "send_a2a_message", send) + monkeypatch.setattr(protocol, "poll_a2a_task", poll) + monkeypatch.setattr(verifier, "_fetch_judge_trajectory", fetch) + return fetched + + +async def _grade(verifier): + return await verifier._run_judge_with_output_retries( + eval_prompt="grade it", criteria=CRITERIA, context=TaskStepContext(), model="m", + judge_a2a_url="http://judge.test", judge_agent_card={}, judge_agent=SimpleNamespace(sandbox_id=None), + judge_effort=None, judge_max_thinking_tokens=None, use_agent_judge=True) + + +@pytest.mark.asyncio +async def test_without_retries_a_timeout_fails_at_once_and_nothing_is_fetched(monkeypatch): + verifier = _verifier() + fetched = _judge(monkeypatch, verifier, ["poll-timeout", ANSWER]) + + with pytest.raises(TimeoutError) as raised: + await _grade(verifier) + + assert type(raised.value) is TimeoutError and fetched == [] + + +@pytest.mark.asyncio +async def test_a_timed_out_call_is_tried_again_and_its_trajectory_kept(monkeypatch): + verifier = _verifier(judge_timeout_retries=1) + fetched = _judge(monkeypatch, verifier, ["poll-timeout", ANSWER]) + + results, output_retries, discrepancies, url = await _grade(verifier) + + assert results[0]["score"] == 1.0 and output_retries == 0 and url == "mem://judge/task-2.json" + assert discrepancies == [{"kind": "judge_timeout", "attempt": 1, "timeout_seconds": 60, + "judge_trajectory_object_url": "mem://judge/task-1.json"}] + assert fetched == ["task-1", "task-2"] + + +@pytest.mark.asyncio +async def test_the_judge_s_own_timeout_code_counts_too_and_the_last_one_fails_with_its_trajectory(monkeypatch): + verifier = _verifier(judge_timeout_retries=1) + _judge(monkeypatch, verifier, ["failed:claude_code.timeout:claude CLI timed out after 60s", "poll-timeout"]) + + with pytest.raises(JudgeTimeoutError, match=re.escape("its trajectory so far: mem://judge/task-2.json")) as raised: + await _grade(verifier) + + assert raised.value.trajectory_object_url == "mem://judge/task-2.json" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [ + "failed:cli_error:claude CLI timed out after 60s", # an image from before the timeout code: its message + "failed::claude CLI timed out after 60s", # a harness that sends no code + "failed:pi.timeout:pi did not finish within 60s", # another harness's code +]) +async def test_a_timeout_the_code_or_else_the_message_tells_is_retried(monkeypatch, failure): + verifier = _verifier(judge_timeout_retries=1) + _judge(monkeypatch, verifier, [failure, ANSWER]) + + results, _, discrepancies, _ = await _grade(verifier) + + assert results[0]["score"] == 1.0 and discrepancies[0]["kind"] == "judge_timeout" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [ + "failed:cli_error:the CLI crashed", + "failed:cli_error:API Error: The operation timed out.", + "failed:pi.model_error:the model timed out after 30s", +]) +async def test_another_failure_is_not_retried(monkeypatch, failure): + verifier = _verifier(judge_timeout_retries=1) + _judge(monkeypatch, verifier, [failure, ANSWER]) + + with pytest.raises(RuntimeError, match="Judge A2A task failed"): + await _grade(verifier) + + +def test_judge_timeout_retries_and_max_turns_round_trip_and_are_left_out_unset(): + assert not {"judge_timeout_retries", "judge_max_turns"} & set(_verifier().to_dict()) + restored = RubricsVerifierTaskStep.from_dict(_verifier(judge_timeout_retries=2, judge_max_turns=300).to_dict()) + assert (restored.judge_timeout_retries, restored.judge_max_turns) == (2, 300) + for bad in (-1, True, "1", 1.5): + with pytest.raises(ValueError, match="judge_timeout_retries must be a non-negative integer"): + _verifier(judge_timeout_retries=bad) + with pytest.raises(ValueError, match="judge_max_turns must be a positive integer"): + _verifier(judge_max_turns=0) + + +@pytest.mark.asyncio +async def test_judge_agent_config_extras_go_to_a_judge_that_takes_them(monkeypatch, caplog): + posted: list[dict] = [] + + async def post(url, payload): + posted.append(payload) + + monkeypatch.setattr(protocol, "post_agent_config", post) + + class Extra(RubricsVerifierTaskStep): + def judge_agent_config(self) -> dict: + return {"builtin_tools": {"disable": ["Agent"]}, "effort_level": "high"} + + card = {"capabilities": {"extensions": [{"uri": A2AAgent.EXT_AGENT_CONFIG, "params": { + "endpoint": "/ext/agent-config", "methods": {"set": {"request": {"supported": [ + "model", "builtin_tools", "max_turns"]}}}}}]}} + verifier = Extra(id="verify", version=1, criteria=CRITERIA, prompt_id="p", verifier_id="vt", judge_max_turns=300) + await verifier._configure_judge_a2a(judge_a2a_url="http://judge.test", judge_agent_card=card, model="m", + judge_effort=None, judge_max_thinking_tokens=None, context=TaskStepContext()) + + assert posted == [{"model": "m", "max_turns": 300, "builtin_tools": {"disable": ["Agent"]}}] + assert "doesn't take ['effort_level']" in caplog.text + + posted.clear() + plain = RubricsVerifierTaskStep(id="verify", version=1, criteria=CRITERIA, prompt_id="p", verifier_id="vt") + await plain._configure_judge_a2a(judge_a2a_url="http://judge.test", judge_agent_card=card, model="m", + judge_effort=None, judge_max_thinking_tokens=None, context=TaskStepContext()) + assert posted == [{"model": "m"}] From a732ea6e553c1bb21eec64a54ad4c1fde23d21eb Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Fri, 13 Nov 2026 18:12:12 -0800 Subject: [PATCH 2/3] fix(claude-code): keep a timed-out turn's error code cli_error The verifier's timeout retry already recognises a Claude Code timeout from its message ("timed out after Ns"), so the new claude_code.timeout code isn't needed. Dropping it leaves this PR with no behaviour change for anyone who doesn't opt in. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/agents/agents/claude-code/a2a_server.py | 4 +--- .../agents/agents/claude-code/claude_bridge.py | 4 +--- .../task_steps/verifiers/rubrics_verifier.py | 6 +++--- .../agents/claude_code/test_claude_bridge_stdin.py | 2 +- tst/unit/agents/claude_code/test_message_contract.py | 11 ----------- .../test_rubrics_verifier_judge_timeout_retry.py | 8 ++++---- 6 files changed, 10 insertions(+), 25 deletions(-) diff --git a/src/agent_env/agents/agents/claude-code/a2a_server.py b/src/agent_env/agents/agents/claude-code/a2a_server.py index 38f4775fd..d01b1dce3 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -105,8 +105,6 @@ INPUTS_DIR = WORKSPACE_DIR / "inputs" _SKIPPABLE_WORKSPACE_ERRORS = (FileNotFoundError, PermissionError, NotADirectoryError, IsADirectoryError) -# The error code of a turn that ran out of its timeout_seconds, as other harnesses name theirs (``.timeout``). -TIMEOUT_CODE = "claude_code.timeout" # What agent-env keeps of a failed task's error_message. _ERROR_MESSAGE_LIMIT = 1000 # How much of a changing workspace file one read stages. @@ -514,7 +512,7 @@ async def _turn(self, request: TaskRequest[ClaudeCodeConfig], cfg: ClaudeCodeCon data["cli_attempts"] = result.attempts builder = TaskResult.builder() if result.error: - builder.failed(TIMEOUT_CODE if result.timed_out else "cli_error", result.error[:_ERROR_MESSAGE_LIMIT]) + builder.failed("cli_error", result.error[:_ERROR_MESSAGE_LIMIT]) else: builder.succeeded() return ( diff --git a/src/agent_env/agents/agents/claude-code/claude_bridge.py b/src/agent_env/agents/agents/claude-code/claude_bridge.py index d7070bb71..e8fcfc651 100644 --- a/src/agent_env/agents/agents/claude-code/claude_bridge.py +++ b/src/agent_env/agents/agents/claude-code/claude_bridge.py @@ -148,8 +148,6 @@ class ClaudeResult: structured_output: dict | None = None # How many `claude -p` runs the turn took; more than one means it was resumed. attempts: int = 1 - # The turn ran out of its timeout_s, rather than failing on its own. - timed_out: bool = False def _attribution_headers(project_id: str | None, task_id: str | None) -> dict[str, str]: @@ -747,7 +745,7 @@ def on_line(line: str) -> None: _kill_cli(proc) await _reap(proc) return ClaudeResult(response="", tool_call_count=tool_count, session_id=session_id, - error=f"claude CLI timed out after {timeout_s}s", timed_out=True) + error=f"claude CLI timed out after {timeout_s}s") except Exception: # Reading failed, as when an event_callback raises: the CLI mustn't keep running. _kill_cli(proc) diff --git a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index 4f3038852..05de8b512 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -337,9 +337,9 @@ def __init__(self, message: str, trajectory_object_url: Optional[str]): self.trajectory_object_url = trajectory_object_url -# A harness that ran out of its own time limit fails its task with an error code that says so (``.timeout``). -# One whose code doesn't say, or an older image, is read from its message, which harnesses phrase "timed out after Ns"; -# that phrase only, since a model API call that timed out ("The operation timed out") is not the judge running out. +# A harness that ran out of its own time limit says so in its error code (``.timeout``, as OpenHands and +# pydantic-ai do) or, when its code doesn't (Claude Code's ``cli_error``), in its message: "timed out after Ns". That +# phrase only, since a model API call that timed out ("The operation timed out") is not the judge running out. _TIMEOUT_CODE = re.compile(r"(^|\.)timeout$") _UNSPECIFIC_CODES = frozenset({None, "", "cli_error"}) _TIMED_OUT_MESSAGE = re.compile(r"\btimed out after \d+") diff --git a/tst/unit/agents/claude_code/test_claude_bridge_stdin.py b/tst/unit/agents/claude_code/test_claude_bridge_stdin.py index bafcad87d..02a775bce 100644 --- a/tst/unit/agents/claude_code/test_claude_bridge_stdin.py +++ b/tst/unit/agents/claude_code/test_claude_bridge_stdin.py @@ -111,7 +111,7 @@ def test_timeout_kills_process_and_returns_error(spawn_capture_hang): claude_bridge.run_claude(prompt="p", session_id="sess-1", resume=False, timeout_s=0) ) assert result.error is not None - assert "timed out" in result.error and result.timed_out is True + assert "timed out" in result.error assert result.response == "" assert spawn_capture_hang["proc"].killed is True diff --git a/tst/unit/agents/claude_code/test_message_contract.py b/tst/unit/agents/claude_code/test_message_contract.py index b0708ce59..d978096af 100644 --- a/tst/unit/agents/claude_code/test_message_contract.py +++ b/tst/unit/agents/claude_code/test_message_contract.py @@ -85,17 +85,6 @@ def test_a_cli_error_is_a_failed_task_with_the_capped_message(calls, client): ] -def test_a_turn_that_ran_out_of_its_time_says_so_in_its_error_code(calls, client): - calls.result_for = lambda kwargs: claude_bridge.ClaudeResult( - response="", tool_call_count=7, session_id=kwargs["session_id"], error="claude CLI timed out after 60s", - timed_out=True, - ) - - state, _, data = terminal(send(client, "hello")) - - assert state == "failed" and data[1]["error_code"] == a2a_server.TIMEOUT_CODE == "claude_code.timeout" - - def test_a_cli_error_keeps_the_response_text_when_there_is_one(calls, client): calls.result_for = lambda kwargs: claude_bridge.ClaudeResult( response="partial answer", tool_call_count=0, session_id=kwargs["session_id"], error="boom", diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py index fc81517c8..332719373 100644 --- a/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_judge_timeout_retry.py @@ -84,9 +84,9 @@ async def test_a_timed_out_call_is_tried_again_and_its_trajectory_kept(monkeypat @pytest.mark.asyncio -async def test_the_judge_s_own_timeout_code_counts_too_and_the_last_one_fails_with_its_trajectory(monkeypatch): +async def test_a_failed_turn_that_ran_out_of_time_counts_too_and_the_last_one_fails_with_its_trajectory(monkeypatch): verifier = _verifier(judge_timeout_retries=1) - _judge(monkeypatch, verifier, ["failed:claude_code.timeout:claude CLI timed out after 60s", "poll-timeout"]) + _judge(monkeypatch, verifier, ["failed:cli_error:claude CLI timed out after 60s", "poll-timeout"]) with pytest.raises(JudgeTimeoutError, match=re.escape("its trajectory so far: mem://judge/task-2.json")) as raised: await _grade(verifier) @@ -96,9 +96,9 @@ async def test_the_judge_s_own_timeout_code_counts_too_and_the_last_one_fails_wi @pytest.mark.asyncio @pytest.mark.parametrize("failure", [ - "failed:cli_error:claude CLI timed out after 60s", # an image from before the timeout code: its message + "failed:cli_error:claude CLI timed out after 60s", # Claude Code: its message "failed::claude CLI timed out after 60s", # a harness that sends no code - "failed:pi.timeout:pi did not finish within 60s", # another harness's code + "failed:openhands.timeout:did not finish in 60s", # a harness whose code says so ]) async def test_a_timeout_the_code_or_else_the_message_tells_is_retried(monkeypatch, failure): verifier = _verifier(judge_timeout_retries=1) From 0f3e2e7741144bb782d90b4d939cce74587d1f82 Mon Sep 17 00:00:00 2001 From: Sarah Chitty Date: Fri, 13 Nov 2026 18:30:16 -0800 Subject: [PATCH 3/3] fix(rubrics_verifier): a judge the verifier deployed is terminated when loading the trajectory fails Co-Authored-By: Claude Opus 5.5 --- .../task_steps/verifiers/rubrics_verifier.py | 61 +++++++++++-------- .../test_rubrics_verifier_staging_cleanup.py | 56 +++++++++++++++++ 2 files changed, 92 insertions(+), 25 deletions(-) create mode 100644 tst/unit/task_step/verifiers/test_rubrics_verifier_staging_cleanup.py diff --git a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index 05de8b512..66bb40169 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -345,6 +345,20 @@ def __init__(self, message: str, trajectory_object_url: Optional[str]): _TIMED_OUT_MESSAGE = re.compile(r"\btimed out after \d+") +async def _terminate_auto_deployed_judge(judge: Any) -> None: + """Terminate the judge sandbox a verifier deployed itself; None (it used a task's agent) has nothing to release.""" + if judge is None: + return + from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider + try: + provider = build_sandbox_provider(judge.sandbox_type) if judge.sandbox_type else get_agent_sandbox_provider() + sandbox = await provider.get_sandbox(judge.sandbox_id) + await sandbox.terminate() + logger.info(f"Terminated auto-deployed judge sandbox {judge.sandbox_id}") + except Exception as e: + logger.warning(f"Failed to terminate auto-deployed judge sandbox: {e}") + + def _judge_ran_out_of_time(tr: Any) -> bool: """Whether a failed judge task failed by running out of its time limit: its error code first, its message only when the code doesn't tell.""" @@ -885,21 +899,27 @@ async def _grade(self, context: TaskStepContext) -> TaskStepContext: judge_agent_card = auto_deployed_judge.agent_card judge_agent = auto_deployed_judge logger.info(f"Auto-deployed judge A2A agent '{resolved_judge_id}' at {judge_a2a_url}") - if self.use_trajectory: - from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider - judge_sandbox_type = auto_deployed_judge.sandbox_type if auto_deployed_judge else (agent.sandbox_type if self.agent_name else None) - provider = build_sandbox_provider(judge_sandbox_type) if judge_sandbox_type else get_agent_sandbox_provider() - sandbox = await provider.get_sandbox(judge_sandbox_id) - staged = await self.stage_trajectory(context, stored_response, run_context, sandbox, judge_container) - if staged is None and multiturn: - await self._load_per_turn_trajectories( - sandbox, per_turn_uris, multiturn_dir, filter_to_apply, judge_container) - elif staged is None: - await sandbox.write_file_from_object(compacted.container_object_url, self.TRAJECTORY_CONTAINER_PATH, - **into_container(judge_container)) - if compacted.tool_result_files: - await self._load_tool_result_files(sandbox, compacted.tool_result_files, judge_container) - compacted.tool_result_files = [] + # Loading the trajectory can fail (an override's staging, a transfer): the judge this step deployed is + # released then too, not only after grading. + try: + if self.use_trajectory: + from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider + judge_sandbox_type = auto_deployed_judge.sandbox_type if auto_deployed_judge else (agent.sandbox_type if self.agent_name else None) + provider = build_sandbox_provider(judge_sandbox_type) if judge_sandbox_type else get_agent_sandbox_provider() + sandbox = await provider.get_sandbox(judge_sandbox_id) + staged = await self.stage_trajectory(context, stored_response, run_context, sandbox, judge_container) + if staged is None and multiturn: + await self._load_per_turn_trajectories( + sandbox, per_turn_uris, multiturn_dir, filter_to_apply, judge_container) + elif staged is None: + await sandbox.write_file_from_object(compacted.container_object_url, self.TRAJECTORY_CONTAINER_PATH, + **into_container(judge_container)) + if compacted.tool_result_files: + await self._load_tool_result_files(sandbox, compacted.tool_result_files, judge_container) + compacted.tool_result_files = [] + except BaseException: + await _terminate_auto_deployed_judge(auto_deployed_judge) + raise try: model = overrides.get("judge_model") or self.default_model @@ -1041,16 +1061,7 @@ async def grade(batch: list[dict]) -> tuple[list[dict], int, list[dict], Optiona _restore_frame_labels(verification_results, real_by_short) _restore_frame_labels(verification_results.checks, real_by_short) finally: - if auto_deployed_judge is not None: - from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider - try: - judge_type = auto_deployed_judge.sandbox_type - provider = build_sandbox_provider(judge_type) if judge_type else get_agent_sandbox_provider() - sandbox = await provider.get_sandbox(auto_deployed_judge.sandbox_id) - await sandbox.terminate() - logger.info(f"Terminated auto-deployed judge sandbox {auto_deployed_judge.sandbox_id}") - except Exception as e: - logger.warning(f"Failed to terminate auto-deployed judge sandbox: {e}") + await _terminate_auto_deployed_judge(auto_deployed_judge) judge_call_fields: dict = { **dual_keyed("compact_trajectory_s3_uri", "compact_trajectory_object_url", compacted.compact_object_url), diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_staging_cleanup.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_staging_cleanup.py new file mode 100644 index 000000000..4de0cad10 --- /dev/null +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_staging_cleanup.py @@ -0,0 +1,56 @@ +"""A judge rubrics_verifier deployed itself is terminated when loading the trajectory fails (an override's staging +here), not only after grading.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +import agent_env.config as agent_env_config +from agent_env.a2a_agent import A2AAgent +from agent_env.providers.sandbox_providers import sandbox_provider +from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.task_steps.verifiers import rubrics_verifier as rv + + +class _StagingFails(rv.RubricsVerifierTaskStep): + async def stage_trajectory(self, *args, **kwargs): + raise ConnectionError("the write to the judge's sandbox dropped") + + +@pytest.mark.asyncio +async def test_a_failed_staging_terminates_the_judge_the_verifier_deployed(monkeypatch): + terminated: list[str] = [] + + class Sandbox: + def __init__(self, sandbox_id): + self.sandbox_id = sandbox_id + + async def terminate(self): + terminated.append(self.sandbox_id) + + class Provider: + async def get_sandbox(self, sandbox_id): + return Sandbox(sandbox_id) + + class Judge: + async def deploy(self, **kwargs): + return SimpleNamespace(a2a_url="http://judge.test", sandbox_id="sb-judge", agent_card={}, sandbox_type=None) + + config = SimpleNamespace(resolve_reference=lambda v: v, get_litellm_api_key=lambda: "k", + get_litellm_base_url=lambda: "http://llm.test", get_default_a2a_agent_id=lambda: "judge") + monkeypatch.setattr(agent_env_config, "get_config", lambda: config) + monkeypatch.setattr(A2AAgent, "get", staticmethod(lambda agent_id: Judge())) + monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", lambda: Provider()) + verifier = _StagingFails(id="verify", version=1, criteria=[{"id": "c1", "title": "asked first"}], prompt_id="ask", + use_trajectory=True, use_agent_judge=True, verifier_id="v") + monkeypatch.setattr(verifier, "_compact_trajectory", lambda *a, **k: rv._CompactedTrajectory.empty()) + ctx = TaskStepContext() + ctx.prompt_responses.append(PromptResponse(prompt_id="ask", response="done", prompt_text="book it", + agent_trajectory_object_url="mem://run/turn.json")) + + with pytest.raises(ConnectionError): + await verifier._grade(ctx) + + assert terminated == ["sb-judge"]