feat(rubrics_verifier): opt-in judge timeout retries, a judge turn cap, and override points for staging and judge config

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
@@ -137,6 +138,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 +484,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,

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)

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"]}}

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
@@ -332,3 +332,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]

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:<error code>:<message>", 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_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: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)
+
+    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",   # Claude Code: its message
+    "failed::claude CLI timed out after 60s",            # a harness that sends no 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)
+    _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"}]

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"]