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..5df5029ca 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -137,6 +137,9 @@ 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 + # What a turn cut short by a transient API error is resumed with; unset, the bridge asks the model to finish the + # request it was given, which suits a task but not a simulated user. + resume_prompt: str | None = None def _rewrite_session_ids(jsonl_text: str, session_id: str) -> tuple[list[str], int]: @@ -492,6 +495,7 @@ async def _turn(self, request: TaskRequest[ClaudeCodeConfig], cfg: ClaudeCodeCon max_turns=cfg.max_turns, max_thinking_tokens=cfg.max_thinking_tokens, builtin_tools=cfg.builtin_tools, + resume_prompt=cfg.resume_prompt, event_callback=request.trajectory.append, ) except Exception as e: 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 cecd8286f..a17e3cc70 100644 --- a/src/agent_env/agents/agents/claude-code/claude_bridge.py +++ b/src/agent_env/agents/agents/claude-code/claude_bridge.py @@ -415,10 +415,12 @@ async def run_claude( max_thinking_tokens: int | None = None, builtin_tools: ClaudeCodeTools | None = None, event_callback: Callable[[dict], None] | None = None, + resume_prompt: str | None = None, ) -> ClaudeResult: """Run one turn, resuming the CLI session when a transient upstream failure cuts it short; a resume spends what is left of ``timeout_s``, so the turn never outlives its caller's deadline. ``CLAUDE_API_ERROR_RETRIES=0`` turns - resumes off. ``event_callback`` gets each stdout event as the CLI prints it; an attempt's trailing `result` is + resumes off. The resume is prompted with ``resume_prompt``, else ``_RESUME_PROMPT``, which asks the model to + finish the request it was given: a caller whose model isn't doing a task, such as a simulated user, sets its own. ``event_callback`` gets each stdout event as the CLI prints it; an attempt's trailing `result` is held until it's known not to be resumed.""" deadline = time.monotonic() + timeout_s retries_left = _configured_api_error_retries() @@ -468,7 +470,7 @@ async def run_claude( trajectory.extend(_without_terminal_result(result.trajectory)) first_failure = first_failure or result.error_detail retries_left -= 1 - attempt_prompt, attempt_resume, attempt_session = _RESUME_PROMPT, True, result.session_id + attempt_prompt, attempt_resume, attempt_session = resume_prompt or _RESUME_PROMPT, True, result.session_id backoff = _RETRY_BACKOFF_S[min(attempts - 1, len(_RETRY_BACKOFF_S) - 1)] logger.warning(f"claude CLI turn cut short ({_redact(failure)!r}); resuming session {attempt_session} " f"in {backoff}s with {int(budget_left)}s of the original {timeout_s}s budget left") diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index c11fb716b..9e8026628 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -115,6 +115,15 @@ def _file_uris(parts: list[dict]) -> list[str]: # The decide body that asks a user agent to record nothing when no agent trigger fires. _QUIET_DECIDE = ("turn", "solver_message", "context_id", "env_triggers", "record_unfired") +# What a user-sim turn cut short by a transient API error is resumed with, for an agent that takes `resume_prompt`. +# An agent's own resume prompt asks it to finish the request it was given; a user-sim's request is the agent's reply, +# so that prompt has it carry on as the agent. +USER_SIM_RESUME_PROMPT = ( + "Your previous turn was interrupted by a transient API error before you sent your message. You are still the " + "user. Write the user's next message, as you were about to; the assistant's last reply above is what you are " + "answering." +) + _DEFAULT_USER_SIM_OUTPUT_FORMAT: dict[str, Any] = { "type": "json_schema", "schema": { @@ -503,7 +512,8 @@ async def _converse( if is_user_sim: user_output_format = self.user_output_format or _DEFAULT_USER_SIM_OUTPUT_FORMAT user_model = self._resolve_user_model(context) - user_config = {k: v for k, v in {"output_format": user_output_format, "model": user_model}.items() if v is not None} + user_config = {k: v for k, v in {"output_format": user_output_format, "model": user_model, + "resume_prompt": USER_SIM_RESUME_PROMPT}.items() if v is not None} negotiated = A2AAgent.negotiate_agent_config(user_sim.a2a_card or {}, user_config) if negotiated: await protocol.post_agent_config(user_url + negotiated.endpoint, negotiated.fields) diff --git a/tst/unit/agents/claude_code/test_card_contract.py b/tst/unit/agents/claude_code/test_card_contract.py index 0ace58b4f..95b15f08e 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", "project_id", "resume_prompt", "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_retry.py b/tst/unit/agents/claude_code/test_claude_bridge_retry.py index 35e1999ea..4ff6210f3 100644 --- a/tst/unit/agents/claude_code/test_claude_bridge_retry.py +++ b/tst/unit/agents/claude_code/test_claude_bridge_retry.py @@ -61,6 +61,18 @@ def _write_session(scripted, session_id="sess-1"): (scripted["session_dir"] / f"{session_id}.jsonl").write_text('{"type":"assistant"}\n') +def test_a_resume_prompt_given_replaces_the_default_one(scripted): + _write_session(scripted) + scripted["queue"].append((_api_error_stdout("API Error: The operation timed out.", tool_uses=1), 1)) + scripted["queue"].append((_ok_stdout("resumed"), 0)) + + result = asyncio.run(claude_bridge.run_claude(prompt="p", session_id="sess-1", resume=False, + resume_prompt="stay the user")) + + assert result.attempts == 2 + assert scripted["procs"][1].stdin.data == b"stay the user" + + def test_success_makes_exactly_one_invocation(scripted): scripted["queue"].append((_ok_stdout(), 0)) result = asyncio.run(claude_bridge.run_claude(prompt="p", session_id="sess-1", resume=False)) diff --git a/tst/unit/task_step/test_prompt_agent_user_sim_resume.py b/tst/unit/task_step/test_prompt_agent_user_sim_resume.py new file mode 100644 index 000000000..df1862746 --- /dev/null +++ b/tst/unit/task_step/test_prompt_agent_user_sim_resume.py @@ -0,0 +1,74 @@ +"""A user simulator is sent the resume prompt that keeps it the user when its card takes one: an agent resumes a +turn cut short by asking its model to finish the request, and a user simulator's request is the agent's reply.""" + +from unittest.mock import MagicMock + +import pytest +from a2a.types import TaskState + +from agent_env.a2a_agent import A2AAgent, protocol +from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.task_steps import prompt_agent as pa_mod +from agent_env.task_step.task_steps.prompt_agent import USER_SIM_RESUME_PROMPT, PromptAgentTaskStep + + +def _card(*supported: str) -> dict: + return {"capabilities": {"extensions": [{"uri": A2AAgent.EXT_AGENT_CONFIG, "params": { + "endpoint": "/ext/agent-config", "methods": {"set": {"request": {"supported": list(supported)}}}}}]}} + + +@pytest.fixture +def posted(monkeypatch): + """Mocks the conversation's HTTP and store boundary, and collects each agent config posted, by URL.""" + cfg = MagicMock() + cfg.get_model_params.return_value = {} + cfg.get_model_for_role.return_value = None + monkeypatch.setattr(pa_mod, "get_config", lambda: cfg) + for fn in ("create_conversation", "add_a2a_task", "complete_a2a_task", "mark_closed", "get_conversation"): + monkeypatch.setattr(pa_mod.conversation_store, fn, lambda *a, **kw: None) + configs = {} + + async def post_agent_config(url, fields): + configs[url] = fields + + async def send(url, parts, message_id, context_id, timeout): + return f"task-{message_id}", None + + async def poll(url, task_id, timeout, interval, **_): + text = '{"message": "thanks", "done": true}' if "user-sim" in url else "done" + return {"status": {"state": TaskState.completed, "message": {"parts": [{"kind": "text", "text": text}]}}} + + monkeypatch.setattr(protocol, "post_agent_config", post_agent_config) + monkeypatch.setattr(pa_mod.protocol, "send_a2a_message", send) + monkeypatch.setattr(pa_mod.protocol, "poll_a2a_task", poll) + return configs + + +def _conversation(solver_card: dict, user_card: dict) -> TaskStepContext: + context = TaskStepContext(instance_id="ti-1") + context.deployed_agents.append(DeployedAgent(agent_name="solver", api_url="http://solver", + a2a_url="http://solver", a2a_card=solver_card)) + context.deployed_agents.append(DeployedAgent(agent_name="usersim", api_url="http://user-sim", + a2a_url="http://user-sim", a2a_card=user_card)) + return context + + +def _step() -> PromptAgentTaskStep: + return PromptAgentTaskStep(id="solve", version=None, prompt="hi", agent_name="solver", model="task-solver", + max_conversation_turns=2, user_agent_name="usersim", user_model="task-user") + + +@pytest.mark.asyncio +async def test_a_user_simulator_whose_card_takes_a_resume_prompt_is_sent_the_one_that_keeps_it_the_user(posted): + card = _card("model", "output_format", "resume_prompt") + await _step().execute(_conversation(card, card)) + + assert posted["http://user-sim/ext/agent-config"]["resume_prompt"] == USER_SIM_RESUME_PROMPT + assert "resume_prompt" not in posted["http://solver/ext/agent-config"] + + +@pytest.mark.asyncio +async def test_a_user_simulator_whose_card_takes_no_resume_prompt_is_sent_none(posted): + await _step().execute(_conversation(_card("model"), _card("model", "output_format"))) + + assert "resume_prompt" not in posted["http://user-sim/ext/agent-config"]