From 8074bdbfdde7087da9606aad226f5f63595b4c2d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 19 Oct 2026 14:56:58 -0700 Subject: [PATCH 01/39] feat(task): play an outside agent through a stand-in agent in its place An episode's outside agent is now played through a model-less stand-in agent that the task's own deploy_agent or install_agent deploys in its place, so the steps that load into, run in, collect from or check the agent's sandbox run unchanged, and the loop's bash runs there as the agent's user in its folder. A task that uses no agent sandbox gets the same stand-in served in-process. Co-Authored-By: Claude Opus 5.5 --- .../agentenv_protocol/a2a_agent/__init__.py | 22 ++ .../agentenv_protocol/a2a_agent/extensions.py | 88 ++++++ .../agentenv_protocol/a2a_agent/stand_in.py | 214 +++++++++++++++ .../tests/test_a2a_stand_in.py | 185 +++++++++++++ src/agent_env/a2a_agent/a2a_agent.py | 3 +- src/agent_env/a2a_agent/stand_in/Dockerfile | 13 + src/agent_env/a2a_agent/stand_in/__init__.py | 67 +++++ .../a2a_agent/stand_in/a2a_server.py | 37 +++ src/agent_env/task/episode/rules.py | 57 ++-- src/agent_env/task/episode/run.py | 250 +++++++++++++----- src/agent_env/task/episode/spec.py | 38 ++- src/agent_env/task/episode/stand_in.py | 89 +++++++ src/agent_env/task/episode/steps.py | 91 +++---- .../task_step/task_steps/install_agent.py | 45 ++++ .../agent/stand_in_modal_vm_test.py | 81 ++++++ .../env/envs/server_env_local_test.py | 51 +++- .../env/envs/server_env_modal_test.py | 2 +- tst/unit/a2a_agent/stand_in_test.py | 38 +++ .../test_reachability_contributions.py | 2 +- tst/unit/store/aws/test_shared_credentials.py | 12 +- tst/unit/task/episode/conftest.py | 12 + tst/unit/task/episode/run_test.py | 111 +++++++- tst/unit/task/episode/spec_test.py | 126 +++++---- tst/unit/task_step/install_agent_test.py | 80 ++++++ 24 files changed, 1475 insertions(+), 239 deletions(-) create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py create mode 100644 packages/agentenv-protocol/tests/test_a2a_stand_in.py create mode 100644 src/agent_env/a2a_agent/stand_in/Dockerfile create mode 100644 src/agent_env/a2a_agent/stand_in/__init__.py create mode 100644 src/agent_env/a2a_agent/stand_in/a2a_server.py create mode 100644 src/agent_env/task/episode/stand_in.py create mode 100644 tst/integration/agent/stand_in_modal_vm_test.py create mode 100644 tst/unit/a2a_agent/stand_in_test.py create mode 100644 tst/unit/task/episode/conftest.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index 93141797a..605fe39d1 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -5,6 +5,7 @@ from .extensions import ( AGENT_CONFIG_V1, ATTRIBUTION_PROBE_V1, + DRIVER_V1, INSTALL_V1, MCP_CONFIG_V1, PEER_AGENTS_V1, @@ -19,6 +20,16 @@ ContextEventsTrajectoryRequest, ContextObjectTrajectoryRequest, ContextTrajectoryRequest, + DriverClaimRequest, + DriverExecRequest, + DriverExecResponse, + DriverNextRequest, + DriverNextResponse, + DriverOutputRequest, + DriverOutputResponse, + DriverReplyRequest, + DriverSkill, + DriverTurn, ExtensionActivation, ExtensionConfiguration, ExtensionDefinition, @@ -117,6 +128,7 @@ __all__ = [ "AGENT_CONFIG_V1", "ATTRIBUTION_PROBE_V1", + "DRIVER_V1", "INSTALL_V1", "MCP_CONFIG_V1", "PEER_AGENTS_V1", @@ -139,6 +151,16 @@ "ContextEventsTrajectoryRequest", "ContextObjectTrajectoryRequest", "ContextTrajectoryRequest", + "DriverClaimRequest", + "DriverExecRequest", + "DriverExecResponse", + "DriverNextRequest", + "DriverNextResponse", + "DriverOutputRequest", + "DriverOutputResponse", + "DriverReplyRequest", + "DriverSkill", + "DriverTurn", "DataPart", "DefaultExtensionHandlers", "ExtensionActivation", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index a35c19bbf..82df2bd90 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -287,6 +287,71 @@ class TriggerDecideRequest(ExtensionRequest): env_triggers: dict[str, dict[str, Any]] | None = None +DRIVER_TOKEN_MIN_LENGTH = 22 + + +class DriverClaimRequest(ExtensionRequest): + token: str = Field(min_length=DRIVER_TOKEN_MIN_LENGTH) + + +class DriverNextRequest(ExtensionRequest): + token: str + after: int = Field(default=0, ge=0) + + +class DriverSkill(ExtensionResponse): + name: str + description: str + path: str | None = None + + +class DriverTurn(ExtensionResponse): + task_id: str + context_id: str + text: str + system_prompt: str | None = None + skills: list[DriverSkill] = Field(default_factory=list) + + +class DriverNextResponse(ExtensionResponse): + version: int + turns: list[DriverTurn] + + +class DriverReplyRequest(ExtensionRequest): + token: str + task_id: str + text: str | None = None + error: str | None = None + trajectory: list[dict[str, Any]] = Field(default_factory=list) + tool_call_count: int = Field(default=0, ge=0) + + @model_validator(mode="after") + def _text_or_error(self) -> DriverReplyRequest: + if (self.text is None) == (self.error is None): + raise ValueError("a reply carries exactly one of text and error") + return self + + +class DriverExecRequest(ExtensionRequest): + token: str + command: str + + +class DriverExecResponse(ExtensionResponse): + exec_id: str + + +class DriverOutputRequest(ExtensionRequest): + token: str + exec_id: str + + +class DriverOutputResponse(ExtensionResponse): + done: bool + output: str | None = None + + def _tuple(values: Iterable[str]) -> tuple[str, ...]: return tuple(dict.fromkeys(values)) @@ -1167,6 +1232,28 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: configuration_validator=_configure_install, ) +DRIVER_V1 = ExtensionDefinition( + uri="urn:agentenv:driver/v1", + description="Hand the agent's turns to a driver outside its sandbox, and run its commands where the agent works.", + endpoint="/ext/driver", + core_operations={ + "claim": _op("claim", "POST", "/ext/driver/claim", RUNTIME, request_model=DriverClaimRequest), + "next": _op( + "next", "POST", "/ext/driver/next", RUNTIME, + request_model=DriverNextRequest, response_model=DriverNextResponse, + ), + "reply": _op("reply", "POST", "/ext/driver/reply", RUNTIME, request_model=DriverReplyRequest), + "exec": _op( + "exec", "POST", "/ext/driver/exec", RUNTIME, + request_model=DriverExecRequest, response_model=DriverExecResponse, + ), + "output": _op( + "output", "POST", "/ext/driver/output", RUNTIME, + request_model=DriverOutputRequest, response_model=DriverOutputResponse, + ), + }, +) + STANDARD_EXTENSIONS: Mapping[str, ExtensionDefinition] = MappingProxyType( { extension.uri: extension @@ -1180,6 +1267,7 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: ATTRIBUTION_PROBE_V1, TRIGGERS_V1, INSTALL_V1, + DRIVER_V1, ) } ) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py new file mode 100644 index 000000000..6d5b5eb5c --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py @@ -0,0 +1,214 @@ +"""The stand-in: an agent with no model, played by a driver outside its sandbox. + +Each message it gets is an open turn until the driver replies to it. Through ``urn:agentenv:driver/v1`` the driver +reads the open turns, replies to them and runs commands where the stand-in works, as it; its first ``claim`` sets the +token every later call carries. A turn its sender cancels closes without a reply. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import hmac +import os +import signal +import uuid +from dataclasses import dataclass, field +from pathlib import Path + +from starlette.exceptions import HTTPException + +from .extensions import ( + DRIVER_V1, + MCP_CONFIG_V1, + SKILL_CONFIG_V1, + TRAJECTORY_V1, + BundleSkillRequest, + DriverClaimRequest, + DriverExecRequest, + DriverExecResponse, + DriverNextRequest, + DriverNextResponse, + DriverOutputRequest, + DriverOutputResponse, + DriverReplyRequest, + DriverSkill, + DriverTurn, + InlineSkillRequest, + extension, +) +from .framework import AgentEnvAgent, AgentIdentity, a2a_agent +from .tasks.v1 import AgentConfig, TaskRequest, TaskResult, TextPart, Usage +from ..transfers import download + +LONG_POLL_SECONDS = 25.0 +COMMAND_TIMEOUT_SECONDS = 3600 +SKILLS_DIR = Path.home() / "skills" +TRAJECTORY_FORMAT = "agentenv-stand-in/v1" + + +class StandInConfig(AgentConfig): + system_prompt: str | None = None + + +@dataclass +class _Turn: + task_id: str + context_id: str + text: str + system_prompt: str | None + skills: list[DriverSkill] + reply: asyncio.Future[DriverReplyRequest] = field( + default_factory=lambda: asyncio.get_running_loop().create_future() + ) + + +IDENTITY = AgentIdentity( + name="agentenv-stand-in", + description="Stands in for an agent played from outside AgentEnv: a driver answers its turns.", + version="1.0.0", +) +EXTENSIONS = (MCP_CONFIG_V1, TRAJECTORY_V1, DRIVER_V1) + + +@a2a_agent(identity=IDENTITY, config=StandInConfig, extensions=EXTENSIONS) +class StandInAgent(AgentEnvAgent): + """``shell=False`` is a stand-in for an agent that works in no sandbox: it keeps no skills and runs no commands. + ``long_poll_seconds`` bounds how long ``next`` and ``output`` wait for something to report.""" + + def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): + self._shell = shell + self._long_poll_seconds = long_poll_seconds + self._token: str | None = None + self._turns: dict[str, _Turn] = {} + self._version = 0 + self._changed: asyncio.Future[None] | None = None + self._commands: dict[str, asyncio.Task[str]] = {} + + async def run(self, request: TaskRequest[StandInConfig]) -> TaskResult: + turn = _Turn( + task_id=request.task_id, + context_id=request.context_id, + text="\n".join(part.text for part in request.parts if isinstance(part, TextPart)), + system_prompt=request.config.system_prompt, + skills=[DriverSkill(name=skill["name"], description=skill.get("description") or "", + path=str(SKILLS_DIR / skill["name"]) if self._shell else None) + for skill in request.skills], + ) + self._turns[turn.task_id] = turn + self._bump() + try: + reply = await turn.reply + finally: + del self._turns[turn.task_id] + self._bump() + request.trajectory.set_format(TRAJECTORY_FORMAT) + for span in reply.trajectory: + request.trajectory.append(span) + if reply.error is not None: + return TaskResult.failure("stand_in.driver_error", reply.error) + return TaskResult.text(reply.text, usage=Usage(tool_call_count=reply.tool_call_count)) + + @extension(DRIVER_V1.claim) + async def claim(self, request: DriverClaimRequest) -> None: + if self._token is None: + self._token = request.token + elif not hmac.compare_digest(self._token, request.token): + raise HTTPException(status_code=409, detail="this stand-in is already claimed") + + @extension(DRIVER_V1.next) + async def next_turns(self, request: DriverNextRequest) -> DriverNextResponse: + self._check(request.token) + if self._version <= request.after: + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(self._next_change()), self._long_poll_seconds) + return DriverNextResponse(version=self._version, turns=[ + DriverTurn(task_id=turn.task_id, context_id=turn.context_id, text=turn.text, + system_prompt=turn.system_prompt, skills=turn.skills) + for turn in self._turns.values() if not turn.reply.done() + ]) + + @extension(DRIVER_V1.reply) + async def reply(self, request: DriverReplyRequest) -> None: + self._check(request.token) + turn = self._turns.get(request.task_id) + if turn is None or turn.reply.done(): + raise HTTPException(status_code=404, detail=f"no open turn {request.task_id!r}") + turn.reply.set_result(request) + + @extension(DRIVER_V1.exec) + async def exec_command(self, request: DriverExecRequest) -> DriverExecResponse: + self._check(request.token) + if not self._shell: + raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it runs no commands") + command_id = uuid.uuid4().hex + self._commands[command_id] = asyncio.create_task(_run(request.command)) + return DriverExecResponse(exec_id=command_id) + + @extension(DRIVER_V1.output) + async def output(self, request: DriverOutputRequest) -> DriverOutputResponse: + self._check(request.token) + if (command := self._commands.get(request.exec_id)) is None: + raise HTTPException(status_code=404, detail=f"no command {request.exec_id!r}") + await asyncio.wait((command,), timeout=self._long_poll_seconds) + if not command.done(): + return DriverOutputResponse(done=False) + self._commands.pop(request.exec_id, None) + return DriverOutputResponse(done=True, output=command.result()) + + @extension(SKILL_CONFIG_V1.add.inline) + async def add_inline_skill(self, request: InlineSkillRequest) -> dict[str, str]: + folder = self._skill_folder(request.name) + folder.mkdir(parents=True) + (folder / "SKILL.md").write_text(request.skill_md) + return {"name": request.name} + + @extension(SKILL_CONFIG_V1.add.bundle) + async def add_bundle_skill(self, request: BundleSkillRequest) -> dict[str, str]: + folder = self._skill_folder(request.name) + for item in request.skill_bundle.files: + await download(item.object, folder / item.path) + return {"name": request.name} + + def _check(self, token: str) -> None: + if self._token is None or not hmac.compare_digest(self._token, token): + raise HTTPException(status_code=401, detail="the driver's token is missing or wrong") + + def _skill_folder(self, name: str) -> Path: + if not self._shell: + raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it keeps no skills") + return SKILLS_DIR / name + + def _next_change(self) -> asyncio.Future[None]: + if self._changed is None: + self._changed = asyncio.get_running_loop().create_future() + return self._changed + + def _bump(self) -> None: + self._version += 1 + if self._changed is not None: + self._changed.set_result(None) + self._changed = None + + +async def _run(command: str) -> str: + """``command``'s stdout then its stderr, as verifiers' bash tool returns them. It runs in a process group of its + own, stopped with whatever it started if it outlives ``COMMAND_TIMEOUT_SECONDS``.""" + process = await asyncio.create_subprocess_exec( + "bash", "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True, + ) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), COMMAND_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" + finally: + if process.returncode is None: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + return (stdout + stderr).decode(errors="replace") + + +if __name__ == "__main__": + StandInAgent().serve() diff --git a/packages/agentenv-protocol/tests/test_a2a_stand_in.py b/packages/agentenv-protocol/tests/test_a2a_stand_in.py new file mode 100644 index 000000000..1908c6859 --- /dev/null +++ b/packages/agentenv-protocol/tests/test_a2a_stand_in.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import os +import time +from typing import Any + +import pytest +from starlette.testclient import TestClient + +from agentenv_protocol.a2a_agent import stand_in +from agentenv_protocol.a2a_agent.stand_in import StandInAgent + +TOKEN = "t" * 32 + + +def _send(client: TestClient, text: str, *, context_id: str = "context-1") -> str: + response = client.post("/a2a", json={ + "jsonrpc": "2.0", "id": "1", "method": "message/send", + "params": { + "message": {"kind": "message", "messageId": f"message-{text}", "role": "user", "contextId": context_id, + "parts": [{"kind": "text", "text": text}]}, + "configuration": {"blocking": False}, + }, + }) + return response.json()["result"]["id"] + + +def _task(client: TestClient, task_id: str) -> dict[str, Any]: + return client.post("/a2a", json={"jsonrpc": "2.0", "id": "1", "method": "tasks/get", + "params": {"id": task_id}}).json()["result"] + + +def _settled(client: TestClient, task_id: str) -> dict[str, Any]: + deadline = time.monotonic() + 10 + while (task := _task(client, task_id))["status"]["state"] not in ("completed", "failed", "canceled"): + assert time.monotonic() < deadline, task + time.sleep(0.02) + return task + + +def _driver(client: TestClient, operation: str, **body: Any): + return client.post(f"/ext/driver/{operation}", json={"token": TOKEN, **body}) + + +def _turns(client: TestClient, after: int = 0) -> dict[str, Any]: + response = _driver(client, "next", after=after) + assert response.status_code == 200, response.text + return response.json() + + +def _ran(client: TestClient, command: str) -> str: + exec_id = _driver(client, "exec", command=command).json()["exec_id"] + while not (result := _driver(client, "output", exec_id=exec_id).json())["done"]: + pass + return result["output"] + + +@pytest.fixture +def client(monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "SKILLS_DIR", tmp_path / "skills") + monkeypatch.chdir(tmp_path) + with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + yield client + + +def test_the_card_advertises_the_driver_and_takes_a_system_prompt(): + card = StandInAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json") + extensions = {extension["uri"]: extension for extension in card["capabilities"]["extensions"]} + + assert sorted(extensions["urn:agentenv:driver/v1"]["params"]["methods"]) == [ + "claim", "exec", "next", "output", "reply"] + assert "system_prompt" in extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"][ + "supported"] + + +def test_the_first_claim_sets_the_token_that_every_driver_call_needs(): + with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: + assert _driver(client, "next").status_code == 401 + assert client.post("/ext/driver/claim", json={"token": "short"}).status_code == 400 + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/driver/claim", json={"token": "u" * 32}).status_code == 409 + assert client.post("/ext/driver/next", json={"token": "u" * 32}).status_code == 401 + assert _driver(client, "next").status_code == 200 + + +def test_a_message_is_an_open_turn_until_the_driver_replies(client): + client.post("/ext/agent-config", json={"system_prompt": "be brief"}) + task_id = _send(client, "fix the build") + + seen = _turns(client) + assert seen["turns"] == [{"task_id": task_id, "context_id": "context-1", "text": "fix the build", + "system_prompt": "be brief", "skills": []}] + span = {"name": "bash", "attributes": {"gen_ai.operation.name": "execute_tool"}} + assert _driver(client, "reply", task_id=task_id, text="fixed", trajectory=[span], tool_call_count=1).status_code \ + == 200 + + task = _settled(client, task_id) + assert task["status"]["state"] == "completed" + parts = task["status"]["message"]["parts"] + assert parts[0] == {"kind": "text", "text": "fixed"} + assert parts[1]["data"]["usage"]["tool_call_count"] == 1 + assert client.post("/ext/trajectory", json={"task_id": task_id}).json()["trajectory"] == [span] + assert _turns(client, after=seen["version"])["turns"] == [] + + +def test_an_error_reply_fails_the_turn(client): + task_id = _send(client, "do it") + + _driver(client, "reply", task_id=task_id, error="the driver gave up") + + task = _settled(client, task_id) + assert task["status"]["state"] == "failed" + assert "the driver gave up" in task["status"]["message"]["parts"][0]["text"] + + +def test_a_reply_names_exactly_one_open_turn(client): + task_id = _send(client, "do it") + + assert _driver(client, "reply", task_id=task_id).status_code == 400 + assert _driver(client, "reply", task_id="nope", text="done").status_code == 404 + _driver(client, "reply", task_id=task_id, text="done") + _settled(client, task_id) + assert _driver(client, "reply", task_id=task_id, text="again").status_code == 404 + + +def test_a_turn_its_sender_cancels_is_withdrawn(client): + task_id = _send(client, "do it") + assert [turn["task_id"] for turn in _turns(client)["turns"]] == [task_id] + + client.post("/a2a", json={"jsonrpc": "2.0", "id": "1", "method": "tasks/cancel", "params": {"id": task_id}}) + + assert _settled(client, task_id)["status"]["state"] == "canceled" + assert _turns(client)["turns"] == [] + + +def test_next_waits_for_a_change_after_the_version_it_was_given(client): + version = _turns(client)["version"] + + started = time.monotonic() + unchanged = _turns(client, after=version) + + assert unchanged["version"] == version + assert time.monotonic() - started >= 0.2 + + +def test_exec_returns_stdout_then_stderr_from_the_folder_it_works_in(client, tmp_path): + (tmp_path / "inputs").mkdir() + + assert _ran(client, "ls; echo out; echo err >&2; exit 3") == "inputs\nout\nerr\n" + + +def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(client, monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + + output = _ran(client, "sleep 60 & echo $! > child; wait") + + assert output == "error: the command timed out after 0.5s" + with pytest.raises(ProcessLookupError): + os.kill(int((tmp_path / "child").read_text()), 0) + + +def test_output_names_a_command_it_started(client): + assert _driver(client, "output", exec_id="nope").status_code == 404 + + +def test_a_skill_is_written_where_its_turns_say_it_is(client, tmp_path): + response = client.post("/ext/skill-config", json={"name": "deploy", "description": "How to deploy.", + "skill_md": "# Deploy\n"}) + assert response.status_code == 200, response.text + _send(client, "go") + + skill = _turns(client)["turns"][0]["skills"] + assert skill == [{"name": "deploy", "description": "How to deploy.", "path": str(tmp_path / "skills" / "deploy")}] + assert (tmp_path / "skills" / "deploy" / "SKILL.md").read_text() == "# Deploy\n" + + +def test_a_stand_in_without_a_shell_keeps_no_skills_and_runs_no_commands(): + with TestClient(StandInAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: + client.post("/ext/driver/claim", json={"token": TOKEN}) + + assert _driver(client, "exec", command="true").status_code == 409 + assert client.post("/ext/skill-config", json={"name": "x", "description": "d", "skill_md": "#"}).status_code \ + == 409 diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 4642fd3ac..d54a57946 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -63,6 +63,7 @@ _DIND_CONTAINER = "agent-dind" _DIND_NETWORK = "agent-docker-net" _DIND_PORT = 2375 +SHARE_STORE_CREDENTIALS = "share_store_credentials" def unwritten_local_agent(agent_id: str, field: str) -> str: @@ -549,7 +550,7 @@ def _build_merged_env(self, resolved_env: dict[str, str], a2a_port: int) -> dict merged_env.update(resolved_env) merged_env["A2A_PORT"] = str(a2a_port) shared = get_config().get_object_store().shared_credentials_env() - if shared.keys().isdisjoint(merged_env): + if self.metadata.get(SHARE_STORE_CREDENTIALS, True) and shared.keys().isdisjoint(merged_env): merged_env.update(shared) merged_env.update(all_sandbox_container_env()) diff --git a/src/agent_env/a2a_agent/stand_in/Dockerfile b/src/agent_env/a2a_agent/stand_in/Dockerfile new file mode 100644 index 000000000..67edcb2ac --- /dev/null +++ b/src/agent_env/a2a_agent/stand_in/Dockerfile @@ -0,0 +1,13 @@ +FROM mirror.gcr.io/library/python:3.12-slim@sha256:423ed6ab25b1921a477529254bfeeabf5855151dc2c3141699a1bfc852199fbf +RUN apt-get update \ + && apt-get install -y --no-install-recommends bash ca-certificates coreutils curl findutils git procps \ + && rm -rf /var/lib/apt/lists/* +COPY requirements.txt /opt/agentenv-stand-in/ +RUN pip install --no-cache-dir -r /opt/agentenv-stand-in/requirements.txt && useradd --create-home agent +COPY agentenv_protocol/ /opt/agentenv-stand-in/agentenv_protocol/ +ENV PYTHONPATH=/opt/agentenv-stand-in +WORKDIR /app +COPY a2a_server.py . +RUN chown agent:agent /app +USER agent +CMD ["python", "/app/a2a_server.py"] diff --git a/src/agent_env/a2a_agent/stand_in/__init__.py b/src/agent_env/a2a_agent/stand_in/__init__.py new file mode 100644 index 000000000..cbf38df10 --- /dev/null +++ b/src/agent_env/a2a_agent/stand_in/__init__.py @@ -0,0 +1,67 @@ +"""The stand-in agent (``agentenv_protocol.a2a_agent.stand_in``) as the Episode API deploys it, in the place of an agent +a loop plays: ``@local/agentenv-stand-in``, an image registered by its build context alone and built where each deploy +runs. Its context is this folder, the installed protocol package, and that package's requirements for serving an +agent, pinned to the versions installed here.""" + +from __future__ import annotations + +import shutil +import tempfile +import threading +from importlib.metadata import requires, version +from pathlib import Path + +import agentenv_protocol +from packaging.requirements import Requirement + +from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS, A2AAgent +from agent_env.artifact.artifacts.docker_image import DockerImageArtifact +from agent_env.store.base import NotFoundError + +STAND_IN_AGENT_ID = "@local/agentenv-stand-in" +_IMAGE_ID = f"{STAND_IN_AGENT_ID}/image" +_PROTOCOL_DISTRIBUTION = "agentenv-framework-protocol" +_HERE = Path(__file__).parent +_writing = threading.Lock() + + +def stand_in_agent() -> A2AAgent: + """The stand-in, written to the store first when it isn't there or its build context has changed.""" + with _writing, tempfile.TemporaryDirectory() as staged: + context = _stage(Path(staged)) + dockerfile = str(context / "Dockerfile") + digest, _ = DockerImageArtifact.context_identity(_IMAGE_ID, str(context), dockerfile, platform=None) + try: + agent = A2AAgent.get(STAND_IN_AGENT_ID) + if agent.docker_image_artifact.source_digest == digest: + return agent + except NotFoundError: + pass + image = DockerImageArtifact.put_context( + _IMAGE_ID, description="The stand-in agent: a driver outside its sandbox answers its turns.", + context_path=str(context), dockerfile_path=dockerfile, platform=None, + ) + return A2AAgent.put( + id=STAND_IN_AGENT_ID, docker_image_artifact=image, + default_env_vars={"LITELLM_API_KEY": "unused", "LITELLM_BASE_URL": "http://unused.invalid"}, + metadata={SHARE_STORE_CREDENTIALS: False}, + ) + + +def _stage(context: Path) -> Path: + for name in ("Dockerfile", "a2a_server.py"): + shutil.copy2(_HERE / name, context / name) + shutil.copytree(Path(agentenv_protocol.__file__).parent, context / "agentenv_protocol", + ignore=shutil.ignore_patterns("__pycache__", "*.pyc")) + (context / "requirements.txt").write_text("".join(f"{line}\n" for line in _requirements())) + return context + + +def _requirements() -> list[str]: + pinned = [] + for spec in requires(_PROTOCOL_DISTRIBUTION) or []: + requirement = Requirement(spec) + if requirement.marker is None or requirement.marker.evaluate({"extra": "agent"}): + extras = f"[{','.join(sorted(requirement.extras))}]" if requirement.extras else "" + pinned.append(f"{requirement.name}{extras}=={version(requirement.name)}") + return sorted(pinned) diff --git a/src/agent_env/a2a_agent/stand_in/a2a_server.py b/src/agent_env/a2a_agent/stand_in/a2a_server.py new file mode 100644 index 000000000..dadb0b388 --- /dev/null +++ b/src/agent_env/a2a_agent/stand_in/a2a_server.py @@ -0,0 +1,37 @@ +"""The stand-in agent as AgentEnv deploys it: served from ``/app`` in its image, or installed on a VM host as +``install_agent`` installs it, in the place of the agent it replaces (its ``{workdir}``, as its ``{user}``). +``AGENT_CARD`` is what ``install_agent`` reads; importing this module serves nothing. ``install_agent`` formats each +command with ``str.format``, so a literal brace is doubled.""" + +from agentenv_protocol.a2a_agent import INSTALL_V1, a2a_agent, enable +from agentenv_protocol.a2a_agent.stand_in import EXTENSIONS, IDENTITY, StandInAgent, StandInConfig + +HOME = "/opt/agentenv-stand-in-{agent_name}" +INSTALL_COMMANDS_HOST = [ + "bash -c 'set -e; python3 -c \"import ensurepip, venv\" 2>/dev/null || {{ export DEBIAN_FRONTEND=noninteractive; " + "apt-get update -qq; apt-get install -y -qq --no-install-recommends python3 python3-venv; }}'", + f"rm -rf {HOME} && mkdir -p {HOME} && tar -xzf {{agent_ctx_tar}} -C {HOME}", + f"python3 -m venv {HOME}/venv && {HOME}/venv/bin/pip install -q -r {HOME}/requirements.txt", + "id -u {user} >/dev/null 2>&1 || useradd -m {user}; mkdir -p {workdir} && chown -R {user}: {workdir}", + f"nohup runuser -u {{user}} -- env -C {{workdir}} A2A_PORT={{a2a_port}} PYTHONPATH={HOME} {HOME}/venv/bin/python " + f"{HOME}/a2a_server.py > /tmp/agentenv-stand-in-{{agent_name}}.log 2>&1 < /dev/null &", +] + + +@a2a_agent( + identity=IDENTITY, + config=StandInConfig, + extensions=( + *EXTENSIONS, + enable(INSTALL_V1, install_commands_host=INSTALL_COMMANDS_HOST, + required_params_host=["agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], a2a_port=8000), + ), +) +class InstalledStandIn(StandInAgent): + pass + + +AGENT_CARD = InstalledStandIn().create_app().state.agentenv_a2a.card + +if __name__ == "__main__": + InstalledStandIn().serve() diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 8d766b428..176ce491e 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -1,8 +1,9 @@ """Which steps a task may have when some of its agents are played from outside AgentEnv, and why the others are refused. -Every registered step type is either checked here (``CHECKS``) or listed as untouched by an outside agent -(``UNAFFECTED``). A test holds the registry to the two, so a new step type is classified before an episode runs it. -Plugin step types aren't inspected. +An outside agent is played through a stand-in agent deployed in its place (``agentenv_protocol.a2a_agent.stand_in``), +so steps that act on the agent, or in its sandbox, act on the stand-in. Every registered step type is either checked +here (``CHECKS``) or listed as untouched by an outside agent (``UNAFFECTED``). A test holds the registry to the two, so +a new step type is classified before an episode runs it. Plugin step types aren't inspected. """ from __future__ import annotations @@ -15,6 +16,7 @@ from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep + from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep from agent_env.task_step.task_steps.peer_agents import PeerAgentsTaskStep from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep @@ -33,10 +35,10 @@ def problems(steps: list[TaskStep], outside: list[str]) -> list[str]: """Why each outside agent, and each step, can't run with those agents played from outside; empty when it can.""" found = [] for name in outside: - deploys = [step for step in steps if step.type == "deploy_agent" and step.agent_name == name] + deploys = [step for step in steps if step.type in DEPLOYS and step.agent_name == name] if len(deploys) != 1: - found.append(f"agent {name!r} is played from outside, so it needs exactly one deploy_agent step, " - f"not {len(deploys)}") + found.append(f"agent {name!r} is played from outside, so it needs exactly one deploy_agent or " + f"install_agent step, not {len(deploys)}") scope = Outside(frozenset(outside), { step.id: step.prompt_id for step in steps if step.type == "prompt_agent" and step.agent_name in outside }) @@ -48,6 +50,17 @@ def problems(steps: list[TaskStep], outside: list[str]) -> list[str]: return found +def works_in_sandbox(steps: list[TaskStep], agent: str) -> bool: + """Whether ``steps`` use ``agent``'s sandbox, so its stand-in is deployed there rather than served in this + process.""" + return any(step.agent_name == agent and ( + step.type in ("install_agent", "load_artifact", "run_code", "add_skills") + or step.type == "deploy_agent" and step.skills + or step.type == "verify_sandbox" and not step.sandbox_name + or step.type == "collect_artifacts" and not (step.env_id or step.container_name or step.sandbox_name) + ) for step in steps if getattr(step, "agent_name", None)) + + def _prompt_agent(step: PromptAgentTaskStep, outside: Outside) -> list[str]: if step.agent_name in outside.agents: return [ @@ -69,7 +82,6 @@ def _deploy_agent(step: DeployAgentTaskStep, outside: Outside) -> list[str]: return [] return [ f"an outside agent can't have {what}" for what, present in ( - ("skills, instructions its player wouldn't get", step.skills), ("agent_snapshot_files_artifact_id", step.agent_snapshot_files_artifact_id), ("agent_changelog_object_url", step.agent_changelog_object_url), ("enable_agent_changelog", step.enable_agent_changelog), @@ -77,6 +89,13 @@ def _deploy_agent(step: DeployAgentTaskStep, outside: Outside) -> list[str]: ] +def _install_agent(step: InstallAgentTaskStep, outside: Outside) -> list[str]: + if step.agent_name not in outside.agents or step.container_name is None: + return [] + return [f"an outside agent's stand-in replaces an agent installed on the VM host, not in container " + f"{step.container_name!r}"] + + def _deploys_an_agent(step: TaskStep, outside: Outside) -> list[str]: if step.agent_name not in outside.agents: return [] @@ -94,12 +113,9 @@ def _validates_an_agent(step: TaskStep, outside: Outside) -> list[str]: def _collect_artifacts(step: CollectArtifactsTaskStep, outside: Outside) -> list[str]: - why = [] - if step.agent_name in outside.agents and not (step.env_id or step.container_name or step.sandbox_name): - why.append(f"collects from {step.agent_name!r}'s sandbox, and it is played from outside") if step.manifest_step_id in outside.prompts: - why.append(f"its manifest is outside prompt {step.manifest_step_id!r}") - return why + return [f"its manifest is outside prompt {step.manifest_step_id!r}"] + return [] def _register_env_triggers(step: RegisterEnvTriggersStep, outside: Outside) -> list[str]: @@ -116,25 +132,18 @@ def _peer_agents(step: PeerAgentsTaskStep, outside: Outside) -> list[str]: def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[str]: - why = [] - if step.use_trajectory and step.prompt_id in outside.prompts.values(): - why.append("grades an outside prompt's trajectory, which an outside player doesn't hand back") if step.use_agent_judge and step.agent_name in outside.agents: - why.append(f"judges in {step.agent_name!r}'s sandbox, and it is played from outside") - return why + return [f"judges in {step.agent_name!r}'s sandbox, and it is played from outside"] + return [] CHECKS: dict[str, Callable[..., list[str]]] = { "prompt_agent": _prompt_agent, "deploy_agent": _deploy_agent, - "install_agent": _deploys_an_agent, + "install_agent": _install_agent, "deploy_human_agent": _deploys_an_agent, - "add_skills": _needs_the_agent, "register_agent_triggers": _needs_the_agent, "snapshot_agent_state": _needs_the_agent, - "load_artifact": _needs_the_agent, - "verify_sandbox": _needs_the_agent, - "run_code": _needs_the_agent, "collect_artifacts": _collect_artifacts, "register_env_triggers": _register_env_triggers, "peer_agents": _peer_agents, @@ -157,4 +166,8 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s "agent_prompt_response_verifier", "verify_mcp_env_assessment", # It skips an agent that isn't deployed. "teardown_sandboxes", + # They act on an outside agent's stand-in, in its sandbox, as on the agent. + "add_skills", "load_artifact", "run_code", "verify_sandbox", }) + +DEPLOYS = ("deploy_agent", "install_agent") diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index 8ee4ddfd4..86abb2810 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -1,37 +1,59 @@ from __future__ import annotations import asyncio +import json import logging from collections import deque -from typing import TYPE_CHECKING, Any +from dataclasses import dataclass +from datetime import datetime, timezone +from typing import Any +import httpx from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url +from agent_env.store.routing import namespace_routing +from agent_env.task.episode.stand_in import InProcessStandIn, StandInClient from agent_env.task.store import record_task_cancelled from agent_env.task.task import Task, recorded_scores, substitute_seed_in_steps from agent_env.task.teardown import teardown_run -from agent_env.task_step.context import TaskStepContext - -if TYPE_CHECKING: - from agent_env.task.episode.steps import _EpisodeTurnStep +from agent_env.task_step.context import DeployedAgent, TaskStepContext logger = logging.getLogger(__name__) _live: dict[str, Episode] = {} +_DRAIN_RETRY_SECONDS = 1.0 +_CARD_TIMEOUT_SECONDS = 10.0 + +BASH = EnvironmentTool( + name="bash", + description="Run a bash command where the agent works, as its user. Returns its stdout, then its stderr.", + inputSchema={"type": "object", "properties": {"command": {"type": "string", "description": "The command to run."}}, + "required": ["command"]}, +) class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" +@dataclass(frozen=True) +class Skill: + """A skill the task gave the agent: its instructions are ``path``'s ``SKILL.md``, in the agent's sandbox.""" + + name: str + description: str + path: str | None + + class Turn: - """An outside agent's turn. Its model reads ``system_prompt`` and ``prompt``, acts through ``tools()`` (or an MCP - client of its own on ``mcp_urls``, sending ``headers`` and ``role``), and hands its reply back with ``end()``.""" + """An outside agent's turn. Its model reads ``system_prompt``, ``prompt`` and ``skills``, acts through ``tools()`` + (or an MCP client of its own on ``mcp_urls``, sending ``headers`` and ``role``), and hands its reply back with + ``end()``.""" def __init__( self, *, agent: str, step_id: str, prompt: str, system_prompt: str | None, role: str | None, - mcp_urls: dict[str, str], + mcp_urls: dict[str, str], skills: list[Skill], shell: bool, stand_in: StandInClient, task_id: str, ): self.agent = agent self.step_id = step_id @@ -40,43 +62,50 @@ def __init__( self.role = role self.mcp_urls = mcp_urls self.headers = {env_id: sandbox_request_headers_for_url(url) for env_id, url in mcp_urls.items()} - self._reply: asyncio.Future[str] = asyncio.get_running_loop().create_future() + self.skills = skills + self._shell = shell + self._stand_in = stand_in + self._task_id = task_id + self._spans: list[dict[str, Any]] = [] + self._started = _now() + self._open = True def tools(self, *, timeout: float = 600.0) -> TurnTools: """A client for the tools of the envs this agent acts in, calling them as its role with the headers each URL - needs; use it with ``async with``. ``timeout`` bounds each call, in seconds.""" + needs, and ``bash`` when it works in a sandbox; use it with ``async with``. ``timeout`` bounds each env tool's + call, in seconds.""" return TurnTools(self, timeout) async def end(self, response: str | None = None) -> None: """Hand the turn back, and the run carries on. No reply records an empty one, as an agent that answers with nothing leaves. Once the turn is over, ended or not, this does nothing.""" + if not self._open: + return + self._open = False + reply = response or "" + spans = [*self._spans, _span("turn", "chain", self._started, completion={"content": [ + {"type": "text", "text": reply}]})] + await self._stand_in.reply(self._task_id, text=reply, trajectory=spans, tool_call_count=len(self._spans)) + + async def _fail(self, error: Exception) -> None: if self._open: - self._reply.set_result(response or "") - - @property - def _open(self) -> bool: - return not self._reply.done() - - async def _replied_within(self, seconds: float) -> str: - return await asyncio.wait_for(self._reply, seconds) - - def _fail(self, error: Exception) -> None: - if self._open: - self._reply.set_exception(error) + self._open = False + await self._stand_in.reply(self._task_id, error=str(error)) class TurnTools: - """The tools of a turn's envs, listed once when it opens: ``list()`` names every one, and ``call()`` routes a - call to the env that offers it.""" + """The tools of a turn, listed once when it opens: its envs', and ``bash`` when its agent works in a sandbox. + ``list()`` names every one, ``call()`` routes a call to what offers it, and each call goes in the turn's + trajectory.""" def __init__(self, turn: Turn, timeout: float): - self._step_id = turn.step_id + self._turn = turn self._sessions = { env_id: ToolSession(mcp_url=url, role=turn.role, headers=turn.headers[env_id], timeout=timeout) for env_id, url in turn.mcp_urls.items() } self._tools: list[EnvironmentTool] = [] - self._offered_by: dict[str, str] = {} + self._offered_by: dict[str, str | None] = {} async def __aenter__(self) -> TurnTools: try: @@ -87,6 +116,12 @@ async def __aenter__(self) -> TurnTools: raise EpisodeError(f"tool {tool.name!r} is offered by both env {other!r} and env {env_id!r}") self._offered_by[tool.name] = env_id self._tools.append(tool) + if self._turn._shell: + if (env_id := self._offered_by.get(BASH.name)) is not None: + raise EpisodeError(f"tool {BASH.name!r} is offered by env {env_id!r}, and the agent's sandbox " + "offers its own") + self._offered_by[BASH.name] = None + self._tools.append(BASH) except BaseException: await self._close() raise @@ -101,8 +136,21 @@ def list(self) -> list[EnvironmentTool]: async def call(self, name: str, arguments: dict[str, Any] | None = None) -> ToolResult: if name not in self._offered_by: raise EpisodeError( - f"no env of turn {self._step_id!r} offers tool {name!r}; they offer {sorted(self._offered_by)}") - return await self._sessions[self._offered_by[name]].call_tool(name, arguments) + f"no env of turn {self._turn.step_id!r} offers tool {name!r}; they offer {sorted(self._offered_by)}") + started = _now() + env_id = self._offered_by[name] + if env_id is None: + result = await self._bash(arguments or {}) + else: + result = await self._sessions[env_id].call_tool(name, arguments) + self._turn._spans.append(_span(name, "execute_tool", started, prompt={"input": arguments or {}}, + completion={"output": result.content})) + return result + + async def _bash(self, arguments: dict[str, Any]) -> ToolResult: + if not isinstance(command := arguments.get("command"), str): + raise EpisodeError(f"tool {BASH.name!r} takes a command string, not {command!r}") + return ToolResult(content=[{"type": "text", "text": await self._turn._stand_in.run(command)}]) async def _close(self) -> None: await asyncio.gather(*(session.close() for session in self._sessions.values())) @@ -119,12 +167,14 @@ def __init__(self, id: str, task: Task, context: TaskStepContext): self._task = task self._context = context self._run: asyncio.Task[None] - self._bindings: dict[str, tuple[str | None, dict[str, str]]] = {} self._waiting: deque[Turn] = deque() self._turns: set[Turn] = set() self._changed: asyncio.Future[None] = asyncio.get_running_loop().create_future() self._tearing_down = False self._abandoned: str | None = None + self._stand_ins: list[StandInClient] = [] + self._drains: list[asyncio.Task[None]] = [] + self._in_process: list[InProcessStandIn] = [] async def next_turn(self) -> Turn | None: """The next outside turn the run reaches, or None once the run has finished. Raises the run's error if it @@ -151,8 +201,7 @@ async def result(self) -> dict[str, float]: ) await self._wait(changed) self._raise_unless_completed() - seeded = substitute_seed_in_steps(self._task.steps, self._context.metadata.get("seed") or {}) - return recorded_scores(seeded, self._context) + return recorded_scores(self._seeded_steps(), self._context) async def abandon(self, reason: str) -> None: """Stop the run and wait for its teardown; it is recorded cancelled with ``reason``. A run a step has @@ -160,8 +209,8 @@ async def abandon(self, reason: str) -> None: Once the run has ended, or while it tears down, this only waits.""" if not (self._run.done() or self._tearing_down or self._abandoned is not None): if any(failure.get("is_fatal") for failure in self._context.metadata.get("failed_steps") or []): - for turn in self._turns: - turn._fail(EpisodeError(f"episode abandoned: {reason}")) + for turn in list(self._turns): + await turn._fail(EpisodeError(f"episode abandoned: {reason}")) else: self._abandoned = reason self._run.cancel(msg=reason) @@ -181,19 +230,24 @@ async def _launch(self) -> None: raise async def _drive(self, running: asyncio.Event) -> None: + """Run the task with namespace routing on, as ``agent-env run`` does, so the stand-in, an ``@local`` agent, + resolves.""" running.set() - try: - await self._task.run(context=self._context, instance_id=self.id) - finally: - self._tearing_down = True + with namespace_routing(): try: - report = await teardown_run(self._context) - if report.still_up: - logger.warning("episode %s: %s may still be up", self.id, [s.sandbox_id for s in report.still_up]) - if self._abandoned is not None: - await asyncio.to_thread(record_task_cancelled, self.id, self._abandoned, Task._utc_now_str()) + await self._task.run(context=self._context, instance_id=self.id) finally: - del _live[self.id] + self._tearing_down = True + try: + await self._release_stand_ins() + report = await teardown_run(self._context) + if report.still_up: + logger.warning("episode %s: %s may still be up", self.id, + [s.sandbox_id for s in report.still_up]) + if self._abandoned is not None: + await asyncio.to_thread(record_task_cancelled, self.id, self._abandoned, Task._utc_now_str()) + finally: + del _live[self.id] async def _wait(self, changed: asyncio.Future[None]) -> None: """Until the run ends or what was ``changed`` when the caller last looked changes again.""" @@ -205,29 +259,93 @@ def _raise_unless_completed(self) -> None: if (error := self._run.exception()) is not None: raise error - def _bind(self, agent: str, role: str | None, mcp_urls: dict[str, str]) -> None: - self._bindings[agent] = (role, mcp_urls) - - async def _play(self, step: _EpisodeTurnStep) -> str: - if step.agent_name not in self._bindings: - raise RuntimeError( - f"agent {step.agent_name!r} isn't deployed yet; turn {step.id!r} must depend on its deploy_agent step" - ) - role, mcp_urls = self._bindings[step.agent_name] - turn = Turn(agent=step.agent_name, step_id=step.id, prompt=step.prompt, system_prompt=step.system_prompt, - role=role, mcp_urls=dict(mcp_urls)) - self._waiting.append(turn) - self._turns.add(turn) - self._changed.set_result(None) - self._changed = asyncio.get_running_loop().create_future() - try: - return await turn._replied_within(step.timeout_seconds) - except TimeoutError: - raise TimeoutError( - f"turn {step.id!r} of agent {step.agent_name!r} wasn't ended within {step.timeout_seconds}s" - ) from None - finally: - self._turns.discard(turn) + async def _serve_in_process(self, agent: str, role: str | None, mcp_urls: dict[str, str]) -> DeployedAgent: + """A stand-in served in this process for ``agent``, which works in no sandbox, played as the others are.""" + server = InProcessStandIn() + self._in_process.append(server) + url = await server.start() + async with httpx.AsyncClient(timeout=_CARD_TIMEOUT_SECONDS) as client: + response = await client.get(f"{url}/.well-known/agent.json") + response.raise_for_status() + deployed = DeployedAgent(agent_name=agent, api_url=url, a2a_url=url, a2a_card=response.json(), role=role) + await self._attach(agent, deployed, role, mcp_urls, shell=False) + return deployed + + async def _attach( + self, agent: str, deployed: DeployedAgent, role: str | None, mcp_urls: dict[str, str], *, shell: bool, + ) -> None: + """Claim ``agent``'s stand-in and hand out the turns it opens.""" + stand_in = StandInClient(deployed.a2a_url, deployed.a2a_card or {}) + self._stand_ins.append(stand_in) + await stand_in.claim() + self._drains.append(asyncio.create_task(self._drain(agent, stand_in, role, mcp_urls, shell))) + + async def _drain( + self, agent: str, stand_in: StandInClient, role: str | None, mcp_urls: dict[str, str], shell: bool, + ) -> None: + """Hand out each turn the stand-in opens, and close each one it drops, until the run ends.""" + after, open_turns = 0, {} + while True: + try: + answer = await stand_in.next(after) + except httpx.HTTPError as e: + logger.warning("episode %s: reading agent %r's turns failed, retrying: %s", self.id, agent, e) + await asyncio.sleep(_DRAIN_RETRY_SECONDS) + continue + after = answer["version"] + listed = {item["task_id"]: item for item in answer["turns"]} + for task_id in [task_id for task_id in open_turns if task_id not in listed]: + turn = open_turns.pop(task_id) + turn._open = False + self._turns.discard(turn) + for task_id, item in listed.items(): + if task_id not in open_turns: + open_turns[task_id] = turn = Turn( + agent=agent, step_id=self._step_of(agent, item["context_id"]), prompt=item["text"], + system_prompt=item.get("system_prompt"), role=role, mcp_urls=dict(mcp_urls), + skills=[Skill(**skill) for skill in item.get("skills", [])], shell=shell, stand_in=stand_in, + task_id=task_id, + ) + self._waiting.append(turn) + self._turns.add(turn) + self._changed.set_result(None) + self._changed = asyncio.get_running_loop().create_future() + + def _step_of(self, agent: str, context_id: str) -> str: + """The ``prompt_agent`` step whose conversation with ``agent`` a message on ``context_id`` is in.""" + conversations = self._context.metadata.get("a2a_conversations", {}) + for step in self._seeded_steps(): + if step.type == "prompt_agent" and step.agent_name == agent and ( + step.context_id or conversations.get(step.id)) == context_id: + return step.id + return context_id + + def _seeded_steps(self) -> list: + return substitute_seed_in_steps(self._task.steps, self._context.metadata.get("seed") or {}) + + async def _release_stand_ins(self) -> None: + for drain in self._drains: + drain.cancel() + await asyncio.gather(*self._drains, return_exceptions=True) + for turn in self._turns: + turn._open = False + for server in self._in_process: + server.stop() + await asyncio.gather(*(stand_in.close() for stand_in in self._stand_ins), return_exceptions=True) + + +def _span(name: str, operation: str, started: str, *, prompt: Any = None, completion: Any = None) -> dict[str, Any]: + """An OTel GenAI span, the trajectory format rubrics read.""" + attributes = {"gen_ai.operation.name": operation} + if prompt is not None: + attributes["gen_ai.prompt"] = json.dumps(prompt, default=str) + if completion is not None: + attributes["gen_ai.completion"] = json.dumps(completion, default=str) + return {"name": name, "start_time": started, "end_time": _now(), "attributes": attributes} + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() def _episode_of(context: TaskStepContext) -> Episode: diff --git a/src/agent_env/task/episode/spec.py b/src/agent_env/task/episode/spec.py index 0ddddda03..c8b0c0263 100644 --- a/src/agent_env/task/episode/spec.py +++ b/src/agent_env/task/episode/spec.py @@ -1,12 +1,13 @@ from __future__ import annotations +import math import uuid from collections.abc import Mapping, Sequence from typing import Any -from agent_env.task.episode.rules import problems +from agent_env.task.episode.rules import DEPLOYS, problems, works_in_sandbox from agent_env.task.episode.run import Episode, EpisodeError -from agent_env.task.episode.steps import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.episode.steps import _EpisodeAgentStep from agent_env.task.task import Task, substitute_seed_in_steps from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_step import TaskStep @@ -28,18 +29,15 @@ def from_task( turn_timeout: float | None = None, ) -> EpisodeSpec: """``trainer_agents`` are the agents the loop plays, by default the task's only prompted agent. - ``turn_timeout`` is every turn's deadline in seconds, in place of its prompt's own ``timeout_seconds``, - which budgets an agent's run rather than the loop playing it.""" + ``turn_timeout`` is every turn's deadline in seconds, rounded up to a whole one, in place of its prompt's own + ``timeout_seconds``, which budgets an agent's run rather than the loop playing it.""" seed = dict(seed) if seed else None seeded = substitute_seed_in_steps(task.steps, seed or {}) outside = _trainer_agents(task, seeded, trainer_agents) if found := problems(seeded, outside): raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(found)) - deploys = { - as_run.agent_name: step for step, as_run in zip(task.steps, seeded) - if step.type == "deploy_agent" and as_run.agent_name in outside - } - steps = [_stand_in(step, as_run, deploys, turn_timeout) for step, as_run in zip(task.steps, seeded)] + sandboxed = {name: works_in_sandbox(seeded, name) for name in outside} + steps = [_stand_in(step, as_run, sandboxed, turn_timeout) for step, as_run in zip(task.steps, seeded)] return cls(copy=Task(id=task.id, version=task.version, steps=steps), trainer_agents=tuple(outside), seed=seed) async def start(self) -> Episode: @@ -71,21 +69,15 @@ def _trainer_agents(task: Task, steps: list[TaskStep], named: Sequence[str] | No return outside -def _stand_in(step: TaskStep, as_run: TaskStep, deploys: dict[str, TaskStep], turn_timeout: float | None) -> TaskStep: +def _stand_in(step: TaskStep, as_run: TaskStep, sandboxed: dict[str, bool], turn_timeout: float | None) -> TaskStep: """The step the episode's copy of the task runs in place of ``step``, judged by ``as_run``, the step as the - seed makes it; ``deploys`` are the outside agents', by name. A stand-in keeps the step's own fields, which the - run seeds.""" - if as_run.type == "deploy_agent" and as_run.agent_name in deploys: + seed makes it; ``sandboxed`` says, for each outside agent by name, whether its stand-in is deployed in its + sandbox. The copy keeps every step's own fields, which the run seeds.""" + if as_run.type in DEPLOYS and as_run.agent_name in sandboxed: return _EpisodeAgentStep( - id=step.id, version=step.version, depends_on=step.depends_on, fail_task_on_error=step.fail_task_on_error, - agent_name=step.agent_name, env_ids=step.env_ids, env_step_id=step.env_step_id, role=step.role, - ) - if as_run.type == "prompt_agent" and as_run.agent_name in deploys: - timeout = step.timeout_seconds if turn_timeout is None else turn_timeout - return _EpisodeTurnStep( - id=step.id, version=step.version, depends_on=step.depends_on, fail_task_on_error=step.fail_task_on_error, - agent_name=step.agent_name, prompt_id=step.prompt_id, prompt=step.prompt, - system_prompt=step.system_prompt or deploys[as_run.agent_name].system_prompt, - timeout_seconds=PromptAgentTaskStep.DEFAULT_PROMPT_TIMEOUT_SECONDS if timeout is None else timeout, + id=step.id, version=step.version, step=step, sandboxed=sandboxed[as_run.agent_name], + depends_on=step.depends_on, fail_task_on_error=step.fail_task_on_error, ) + if as_run.type == "prompt_agent" and as_run.agent_name in sandboxed and turn_timeout is not None: + return PromptAgentTaskStep.from_dict({**step.to_dict(), "timeout_seconds": math.ceil(turn_timeout)}) return step diff --git a/src/agent_env/task/episode/stand_in.py b/src/agent_env/task/episode/stand_in.py new file mode 100644 index 000000000..dbff4ecbe --- /dev/null +++ b/src/agent_env/task/episode/stand_in.py @@ -0,0 +1,89 @@ +"""The episode's end of a stand-in agent: the driver calls of ``urn:agentenv:driver/v1``, and a stand-in served in this +process for an outside agent that works in no sandbox.""" + +from __future__ import annotations + +import asyncio +import secrets +import threading +from typing import Any + +import httpx +import uvicorn +from agentenv_protocol.a2a_agent import DRIVER_V1 +from agentenv_protocol.a2a_agent.stand_in import StandInAgent + +from agent_env.a2a_agent.a2a_agent import A2AAgent + +_CALL_TIMEOUT = httpx.Timeout(60.0, connect=10.0) +_SERVE_TIMEOUT_SECONDS = 30.0 +# Short, so a long-poll in flight when the episode ends finishes within the server's shutdown. +_IN_PROCESS_LONG_POLL_SECONDS = 1.0 + + +class StandInClient: + """One stand-in, claimed with a token of this episode's, at the endpoints its card names.""" + + def __init__(self, a2a_url: str, card: dict[str, Any]): + driver = A2AAgent.find_extension(card, DRIVER_V1.uri) + if driver is None: + raise RuntimeError(f"the agent at {a2a_url} isn't a stand-in: its card has no {DRIVER_V1.uri}") + self._endpoints = {operation: a2a_url + A2AAgent.operation(driver, operation)[1] + for operation in DRIVER_V1.operations} + self._token = secrets.token_urlsafe(24) + self._client = httpx.AsyncClient(timeout=_CALL_TIMEOUT) + + async def claim(self) -> None: + await self._call("claim") + + async def next(self, after: int) -> dict[str, Any]: + return await self._call("next", after=after) + + async def reply(self, task_id: str, **reply: Any) -> bool: + """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone.""" + response = await self._client.post(self._endpoints["reply"], json={"token": self._token, "task_id": task_id, + **reply}) + if response.status_code == 404: + return False + response.raise_for_status() + return True + + async def run(self, command: str) -> str: + exec_id = (await self._call("exec", command=command))["exec_id"] + while not (output := await self._call("output", exec_id=exec_id))["done"]: + pass + return output["output"] + + async def close(self) -> None: + await self._client.aclose() + + async def _call(self, operation: str, **body: Any) -> dict[str, Any]: + response = await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) + response.raise_for_status() + return response.json() + + +class InProcessStandIn: + """A stand-in served on loopback in a thread of its own, with no shell: one for an outside agent that works in no + sandbox, which only this process's steps message.""" + + def __init__(self) -> None: + agent = StandInAgent(shell=False, long_poll_seconds=_IN_PROCESS_LONG_POLL_SECONDS) + self._server = uvicorn.Server(uvicorn.Config(agent.create_app(), host="127.0.0.1", port=0, log_level="warning", + timeout_graceful_shutdown=2 * _IN_PROCESS_LONG_POLL_SECONDS)) + self._thread = threading.Thread(target=self._server.run, name="agentenv-stand-in", daemon=True) + + async def start(self) -> str: + """Serve it, and return its URL once it answers.""" + self._thread.start() + async with asyncio.timeout(_SERVE_TIMEOUT_SECONDS): + while not self._server.started: + if not self._thread.is_alive(): + raise RuntimeError("the in-process stand-in agent stopped as it started") + await asyncio.sleep(0.01) + port = self._server.servers[0].sockets[0].getsockname()[1] + return f"http://127.0.0.1:{port}" + + def stop(self) -> None: + """Have it stop; its thread ends once the requests in flight have.""" + self._server.should_exit = True diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py index f6f55d332..f6df0f33c 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -1,80 +1,63 @@ from __future__ import annotations +import asyncio from typing import Any, ClassVar, Optional -from agent_env.entity.refs import EntityRef +from agent_env.a2a_agent import protocol +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent from agent_env.task.episode.run import _episode_of -from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url class _EpisodeAgentStep(TaskStep): - """An outside agent's ``deploy_agent``: deploys nothing, and binds the agent to its envs where it stood.""" + """An outside agent's ``deploy_agent`` or ``install_agent``. In a task that uses the agent's sandbox it runs that + step with the stand-in agent in the agent's place; otherwise it serves a stand-in in this process. Either way the + episode then plays the stand-in, which the agent's ``prompt_agent`` steps message as they would the agent.""" type: ClassVar[str] = "episode_agent" - entity_refs = (EntityRef.env("env_ids[]"),) + entity_refs = () def __init__( - self, id: str, version: Optional[int], *, agent_name: str, env_ids: list[str], env_step_id: str | None, - role: str | None, depends_on: Optional[list] = None, fail_task_on_error: bool = True, + self, id: str, version: Optional[int], *, step: TaskStep, sandboxed: bool, + depends_on: Optional[list] = None, fail_task_on_error: bool = True, ): super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) - self.agent_name = agent_name - self.env_ids = list(env_ids) - self.env_step_id = env_step_id - self.role = role + self.step = step + self.sandboxed = sandboxed def to_dict(self) -> dict[str, Any]: - return {**super().to_dict(), "agent_name": self.agent_name, "env_ids": self.env_ids, - "env_step_id": self.env_step_id, "role": self.role} + return {**super().to_dict(), "step": self.step.to_dict(), "sandboxed": self.sandboxed} @classmethod def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: - return cls(**cls._base_from_dict(data), agent_name=data["agent_name"], env_ids=data["env_ids"], - env_step_id=data["env_step_id"], role=data["role"]) + step = get_task_step_registry()[data["step"]["type"]].from_dict(data["step"]) + return cls(**cls._base_from_dict(data), step=step, sandboxed=data["sandboxed"]) async def execute(self, context: TaskStepContext) -> TaskStepContext: + episode = _episode_of(context) + agent = self.step.agent_name + role = getattr(self.step, "role", None) mcp_urls = {} - for env_id in self.env_ids: - deployed = bound_deployment(context, env_id, env_step_id=self.env_step_id, step_id=self.id) + for env_id in getattr(self.step, "env_ids", None) or []: + deployed = bound_deployment(context, env_id, env_step_id=self.step.env_step_id, step_id=self.id) mcp_urls[env_id] = deployed.mcp_url if deployed is not None else live_mcp_url(env_id) - _episode_of(context)._bind(self.agent_name, self.role, mcp_urls) - return context - - -class _EpisodeTurnStep(TaskStep): - """An outside agent's ``prompt_agent``: hands the turn to the caller and records the reply it gets back.""" - - type: ClassVar[str] = "episode_turn" - entity_refs = () - - def __init__( - self, id: str, version: Optional[int], *, agent_name: str, prompt_id: str, prompt: str, - system_prompt: str | None, timeout_seconds: float, depends_on: Optional[list] = None, - fail_task_on_error: bool = True, - ): - super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) - self.agent_name = agent_name - self.prompt_id = prompt_id - self.prompt = prompt - self.system_prompt = system_prompt - self.timeout_seconds = timeout_seconds - - def to_dict(self) -> dict[str, Any]: - return {**super().to_dict(), "agent_name": self.agent_name, "prompt_id": self.prompt_id, - "prompt": self.prompt, "system_prompt": self.system_prompt, "timeout_seconds": self.timeout_seconds} - - @classmethod - def from_dict(cls, data: dict[str, Any]) -> _EpisodeTurnStep: - return cls(**cls._base_from_dict(data), agent_name=data["agent_name"], prompt_id=data["prompt_id"], - prompt=data["prompt"], system_prompt=data["system_prompt"], - timeout_seconds=data["timeout_seconds"]) - - async def execute(self, context: TaskStepContext) -> TaskStepContext: - reply = await _episode_of(context)._play(self) - context.prompt_responses.append(PromptResponse( - prompt_id=self.prompt_id, response=reply, prompt_text=self.prompt, agent_name=self.agent_name, - step_id=self.id, - )) + if not self.sandboxed: + deployed = await episode._serve_in_process(agent, role, mcp_urls) + context.deployed_agents.append(deployed) + if system_prompt := getattr(self.step, "system_prompt", None): + config = A2AAgent.negotiate_agent_config(deployed.a2a_card, {"system_prompt": system_prompt}) + await protocol.post_agent_config(deployed.a2a_url + config.endpoint, config.fields) + return context + await asyncio.to_thread(stand_in_agent) + override = {"a2a_agent_id": STAND_IN_AGENT_ID} if self.step.type == "deploy_agent" else { + "stand_in": STAND_IN_AGENT_ID} + step_params = context.metadata.setdefault("user_overrides", {}).setdefault("step_params", {}) + step_params[self.step.id] = {**step_params.get(self.step.id, {}), **override} + context = await self.step.execute(context) + deployed = next(deployed for deployed in context.deployed_agents if deployed.agent_name == agent) + await episode._attach(agent, deployed, role, mcp_urls, shell=True) return context diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 40df1eb27..07d2e5540 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -43,6 +43,8 @@ REQUIRED_PARAMS_KEY = "required_params" REQUIRED_PARAMS_HOST_KEY = "required_params_host" A2A_PORT_KEY = "a2a_port" +WORKDIR_HOST_KEY = "workdir_host" +USER_HOST_KEY = "user_host" def _install_problem(agent: A2AAgent) -> str | None: @@ -66,6 +68,7 @@ def _build_param_resolvers( litellm_api_key_override: Optional[str] = None, sandbox_type: Optional[str] = None, trust_lines: str = "", + place: Optional[dict[str, str]] = None, ) -> dict[str, Callable[[], str]]: """Per-install registry: required_params name -> zero-arg resolver returning the raw value. @@ -86,6 +89,14 @@ def _resolve_workspace_dir() -> str: ) return workspace_dir + def _placed(name: str) -> str: + if place is None: + raise KeyError( + f"agent declares required_param {name!r}, which only an install in another agent's place " + "(the step's `stand_in` param) gives" + ) + return place[name] + def _resolve_container() -> str: if container_name is None: raise KeyError( @@ -102,6 +113,9 @@ def _resolve_container() -> str: "a2a_port": lambda: str(a2a_port), "agent_name": lambda: agent_name, "workspace_dir": _resolve_workspace_dir, + # Where the agent an install replaces works on the VM host, and as which user. + "workdir": lambda: _placed("workdir"), + "user": lambda: _placed("user"), # Config-sourced (from agent-env config) "litellm_api_key": lambda: (config.resolve_reference(litellm_api_key_override) if litellm_api_key_override else config.get_litellm_api_key()), @@ -225,6 +239,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: agent = A2AAgent.get(a2a_agent_id, self.a2a_agent_version) if problem := _install_problem(agent): raise RuntimeError(problem) + place = None + if stand_in := self.step_param_overrides(context).get("stand_in"): + place = await self._place_of(sandbox, agent) + agent = A2AAgent.get(stand_in) + if problem := _install_problem(agent): + raise RuntimeError(problem) logger.info(f"Loading agent image {agent.docker_image_artifact.image_name} onto VM (needed to read agent card)...") await sandbox.load_docker_images([agent.docker_image_artifact]) @@ -283,6 +303,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # A host install on a local sandbox runs on this machine itself, not in a container. sandbox_type=None if host_mode else sandbox.type, trust_lines=trust_lines, + place=place, ) resolved: dict[str, str] = {} for name in required_params: @@ -356,6 +377,30 @@ async def _trust_lines(sandbox, container: Optional[str], work_dir: str) -> str: env = trust_env(directory) return "".join(f"{name}={shlex.quote(value)}\n" for name, value in env.items()) + async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: + """Where ``agent``, which a stand-in replaces, works on the VM host and as which user: its install/v1 params + ``workdir_host`` and ``user_host``, templates of ``{agent_name}`` as its install commands are.""" + if self.container_name is not None: + raise RuntimeError( + f"install_agent '{self.id}': a stand-in replaces an agent installed on the VM host, not in a container" + ) + if sandbox.ON_THIS_MACHINE: + raise RuntimeError( + f"install_agent '{self.id}': a stand-in isn't installed on a sandbox that runs on this machine, where " + "its player's commands would run here; deploy the sandbox on a remote VM provider" + ) + logger.info(f"Loading agent image {agent.docker_image_artifact.image_name} onto VM (to read where it works)...") + await sandbox.load_docker_images([agent.docker_image_artifact]) + install_ext = await self._extract_install_extension(sandbox, agent.docker_image_artifact.image_name) + params = install_ext.get("params") or {} + place = {"workdir": params.get(WORKDIR_HOST_KEY), "user": params.get(USER_HOST_KEY)} + if not all(isinstance(value, str) and value for value in place.values()): + raise RuntimeError( + f"Agent {agent.id} v{agent.version} install/v1 params don't say where it works on the VM host: a " + f"stand-in installed in its place needs '{WORKDIR_HOST_KEY}' and '{USER_HOST_KEY}'" + ) + return {name: value.format(agent_name=self.agent_name) for name, value in place.items()} + def _register_instance(self, agent, a2a_url: str, ds, agent_card: dict): """Register the installed agent in the instance store, same as A2AAgent.deploy(), so store-resolving consumers (AddSkillsTaskStep, the a2a-agent CLIs) can find it.""" diff --git a/tst/integration/agent/stand_in_modal_vm_test.py b/tst/integration/agent/stand_in_modal_vm_test.py new file mode 100644 index 000000000..425671b63 --- /dev/null +++ b/tst/integration/agent/stand_in_modal_vm_test.py @@ -0,0 +1,81 @@ +"""An agent installed on a Modal VM's host, played from this machine through the stand-in installed in its place: the +shape of a task that loads its inputs onto the VM, installs a harness there and checks the outputs it leaves.""" + +import asyncio +import tempfile +import uuid +from pathlib import Path + +import modal +import pytest + +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.artifact import DockerImageArtifact, FileArtifact +from agent_env.task import EpisodeSpec, Task +from tst.util.capabilities import skip_without_remote_sandbox + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_remote_sandbox("modal_vm")] + +# An agent the stand-in replaces needs only a card declaring where it works on a VM host. +REPLACED_SERVER = """ +class _Card: + def model_dump(self): + return {"name": "replaced", "capabilities": {"extensions": [{"uri": "urn:agentenv:install/v1", "params": { + "workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", + "install_commands_host": ["false"], "a2a_port": 8000}}]}} + + +AGENT_CARD = _Card() +""" +REPLACED_DOCKERFILE = """FROM mirror.gcr.io/library/python:3.12-slim +WORKDIR /app +COPY a2a_server.py . +""" + + +@pytest.mark.asyncio +async def test_an_agent_installed_on_a_vm_host_is_played_through_a_stand_in_installed_in_its_place(): + uid = uuid.uuid4().hex[:8] + greeting = FileArtifact.put_bytes(id=f"greeting-vm-{uid}", description="stand-in on a VM", filename="greeting.txt", + content=b"hello", content_type="text/plain") + task = Task.from_dict({"id": f"copy-vm-{uid}", "steps": [ + {"id": "vm", "type": "deploy_sandbox", "sandbox_name": "vm", "sandbox_mode": "vm", "sandbox_type": "modal_vm", + "exposed_ports": [8000], "cpu": 2, "memory_mb": 8192, "ttl_seconds": 1800}, + {"id": "load", "type": "load_artifact", "sandbox_name": "vm", "artifact_id": greeting.id, + "destination_path": "/opt/agent-solver/workspace"}, + {"id": "install", "type": "install_agent", "sandbox_name": "vm", "a2a_agent_id": _replaced(uid).id, + "agent_name": "solver"}, + {"id": "solve", "type": "prompt_agent", "agent_name": "solver", "prompt_id": "solve", + "prompt": "Copy greeting.txt to outputs/answer.txt."}, + {"id": "check", "type": "verify_sandbox", "sandbox_name": "vm", "base_dir": "/opt/agent-solver/workspace", + "criteria": [{"type": "probe_file_contains", "criterion": "the greeting is copied", + "paths": ["outputs/answer.txt"], "expected": "^hello$"}]}, + ]}) + + episode = await EpisodeSpec.from_task(task).start() + try: + turn = await episode.next_turn() + async with turn.tools() as tools: + assert [tool.name for tool in tools.list()] == ["bash"] + assert (await tools.call("bash", {"command": "pwd; whoami; cat greeting.txt"})).text == ( + "/opt/agent-solver/workspace\nagent-solver\nhello") + await tools.call("bash", {"command": "mkdir -p outputs && cp greeting.txt outputs/answer.txt"}) + await turn.end("Copied.") + assert await episode.result() == {"check": 1.0} + finally: + await episode.abandon("the test ended") + [vm] = Task.get_instance(episode.id).context["deployed_sandboxes"] + sandbox = await modal.Sandbox.from_id.aio(vm["sandbox_id"]) + async with asyncio.timeout(60): + while await sandbox.poll.aio() is None: + await asyncio.sleep(2) + + +def _replaced(uid: str) -> A2AAgent: + with tempfile.TemporaryDirectory() as build_dir: + context = Path(build_dir) + (context / "a2a_server.py").write_text(REPLACED_SERVER) + (context / "Dockerfile").write_text(REPLACED_DOCKERFILE) + image = DockerImageArtifact.put_context(f"replaced-agent-image-{uid}", description="stand-in on a VM", + context_path=str(context), dockerfile_path="Dockerfile", platform=None) + return A2AAgent.put(id=f"replaced-agent-{uid}", docker_image_artifact=image) diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index 2481076c0..e15f4c1b5 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -211,7 +211,8 @@ async def verify(mcp_url): @pytest.mark.asyncio async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(local_stack, tmp_path): """A trainer's view of a seeded task: the env deployed and loaded and the echo agent prompted around the outside - agent, which plays its turn through the env's tools as its role and is never deployed; its reply recorded and + agent, which plays its turn through the env's tools as its role through a stand-in served in this process; its + reply recorded and scored under the verifier's step; and every episode, played, not played, abandoned or forgotten past its turn's deadline, leaving nothing behind.""" uid = uuid.uuid4().hex[:8] @@ -237,7 +238,9 @@ async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(loc async def play(episode, add): turn = await episode.next_turn() assert turn.prompt == "Add a crate." and turn.role == "alice@example.com" - assert [agent["agent_name"] for agent in Task.get_instance(episode.id).context["deployed_agents"]] == ["echo"] + assert [(agent["agent_name"], bool(agent["sandbox_id"])) + for agent in Task.get_instance(episode.id).context["deployed_agents"]] == [("policy", False), + ("echo", True)] if add: async with turn.tools() as tools: assert {"items_add_item", "items_whoami"} <= {tool.name for tool in tools.list()} @@ -252,7 +255,7 @@ async def play(episode, add): assert list(sandboxes.glob("agent-env-*")) == [] record = Task.get_instance(played.id) assert record.status == "completed" and record.context["metadata"]["episode"]["trainer_agents"] == ["policy"] - assert [agent["agent_name"] for agent in record.context["deployed_agents"]] == ["echo"] + assert [agent["agent_name"] for agent in record.context["deployed_agents"]] == ["policy", "echo"] replies = {response["step_id"]: response["response"] for response in record.context["prompt_responses"]} assert replies["play"] == "Added a crate." and replies["ask-echo"].startswith("Echo:") @@ -265,7 +268,7 @@ async def play(episode, add): forgotten = await EpisodeSpec.from_task(task, trainer_agents=["policy"], seed=seed, turn_timeout=3).start() await forgotten.next_turn() - with pytest.raises(TimeoutError, match="turn 'play' of agent 'policy' wasn't ended within 3s"): + with pytest.raises(TimeoutError, match="did not complete within 3s"): await forgotten.next_turn() assert Task.get_instance(forgotten.id).status == "failed" assert list(sandboxes.glob("agent-env-*")) == [] @@ -299,3 +302,43 @@ def _items_artifact(uid: str) -> EnvironmentArtifact: file_artifact=FileArtifact.put_bytes(id=f"server-items-file-{uid}", description="server env e2e", filename="items.json", content=json.dumps({"items": ["snap-x", "snap-y"]}).encode(), content_type="application/json"), ) + + +@pytest.mark.asyncio +async def test_an_agent_that_works_in_its_sandbox_is_played_through_a_stand_in_deployed_there(local_stack, tmp_path): + """A task that loads files into its agent and checks what the agent left there: the agent is never deployed, + a stand-in is, in its place, and the loop's ``bash`` runs in it, as its user and in its folder, with no model key + of its own; the check scores what the commands left, and nothing is left running.""" + uid = uuid.uuid4().hex[:8] + greeting = FileArtifact.put_bytes(id=f"greeting-{uid}", description="stand-in e2e", filename="greeting.txt", + content=b"hello", content_type="text/plain") + task = Task.from_dict({"id": f"copy-{uid}", "steps": [ + {"id": "agent", "type": "deploy_agent", "agent_name": "solver", "a2a_agent_id": f"never-deployed-{uid}"}, + {"id": "load", "type": "load_artifact", "agent_name": "solver", "artifact_id": greeting.id, + "destination_path": "/app/inputs"}, + {"id": "solve", "type": "prompt_agent", "agent_name": "solver", "prompt_id": "solve", + "prompt": "Copy inputs/greeting.txt to outputs/answer.txt.", "poll_interval_seconds": 0.5}, + {"id": "check", "type": "verify_sandbox", "agent_name": "solver", "criteria": [ + {"type": "probe_file_contains", "criterion": "the greeting is copied", "paths": ["outputs/answer.txt"], + "expected": "^hello$"}]}, + ]}) + + episode = await EpisodeSpec.from_task(task).start() + turn = await episode.next_turn() + async with turn.tools() as tools: + assert [tool.name for tool in tools.list()] == ["bash"] + assert (await tools.call("bash", {"command": "pwd; whoami; cat inputs/greeting.txt"})).text == ( + "/app\nagent\nhello") + assert (await tools.call("bash", {"command": "printenv LITELLM_API_KEY"})).text == "unused\n" + await tools.call("bash", {"command": "mkdir -p outputs && cp inputs/greeting.txt outputs/answer.txt"}) + await turn.end("Copied.") + + assert await episode.result() == {"check": 1.0} + record = Task.get_instance(episode.id) + [solver] = record.context["deployed_agents"] + assert solver["agent_name"] == "solver" and solver["sandbox_id"] + [response] = record.context["prompt_responses"] + assert (response["response"], response["tool_call_count"]) == ("Copied.", 3) + assert response["agent_trajectory_object_url"] + assert list((tmp_path / "sandboxes").glob("agent-env-*")) == [] + assert _docker("ps", "-aq", "--filter", f"name=^/agent-{solver['sandbox_id']}$").stdout.strip() == "" diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index ed7ef8742..2eb2c0023 100644 --- a/tst/integration/env/envs/server_env_modal_test.py +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -112,7 +112,7 @@ async def start(spec): forgotten = await start(EpisodeSpec.from_task(task, turn_timeout=3)) await forgotten.next_turn() - with pytest.raises(TimeoutError, match="wasn't ended within 3s"): + with pytest.raises(TimeoutError, match="did not complete within 3s"): await forgotten.next_turn() for episode in (played, abandoned, forgotten): diff --git a/tst/unit/a2a_agent/stand_in_test.py b/tst/unit/a2a_agent/stand_in_test.py new file mode 100644 index 000000000..84db9d249 --- /dev/null +++ b/tst/unit/a2a_agent/stand_in_test.py @@ -0,0 +1,38 @@ +"""The stand-in agent the Episode API deploys in place of an agent played from outside: written to the store as an +image's build context, and written again only when that context changes.""" + +from agent_env.a2a_agent import stand_in +from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS +from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent + + +def test_the_stand_in_is_a_build_context_given_no_model_key_or_store_credentials(local_stores, cli_routing): + agent = stand_in_agent() + + assert agent.id == STAND_IN_AGENT_ID + assert agent.docker_image_artifact.context_only + assert agent.default_env_vars == {"LITELLM_API_KEY": "unused", "LITELLM_BASE_URL": "http://unused.invalid"} + assert agent.metadata == {SHARE_STORE_CREDENTIALS: False} + + +def test_an_unchanged_context_reuses_the_stand_in_and_a_changed_one_writes_it_again( + local_stores, cli_routing, monkeypatch): + first = stand_in_agent() + assert stand_in_agent().version == first.version + + monkeypatch.setattr(stand_in, "_requirements", lambda: ["httpx==0.0.1"]) + + assert stand_in_agent().version == first.version + 1 + + +def test_its_context_holds_the_server_the_protocol_package_and_the_pinned_requirements_to_serve(tmp_path): + context = stand_in._stage(tmp_path) + + assert sorted(path.name for path in context.iterdir()) == [ + "Dockerfile", "a2a_server.py", "agentenv_protocol", "requirements.txt"] + assert (context / "agentenv_protocol" / "a2a_agent" / "stand_in.py").is_file() + assert not list(context.rglob("__pycache__")) + requirements = (context / "requirements.txt").read_text().splitlines() + assert {line.split("==")[0] for line in requirements} == { + "a2a-sdk[http-server]", "httpx", "pydantic", "python-multipart", "regex", "starlette", "uvicorn"} + assert all(line.split("==")[1] for line in requirements) diff --git a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py index e22b6e5c9..2b931d907 100644 --- a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py +++ b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py @@ -245,7 +245,7 @@ def test_a2a_merged_env_carries_contributions_and_they_win_over_caller_env(cf_cr from agent_env.a2a_agent.a2a_agent import A2AAgent - record = SimpleNamespace(default_env_vars={"AWS_ACCESS_KEY_ID": "preset"}) + record = SimpleNamespace(default_env_vars={"AWS_ACCESS_KEY_ID": "preset"}, metadata={}) merged = A2AAgent._build_merged_env(record, {_CF_ID: "caller-supplied"}, 8000) assert merged[_CF_ID] == "cf-id" assert merged[_CF_SECRET] == "cf-secret" diff --git a/tst/unit/store/aws/test_shared_credentials.py b/tst/unit/store/aws/test_shared_credentials.py index e67c7390b..65efc4602 100644 --- a/tst/unit/store/aws/test_shared_credentials.py +++ b/tst/unit/store/aws/test_shared_credentials.py @@ -5,7 +5,7 @@ import pytest from agentenv_protocol import client as protocol_v1 -from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS, A2AAgent from agent_env.config import set_object_store from agent_env.store import routing from agent_env.store.object_store import s3_object_store @@ -33,8 +33,9 @@ def host_credentials(monkeypatch): monkeypatch.setattr(s3_object_store.boto3, "Session", lambda: SimpleNamespace(get_credentials=lambda: creds)) -def _merged_env(caller_env: dict[str, str] | None = None) -> dict[str, str]: - return A2AAgent._build_merged_env(SimpleNamespace(default_env_vars={}), caller_env or {}, 8000) +def _merged_env(caller_env: dict[str, str] | None = None, metadata: dict | None = None) -> dict[str, str]: + return A2AAgent._build_merged_env(SimpleNamespace(default_env_vars={}, metadata=metadata or {}), caller_env or {}, + 8000) @pytest.mark.asyncio @@ -80,3 +81,8 @@ def test_a_caller_supplied_key_suppresses_the_whole_shared_set(host_credentials) merged = _merged_env({"AWS_ACCESS_KEY_ID": "AKIA-CALLER", "AWS_SECRET_ACCESS_KEY": "caller-secret"}) assert merged["AWS_ACCESS_KEY_ID"] == "AKIA-CALLER" assert "AWS_SESSION_TOKEN" not in merged + + +def test_an_agent_that_opts_out_gets_none_of_the_shared_set(host_credentials): + set_object_store(S3ObjectStore(object(), "bucket", region="us-west-2", share_credentials=True)) + assert not any(key.startswith("AWS_") for key in _merged_env(metadata={SHARE_STORE_CREDENTIALS: False})) diff --git a/tst/unit/task/episode/conftest.py b/tst/unit/task/episode/conftest.py new file mode 100644 index 000000000..15854fe44 --- /dev/null +++ b/tst/unit/task/episode/conftest.py @@ -0,0 +1,12 @@ +"""An episode's turns reach a stand-in agent this process serves on loopback, so these tests allow loopback +connections and nothing else.""" + +import pytest +from pytest_socket import enable_socket, socket_allow_hosts + + +@pytest.fixture(autouse=True) +def _disable_network(): + socket_allow_hosts(["127.0.0.1"], allow_unix_socket=True) + yield + enable_socket() diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index 4cd5ccdb3..f5ac223da 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -1,12 +1,14 @@ """The run of an episode, on fake steps: its turns, its scores, abandoning it, teardown and its record.""" import asyncio +import json import pytest import pytest_asyncio from agentenv_protocol import EnvironmentTool, ToolResult import agent_env.task.episode.run as episode_run +from agent_env.config import get_config from agent_env.config.runtime import Config from agent_env.env.env import DeployedEnv from agent_env.task import EpisodeError, EpisodeSpec, Task @@ -156,7 +158,7 @@ async def teardowns(monkeypatch, local_stores): ALICE = {"id": "alice", "type": "deploy_agent", "agent_name": "alice", "env_ids": ["items"], "role": "alice@example.com", "system_prompt": "You keep the items."} TURN = {"id": "turn", "type": "prompt_agent", "agent_name": "alice", "prompt_id": "turn", "prompt": "Add a crate.", - "timeout_seconds": 5} + "timeout_seconds": 5, "poll_interval_seconds": 0.05} REPLIES = {"id": "replies", "type": "episode_test_replies", "verifier_id": "replies"} @@ -194,7 +196,8 @@ async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(tear assert record.context["metadata"]["episode"] == {"id": episode.id, "trainer_agents": ["alice"]} assert [(r["step_id"], r["prompt_text"], r["response"]) for r in record.context["prompt_responses"]] == [ ("turn", "Add a crate.", "Added a crate.")] - assert record.context["deployed_agents"] == [] + assert [(agent["agent_name"], agent["sandbox_id"]) for agent in record.context["deployed_agents"]] == [ + ("alice", None)] @pytest.mark.asyncio @@ -322,32 +325,32 @@ async def test_a_setup_failure_reaches_the_caller_waiting_for_a_turn(teardowns): @pytest.mark.asyncio async def test_a_turn_not_taken_by_its_deadline_fails_the_run_and_is_never_handed_out(teardowns): - episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.05}, REPLIES).start() + episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 1}, REPLIES).start() await _ended(episode) - with pytest.raises(TimeoutError, match="turn 'turn' of agent 'alice' wasn't ended within 0.05s"): + with pytest.raises(TimeoutError, match="did not complete within 1s"): await episode.next_turn() assert Task.get_instance(episode.id).status == "failed" @pytest.mark.asyncio async def test_a_turn_ended_after_its_deadline_changes_nothing(teardowns): - episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 0.5}, REPLIES).start() + episode = await _spec(ITEMS, ALICE, {**TURN, "timeout_seconds": 1}, REPLIES).start() turn = await episode.next_turn() await _ended(episode) await turn.end("too late") - with pytest.raises(TimeoutError, match="wasn't ended within 0.5s"): + with pytest.raises(TimeoutError, match="did not complete within 1s"): await episode.result() record = Task.get_instance(episode.id) - assert record.status == "failed" and "wasn't ended within 0.5s" in record.error + assert record.status == "failed" and "did not complete within 1s" in record.error assert record.context["prompt_responses"] == [] @pytest.mark.asyncio async def test_a_tolerant_turn_past_its_deadline_is_scored_without_a_reply(teardowns): - tolerant = {**TURN, "timeout_seconds": 0.05, "fail_task_on_error": False} + tolerant = {**TURN, "timeout_seconds": 1, "fail_task_on_error": False} episode = await _spec(ITEMS, ALICE, tolerant, REPLIES).start() await _ended(episode) @@ -586,3 +589,95 @@ async def test_a_cancelled_start_takes_its_run_down_with_it(teardowns): assert not episode_run._live assert teardowns.finished == teardowns.started and len(teardowns.finished) == 1 + + +class _SandboxedStandIn: + """Stands in for the client of a stand-in deployed in a sandbox: runs each command as ``ran ``, and + records each reply.""" + + def __init__(self): + self.commands = [] + self.replies = [] + + async def run(self, command): + self.commands.append(command) + return f"ran {command}" + + async def reply(self, task_id, **reply): + self.replies.append((task_id, reply)) + return True + + +def _sandboxed_turn(stand_in): + return episode_run.Turn(agent="alice", step_id="turn", prompt="Add a crate.", system_prompt=None, role=None, + mcp_urls={"items": ITEMS_MCP}, skills=[], shell=True, stand_in=stand_in, task_id="t-1") + + +@pytest.mark.asyncio +async def test_a_turn_in_the_agents_sandbox_offers_bash_beside_its_envs_tools(monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) + stand_in = _SandboxedStandIn() + + async with _sandboxed_turn(stand_in).tools() as tools: + assert [tool.name for tool in tools.list()] == ["items_add_item", "bash"] + assert (await tools.call("bash", {"command": "ls inputs"})).text == "ran ls inputs" + with pytest.raises(EpisodeError, match="tool 'bash' takes a command string, not 3"): + await tools.call("bash", {"command": 3}) + + assert stand_in.commands == ["ls inputs"] + + +@pytest.mark.asyncio +async def test_an_env_tool_named_bash_clashes_with_the_sandboxes_own(monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["bash"]}) + + with pytest.raises(EpisodeError, match="tool 'bash' is offered by env 'items', and the agent's sandbox offers"): + async with _sandboxed_turn(_SandboxedStandIn()).tools(): + pass + + +@pytest.mark.asyncio +async def test_a_turns_calls_and_its_reply_are_handed_back_as_its_trajectory(monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) + stand_in = _SandboxedStandIn() + turn = _sandboxed_turn(stand_in) + + async with turn.tools() as tools: + await tools.call("items_add_item", {"name": "crate"}) + await tools.call("bash", {"command": "ls"}) + await turn.end("Added a crate.") + + [(task_id, reply)] = stand_in.replies + assert (task_id, reply["text"], reply["tool_call_count"]) == ("t-1", "Added a crate.", 2) + spans = reply["trajectory"] + assert [(span["name"], span["attributes"]["gen_ai.operation.name"]) for span in spans] == [ + ("items_add_item", "execute_tool"), ("bash", "execute_tool"), ("turn", "chain")] + assert json.loads(spans[0]["attributes"]["gen_ai.prompt"]) == {"input": {"name": "crate"}} + assert json.loads(spans[1]["attributes"]["gen_ai.completion"]) == {"output": [{"type": "text", "text": "ran ls"}]} + assert json.loads(spans[2]["attributes"]["gen_ai.completion"]) == { + "content": [{"type": "text", "text": "Added a crate."}]} + assert all(span["start_time"] <= span["end_time"] for span in spans) + + +@pytest.mark.asyncio +async def test_the_turns_trajectory_is_stored_as_any_agents_is(teardowns, monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() + + async with turn.tools() as tools: + await tools.call("items_add_item", {"name": "crate"}) + await turn.end("Added a crate.") + await episode.result() + + [response] = Task.get_instance(episode.id).context["prompt_responses"] + stored = json.loads(get_config().get_object_store().get(response["agent_trajectory_object_url"])) + assert [span["name"] for span in stored] == ["items_add_item", "turn"] + assert response["tool_call_count"] == 1 + + +@pytest.mark.asyncio +async def test_a_prompts_own_system_prompt_wins_over_the_agents(teardowns): + episode = await _spec(ITEMS, ALICE, {**TURN, "system_prompt": "Count first."}, REPLIES).start() + + assert (await episode.next_turn()).system_prompt == "Count first." diff --git a/tst/unit/task/episode/spec_test.py b/tst/unit/task/episode/spec_test.py index 14e4d4f2a..1246b673c 100644 --- a/tst/unit/task/episode/spec_test.py +++ b/tst/unit/task/episode/spec_test.py @@ -5,7 +5,7 @@ from agent_env.task import EpisodeError, EpisodeSpec, Task from agent_env.task.episode.rules import CHECKS, UNAFFECTED -from agent_env.task.episode.steps import _EpisodeAgentStep, _EpisodeTurnStep +from agent_env.task.episode.steps import _EpisodeAgentStep from agent_env.task.task import substitute_seed_in_steps from agent_env.task_step.registry import _builtin_registry from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep @@ -53,20 +53,7 @@ def test_an_agent_named_by_the_seed_is_the_one_played_from_outside(): assert spec.trainer_agents == ("planner",) assert [step.type for step in spec._copy.steps] == [ - "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier"] - - -def test_a_prompts_own_system_prompt_wins_over_the_agents(): - spec = EpisodeSpec.from_task(_portsim(play={**PLAY, "system_prompt": "Plan only the berths."})) - - assert spec._copy.steps[3].system_prompt == "Plan only the berths." - - -def test_each_prompt_of_a_trainer_agent_becomes_a_turn(): - again = {**PLAY, "id": "replan", "prompt_id": "replan", "prompt": "Replan.", "depends_on": ["play"]} - spec = EpisodeSpec.from_task(_portsim(again)) - - assert [step.id for step in spec._copy.steps if step.type == "episode_turn"] == ["play", "replan"] + "deploy_env", "load_artifact", "episode_agent", "prompt_agent", "env_outcome_verifier"] def test_a_task_prompting_several_agents_needs_the_ones_to_play_named(): @@ -97,60 +84,97 @@ def test_the_copy_stands_in_for_the_outside_agents_steps_and_keeps_everything_el assert (copy.id, copy.version) == ("dock", 4) assert [step.type for step in copy.steps] == [ - "deploy_env", "load_artifact", "episode_agent", "episode_turn", "env_outcome_verifier", "deploy_agent", + "deploy_env", "load_artifact", "episode_agent", "prompt_agent", "env_outcome_verifier", "deploy_agent", "prompt_agent", ] assert [copy.steps[i] is task.steps[i] for i in range(len(task.steps))] == [ - True, True, False, False, True, True, True] + True, True, False, True, True, True, True] for original, stand_in in zip(task.steps, copy.steps): assert (stand_in.id, stand_in.version, stand_in.depends_on, stand_in.fail_task_on_error) == ( original.id, original.version, original.depends_on, original.fail_task_on_error) -def test_the_stand_ins_carry_what_the_run_needs_unseeded_and_round_trip(): - agent, turn = EpisodeSpec.from_task(_portsim(), seed=SEED)._copy.steps[2:4] +def test_the_outside_agents_deploy_runs_in_a_stand_in_unseeded_and_round_trips(): + task = _portsim() + agent = EpisodeSpec.from_task(task, seed=SEED)._copy.steps[2] - assert isinstance(agent, _EpisodeAgentStep) and isinstance(turn, _EpisodeTurnStep) - assert (agent.agent_name, agent.env_ids, agent.env_step_id, agent.role) == ( - "default-agent", ["portsim"], None, "planner@") - assert (turn.agent_name, turn.prompt_id, turn.prompt, turn.system_prompt, turn.timeout_seconds) == ( - "default-agent", "play", "Plan week .", "You plan .", 600) - for step in (agent, turn): - assert type(step).from_dict(step.to_dict()).to_dict() == step.to_dict() + assert isinstance(agent, _EpisodeAgentStep) + assert (agent.step, agent.sandboxed) == (task.steps[2], False) + assert agent.step.role == "planner@" + assert _EpisodeAgentStep.from_dict(agent.to_dict()).to_dict() == agent.to_dict() -@pytest.mark.parametrize( - ("timeout", "deadline"), [(30, 30), (None, PromptAgentTaskStep.DEFAULT_PROMPT_TIMEOUT_SECONDS)]) -def test_a_turns_deadline_is_its_prompts_timeout(timeout, deadline): - turn = EpisodeSpec.from_task(_portsim(play={**PLAY, "timeout_seconds": timeout}))._copy.steps[3] +def test_without_a_turn_timeout_a_turns_deadline_is_its_prompts_own(): + task = _portsim(play={**PLAY, "timeout_seconds": 30}) - assert turn.timeout_seconds == deadline + assert EpisodeSpec.from_task(task)._copy.steps[3] is task.steps[3] -def test_a_turn_timeout_replaces_every_outside_prompts_own(): +@pytest.mark.parametrize(("turn_timeout", "deadline"), [(90, 90), (0.5, 1)]) +def test_a_turn_timeout_replaces_every_outside_prompts_own_in_whole_seconds(turn_timeout, deadline): again = {**PLAY, "id": "replan", "prompt_id": "replan", "timeout_seconds": 60} - copy = EpisodeSpec.from_task(_portsim(again, play={**PLAY, "timeout_seconds": 7200}), turn_timeout=90)._copy + task = _portsim(again, play={**PLAY, "timeout_seconds": 7200}) + copy = EpisodeSpec.from_task(task, turn_timeout=turn_timeout)._copy - assert [step.timeout_seconds for step in copy.steps if step.type == "episode_turn"] == [90, 90] + assert [step.timeout_seconds for step in copy.steps if step.type == "prompt_agent"] == [deadline, deadline] + assert copy.steps[3].to_dict() == {**task.steps[3].to_dict(), "timeout_seconds": deadline} def test_an_inside_agent_keeps_its_own_steps(): copy = EpisodeSpec.from_task(_portsim(BOB, ASK_BOB), trainer_agents=["bob"])._copy assert type(copy.steps[2]) is DeployAgentTaskStep and type(copy.steps[3]) is PromptAgentTaskStep - assert [step.type for step in copy.steps[5:]] == ["episode_agent", "episode_turn"] + assert [step.type for step in copy.steps[5:]] == ["episode_agent", "prompt_agent"] + + +INSTALL = {"id": "agent", "type": "install_agent", "sandbox_name": "box", "a2a_agent_id": "codex"} +IN_THE_SANDBOX = { + "an install": ({}, {"agent": INSTALL}), + "skills on the deploy": ({}, {"agent": {**AGENT, "skills": [{"name": "s", "skill_md": "Plan berths first."}]}}), + "add_skills": (({"id": "s", "type": "add_skills", "skills": [{"name": "s", "description": "d", "body": "x"}]},), + {}), + "a load into the agent": (({"id": "s", "type": "load_artifact", "artifact_id": "a", + "agent_name": "default-agent"},), {}), + "a collect from the agent": (({"id": "s", "type": "collect_artifacts", "artifact_paths": ["/out"]},), {}), + "run_code in the agent": (({"id": "s", "type": "run_code", "script_artifact_id": "s", "entrypoint": "main.py"},), + {}), + "verify_sandbox": (({"id": "s", "type": "verify_sandbox", "criteria": []},), {}), +} +OUTSIDE_THE_SANDBOX = { + "env tools alone": ((), {}), + "a load into the env": (({**LOAD, "id": "load-2"},), {}), + "a collect from the env": (({"id": "s", "type": "collect_artifacts", "env_id": "portsim", + "artifact_paths": ["/out"]},), {}), + "run_code in the env": (({"id": "s", "type": "run_code", "script_artifact_id": "s", "entrypoint": "main.py", + "env_id": "portsim"},), {}), + "verify_sandbox on a sandbox": (({"id": "s", "type": "verify_sandbox", "sandbox_name": "box", "criteria": []},), + {}), +} + + +@pytest.mark.parametrize(("extra", "replace"), IN_THE_SANDBOX.values(), ids=IN_THE_SANDBOX.keys()) +def test_a_task_using_the_outside_agents_sandbox_deploys_its_stand_in_there(extra, replace): + assert EpisodeSpec.from_task(_portsim(*extra, **replace))._copy.steps[2].sandboxed + + +@pytest.mark.parametrize(("extra", "replace"), OUTSIDE_THE_SANDBOX.values(), ids=OUTSIDE_THE_SANDBOX.keys()) +def test_a_task_using_only_the_outside_agents_envs_serves_its_stand_in_here(extra, replace): + assert not EpisodeSpec.from_task(_portsim(*extra, **replace))._copy.steps[2].sandboxed AT_THE_AGENT = {"agent_name": "default-agent"} REFUSED = { "no deploy_agent": ({}, {"agent": None}, "agent 'default-agent' is played from outside, so it needs exactly one " - "deploy_agent step, not 0"), - "two deploy_agents": (({**AGENT, "id": "agent-2"},), {}, "needs exactly one deploy_agent step, not 2"), + "deploy_agent or install_agent step, not 0"), + "two deploy_agents": (({**AGENT, "id": "agent-2"},), {}, + "needs exactly one deploy_agent or install_agent step, not 2"), + "a deploy and an install": (({**INSTALL, "id": "install"},), {}, + "needs exactly one deploy_agent or install_agent step, not 2"), "deploy_human_agent": (({"id": "human", "type": "deploy_human_agent", **AT_THE_AGENT},), {"agent": None}, "step 'human': deploy_human_agent deploys 'default-agent', which is played from outside"), - "install_agent": (({"id": "install", "type": "install_agent", "sandbox_name": "box", "a2a_agent_id": "x", - **AT_THE_AGENT},), {}, - "step 'install': install_agent deploys 'default-agent'"), + "an install into a container": ({}, {"agent": {**INSTALL, "container_name": "task"}}, + "step 'agent': an outside agent's stand-in replaces an agent installed on the VM " + "host, not in container 'task'"), "parts": ({}, {"play": {**PLAY, "prompt": None, "parts": [{"kind": "text", "text": "Plan."}]}}, "step 'play': an outside agent's prompt can't have parts, only a text prompt"), "user simulator": ({}, {"play": {**PLAY, "max_conversation_turns": 2}}, @@ -162,28 +186,16 @@ def test_an_inside_agent_keeps_its_own_steps(): "step 'play': an outside agent's prompt can't have agentenv_tools"), "snapshot_config": ({}, {"play": {**PLAY, "snapshot_config": {"env_id": "portsim"}}}, "step 'play': an outside agent's prompt can't have snapshot_config"), - "skills": ({}, {"agent": {**AGENT, "skills": [{"name": "s", "skill_md": "Plan berths first."}]}}, - "step 'agent': an outside agent can't have skills, instructions its player wouldn't get"), "snapshot files": ({}, {"agent": {**AGENT, "agent_snapshot_files_artifact_id": "files"}}, "step 'agent': an outside agent can't have agent_snapshot_files_artifact_id"), "changelog url": ({}, {"agent": {**AGENT, "agent_changelog_object_url": "file:///changelog"}}, "step 'agent': an outside agent can't have agent_changelog_object_url"), "changelog": ({}, {"agent": {**AGENT, "enable_agent_changelog": True}}, "step 'agent': an outside agent can't have enable_agent_changelog"), - "add_skills": (({"id": "s", "type": "add_skills", "skills": [{"name": "s", "description": "d", "body": "x"}]},), {}, - "step 's': add_skills needs 'default-agent' as a real agent, and it is played from outside"), "register_agent_triggers": (({"id": "s", "type": "register_agent_triggers", "triggers": [], **AT_THE_AGENT},), {}, "step 's': register_agent_triggers needs 'default-agent' as a real agent"), "snapshot_agent_state": (({"id": "s", "type": "snapshot_agent_state", "artifact_id": "a", "prompt_id": "play"},), {}, "step 's': snapshot_agent_state needs 'default-agent' as a real agent"), - "load_artifact": (({"id": "s", "type": "load_artifact", "artifact_id": "a", **AT_THE_AGENT},), {}, - "step 's': load_artifact needs 'default-agent' as a real agent"), - "verify_sandbox": (({"id": "s", "type": "verify_sandbox", "criteria": []},), {}, - "step 's': verify_sandbox needs 'default-agent' as a real agent"), - "run_code": (({"id": "s", "type": "run_code", "script_artifact_id": "s", "entrypoint": "main.py"},), {}, - "step 's': run_code needs 'default-agent' as a real agent"), - "collect from the agent": (({"id": "s", "type": "collect_artifacts", "artifact_paths": ["/out"]},), {}, - "step 's': collects from 'default-agent''s sandbox, and it is played from outside"), "collect a manifest": (({"id": "s", "type": "collect_artifacts", "env_id": "portsim", "manifest_step_id": "play"},), {}, "step 's': its manifest is outside prompt 'play'"), "trigger executor": (({"id": "s", "type": "register_env_triggers", "env_id": "portsim", "triggers": [], @@ -192,10 +204,6 @@ def test_an_inside_agent_keeps_its_own_steps(): "peered": ((BOB, {"id": "s", "type": "peer_agents", "peerings": [ {"source_agent_name": "bob", "peer_agent_names": ["default-agent"]}]}), {}, "step 's': peers 'default-agent', which is played from outside"), - "rubric on the trajectory": (({"id": "s", "type": "rubrics_verifier", "prompt_id": "play", "criteria": [], - "use_agent_judge": False},), {}, - "step 's': grades an outside prompt's trajectory, which an outside player doesn't " - "hand back"), "rubric judged in the agent": (({"id": "s", "type": "rubrics_verifier", "prompt_id": "other", "criteria": [], "use_trajectory": False, **AT_THE_AGENT},), {}, "step 's': judges in 'default-agent''s sandbox, and it is played from outside"), @@ -238,6 +246,12 @@ def test_every_refusal_is_reported_at_once(): "use_trajectory": False, "use_agent_judge": False}},), "collecting from an env": ({"collect": {"id": "collect", "type": "collect_artifacts", "env_id": "portsim", "artifact_paths": ["/out"]}},), + "a rubric of the trajectory": ( + {"verify": {"id": "rubric", "type": "rubrics_verifier", "prompt_id": "play", "criteria": [], + "use_agent_judge": False}},), + "installing the outside agent on a VM host": ({"agent": INSTALL},), + **{name: ({"agent": replace["agent"]} if replace else {"extra": extra[0]},) + for name, (extra, replace) in IN_THE_SANDBOX.items() if name != "an install"}, "tearing down an outside agent": ({"down": {"id": "down", "type": "teardown_sandboxes", "agent_names": ["default-agent"]}},), } diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index 17ba06915..c2525076e 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -129,6 +129,8 @@ def test_a2a_port_roundtrip(): class _InstallSandbox: + ON_THIS_MACHINE = False + def __init__(self): self.scripts: list[str] = [] @@ -310,3 +312,81 @@ def test_the_trust_lines_resolve_as_the_trust_env_param(): resolvers = _build_param_resolvers(container_name="c", work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, config=None, agent_name="a", trust_lines="SSL_CERT_FILE=/x\n") assert resolvers["trust_env"]() == "SSL_CERT_FILE=/x\n" +STAND_IN = "@local/agentenv-stand-in" +CODEX_INSTALL = {"params": {"workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", + "install_commands_host": ["install codex"], "a2a_port": 8000}} +STAND_IN_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], + "required_params_host": ["user", "workdir", "a2a_port"], "a2a_port": 8000}} + + +def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, container_name=None): + """An install of ``agent-x`` that the step's ``stand_in`` param replaces with the stand-in.""" + sandbox = sandbox or _InstallSandbox() + monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", + lambda name: SimpleNamespace(get_sandbox=_resolved(sandbox))) + agents = {agent_id: SimpleNamespace(id=agent_id, version=1, metadata={}, docker_image_artifact=SimpleNamespace( + image_name=f"{agent_id}-image", build_context_object_url=f"s3://bucket/{agent_id}.tar.gz")) + for agent_id in ("agent-x", STAND_IN)} + monkeypatch.setattr(A2AAgent, "get", staticmethod(lambda agent_id, version=None: agents[agent_id])) + step = InstallAgentTaskStep(id="i", version=None, sandbox_name="vm", a2a_agent_id="agent-x", agent_name="solver", + container_name=container_name) + cards = {"agent-x-image": codex, f"{STAND_IN}-image": STAND_IN_INSTALL} + monkeypatch.setattr(step, "_extract_install_extension", lambda sandbox, image: _resolved(cards[image])()) + monkeypatch.setattr(step, "_wait_for_agent_card", _resolved({})) + monkeypatch.setattr(step, "_register_instance", lambda *args: SimpleNamespace(instance_id="solver-1")) + context = TaskStepContext( + deployed_sandboxes=[DeployedSandbox( + sandbox_name="vm", sandbox_id="sb-1", sandbox_mode="vm", sandbox_type="modal_vm", + tunnel_urls={"8000": "https://vm-8000.example"})], + metadata={"user_overrides": {"step_params": {"i": {"stand_in": STAND_IN}}}, + "deployed_docker_containers": [{"container_name": "task", "sandbox_name": "vm"}]}, + ) + return step, context, sandbox + + +@pytest.mark.asyncio +async def test_a_stand_in_is_installed_where_the_agent_it_replaces_says_it_works(monkeypatch, local_stores): + step, context, sandbox = _installing_in_place(monkeypatch) + + await step.execute(context) + + assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ + "serve as agent-solver in /opt/agent-solver/workspace on 8000"] + assert [agent.agent_name for agent in context.deployed_agents] == ["solver"] + + +@pytest.mark.asyncio +async def test_a_stand_in_needs_the_replaced_agent_to_say_where_it_works(monkeypatch, local_stores): + step, context, _ = _installing_in_place(monkeypatch, codex={"params": {"install_commands_host": ["install codex"], + "workdir_host": "/opt/agent"}}) + + with pytest.raises(RuntimeError, match="needs 'workdir_host' and 'user_host'"): + await step.execute(context) + + +@pytest.mark.asyncio +async def test_a_stand_in_isnt_installed_into_a_container(monkeypatch, local_stores): + step, context, _ = _installing_in_place(monkeypatch, container_name="task") + + with pytest.raises(RuntimeError, match="installed on the VM host, not in a container"): + await step.execute(context) + + +@pytest.mark.asyncio +async def test_a_stand_in_isnt_installed_on_a_sandbox_that_runs_on_this_machine(monkeypatch, local_stores): + here = _InstallSandbox() + here.ON_THIS_MACHINE = True + step, context, _ = _installing_in_place(monkeypatch, sandbox=here) + + with pytest.raises(RuntimeError, match="runs on this machine"): + await step.execute(context) + + +def test_only_an_install_in_another_agents_place_gives_its_workdir_and_user(): + with pytest.raises(KeyError, match="stand_in"): + _resolvers(None)["workdir"]() + placed = _build_param_resolvers( + container_name=None, work_dir="/tmp/w", agent_ctx_tar="/tmp/w/c.tgz", a2a_port=8000, config=object(), + agent_name="solver", place={"workdir": "/opt/agent-solver/workspace", "user": "agent-solver"}, + ) + assert (placed["workdir"](), placed["user"]()) == ("/opt/agent-solver/workspace", "agent-solver") From f0c609b002a0543994e7f5695a8374eddda6db91 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 19 Oct 2026 15:12:37 -0700 Subject: [PATCH 02/39] docs: an episode plays its outside agents through a stand-in agent Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 25ae1d162..22a0da6f2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -202,11 +202,15 @@ orchestrators call it the same way. What callers rely on: - A runner that must hear from a long step puts a zero-argument callable at `context.metadata[HEARTBEAT_KEY]` (`task_step/context.py`); a step that blocks for long calls `context.heartbeat()`, a no-op without one. The key is in `_REDACTED_KEYS`, so it is never persisted. -- An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task in which - the `deploy_agent` and `prompt_agent` steps of each of its `trainer_agents` are the plain-data - `episode_agent` and `episode_turn`, which wait for the caller instead of deploying or prompting. Its - record is tagged `context.metadata["episode"]` and names the original task, so it isn't resumable - with `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded +- An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task, with + namespace routing on, in which each of its `trainer_agents` is played through a stand-in agent + (`agentenv_protocol.a2a_agent.stand_in`, written as `@local/agentenv-stand-in` from its build context). + The agent's `deploy_agent` or `install_agent`, wrapped as `episode_agent`, deploys the stand-in in the + agent's place when the task uses the agent's sandbox and serves it in this process when it doesn't; + the agent's `prompt_agent` steps message it as they would the agent, and the caller plays each + message as a `Turn` through its `urn:agentenv:driver/v1`, where a turn's `bash` runs. Its record is + tagged `context.metadata["episode"]` and names the original task, so it isn't resumable with + `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded `cancelled`. ## Tests From c7773c77240f26599c2433466949320b30c6b7c9 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 19 Oct 2026 15:13:40 -0700 Subject: [PATCH 03/39] fix(task): a turn's reply that can't reach its stand-in leaves the failure to the run Turn.end() never raised before the stand-in; a reply the stand-in can't take is logged, and the run's own failure reaches the caller through next_turn() and result(). Skill, the type of Turn.skills, is exported beside Turn. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/__init__.py | 3 ++- src/agent_env/task/episode/__init__.py | 4 ++-- src/agent_env/task/episode/run.py | 13 ++++++++++--- tst/unit/task/episode/run_test.py | 14 ++++++++++++++ 4 files changed, 28 insertions(+), 6 deletions(-) diff --git a/src/agent_env/task/__init__.py b/src/agent_env/task/__init__.py index f7d2f7c27..3c5b634de 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,7 +1,7 @@ """Task module for AgentEnv.""" from .task import Task -from .episode import Episode, EpisodeError, EpisodeSpec, Turn +from .episode import Episode, EpisodeError, EpisodeSpec, Skill, Turn from .store import ( TaskStore, TaskQuery, get_task_store, set_task_store, reset_task_store, TaskInstance, TaskInstanceStore, get_task_instance_store, set_task_instance_store, reset_task_instance_store, @@ -13,6 +13,7 @@ "Episode", "EpisodeError", "EpisodeSpec", + "Skill", "Turn", "TaskStore", "TaskQuery", diff --git a/src/agent_env/task/episode/__init__.py b/src/agent_env/task/episode/__init__.py index bd00530f0..1b8fd59d8 100644 --- a/src/agent_env/task/episode/__init__.py +++ b/src/agent_env/task/episode/__init__.py @@ -3,7 +3,7 @@ AgentEnv runs everything else exactly as ``Task.run`` does, so its scores are the ones ``agent-env run`` would report. """ -from .run import Episode, EpisodeError, Turn +from .run import Episode, EpisodeError, Skill, Turn from .spec import EpisodeSpec -__all__ = ["Episode", "EpisodeError", "EpisodeSpec", "Turn"] +__all__ = ["Episode", "EpisodeError", "EpisodeSpec", "Skill", "Turn"] diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index 86abb2810..c6e4a4641 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -78,19 +78,26 @@ def tools(self, *, timeout: float = 600.0) -> TurnTools: async def end(self, response: str | None = None) -> None: """Hand the turn back, and the run carries on. No reply records an empty one, as an agent that answers with - nothing leaves. Once the turn is over, ended or not, this does nothing.""" + nothing leaves. Once the turn is over, ended or not, this does nothing; a stand-in that can't be reached + fails the run, which ``next_turn()`` and ``result()`` report.""" if not self._open: return self._open = False reply = response or "" spans = [*self._spans, _span("turn", "chain", self._started, completion={"content": [ {"type": "text", "text": reply}]})] - await self._stand_in.reply(self._task_id, text=reply, trajectory=spans, tool_call_count=len(self._spans)) + await self._reply(text=reply, trajectory=spans, tool_call_count=len(self._spans)) async def _fail(self, error: Exception) -> None: if self._open: self._open = False - await self._stand_in.reply(self._task_id, error=str(error)) + await self._reply(error=str(error)) + + async def _reply(self, **reply: Any) -> None: + try: + await self._stand_in.reply(self._task_id, **reply) + except httpx.HTTPError as e: + logger.warning("turn %r of agent %r: its reply didn't reach the stand-in: %s", self.step_id, self.agent, e) class TurnTools: diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index f5ac223da..35035f812 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -3,6 +3,7 @@ import asyncio import json +import httpx import pytest import pytest_asyncio from agentenv_protocol import EnvironmentTool, ToolResult @@ -681,3 +682,16 @@ async def test_a_prompts_own_system_prompt_wins_over_the_agents(teardowns): episode = await _spec(ITEMS, ALICE, {**TURN, "system_prompt": "Count first."}, REPLIES).start() assert (await episode.next_turn()).system_prompt == "Count first." + + +@pytest.mark.asyncio +async def test_a_reply_that_cant_reach_the_stand_in_leaves_the_runs_own_failure_to_report(caplog): + class _Unreachable(_SandboxedStandIn): + async def reply(self, task_id, **reply): + raise httpx.ConnectError("the sandbox is gone") + + turn = _sandboxed_turn(_Unreachable()) + + await turn.end("done") + + assert "its reply didn't reach the stand-in: the sandbox is gone" in caplog.text From 4072d2c4f804b00044981e86696b1b7922d1a7a0 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 20 Oct 2026 14:20:59 -0700 Subject: [PATCH 04/39] test: the stand-in's pinned requirements follow the protocol package's agent extra Co-Authored-By: Claude Opus 5.5 --- tst/unit/a2a_agent/stand_in_test.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tst/unit/a2a_agent/stand_in_test.py b/tst/unit/a2a_agent/stand_in_test.py index 84db9d249..11eaebb88 100644 --- a/tst/unit/a2a_agent/stand_in_test.py +++ b/tst/unit/a2a_agent/stand_in_test.py @@ -1,6 +1,8 @@ """The stand-in agent the Episode API deploys in place of an agent played from outside: written to the store as an image's build context, and written again only when that context changes.""" +from importlib.metadata import version + from agent_env.a2a_agent import stand_in from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent @@ -32,7 +34,7 @@ def test_its_context_holds_the_server_the_protocol_package_and_the_pinned_requir "Dockerfile", "a2a_server.py", "agentenv_protocol", "requirements.txt"] assert (context / "agentenv_protocol" / "a2a_agent" / "stand_in.py").is_file() assert not list(context.rglob("__pycache__")) - requirements = (context / "requirements.txt").read_text().splitlines() - assert {line.split("==")[0] for line in requirements} == { - "a2a-sdk[http-server]", "httpx", "pydantic", "python-multipart", "regex", "starlette", "uvicorn"} - assert all(line.split("==")[1] for line in requirements) + pinned = dict(line.split("==") for line in (context / "requirements.txt").read_text().splitlines()) + assert {"a2a-sdk[http-server]", "httpx", "pydantic", "starlette", "uvicorn"} <= pinned.keys() + assert not pinned.keys() & {"fastmcp", "pytest", "pytest-asyncio"} + assert all(pinned[name] == version(name.split("[")[0]) for name in pinned) From 0ad10aca647aee9251a33c514cd7269717543948 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 20 Oct 2026 14:21:01 -0700 Subject: [PATCH 05/39] feat(agents): Claude Code declares where it works on a VM host, so a stand-in can take its place Co-Authored-By: Claude Opus 5.5 --- src/agent_env/agents/agents/claude-code/a2a_server.py | 3 +++ tst/unit/agents/claude_code/test_card_contract.py | 1 + 2 files changed, 4 insertions(+) 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 7d13710c8..f209bd43a 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -375,6 +375,9 @@ async def _turns_stop_with_the_server(app: Starlette) -> AsyncIterator[None]: required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], + # Where an install on the host works, and as whom, for a stand-in installed in this agent's place. + workdir_host="/app", + user_host="root", ), enable(TRIGGERS_V1, description=( "Register a typed trigger script and decide the stakeholder reply per turn " diff --git a/tst/unit/agents/claude_code/test_card_contract.py b/tst/unit/agents/claude_code/test_card_contract.py index bd964ead3..1b48c9104 100644 --- a/tst/unit/agents/claude_code/test_card_contract.py +++ b/tst/unit/agents/claude_code/test_card_contract.py @@ -141,6 +141,7 @@ def test_install_agent_reads_the_card_and_its_install_commands_from_the_image(): (install,) = [ext for ext in json.loads(out.stdout)["capabilities"]["extensions"] if ext["uri"] == "urn:agentenv:install/v1"] assert install["params"]["install_commands"] and install["params"]["install_commands_host"] + assert (install["params"]["workdir_host"], install["params"]["user_host"]) == ("/app", "root") def test_each_install_mode_asks_only_for_params_install_agent_supplies(app): From 89f819fd1de3ab6ab96304fc1ccd590b0125fdd8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 20 Oct 2026 15:13:07 -0700 Subject: [PATCH 06/39] fix(episode): keep credentials from the stand-in, record failed calls, close turns their step gave up on - An agent whose metadata sets takes_credentials False is given no model key, endpoint or headers, from any source, and none of the object store's shared credentials. The stand-in sets it and carries no placeholder keys. - A tool call that raises is recorded in the turn's trajectory with its error. - A message on a conversation maps to the first of its prompt steps not yet played. - A command that times out or is cancelled has its whole process group killed, so a background child holding its output can't keep it waiting. - A turn whose prompt_agent step has failed is closed at once, rather than a long-poll later when the stand-in drops it, so it is never handed out once the run's record says failed; abandon() closes the turns of the run it cancels. Co-Authored-By: Claude Opus 5.5 --- .../agentenv_protocol/a2a_agent/stand_in.py | 17 +++---- .../tests/test_a2a_stand_in.py | 10 ++++ src/agent_env/a2a_agent/a2a_agent.py | 35 +++++++++----- src/agent_env/a2a_agent/stand_in/__init__.py | 8 +--- src/agent_env/task/episode/run.py | 48 +++++++++++++------ .../env/envs/server_env_local_test.py | 6 +-- tst/unit/a2a_agent/stand_in_test.py | 20 ++++++-- tst/unit/store/aws/test_shared_credentials.py | 4 +- tst/unit/task/episode/run_test.py | 48 +++++++++++++++++-- 9 files changed, 142 insertions(+), 54 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py index 6d5b5eb5c..00e915be2 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py @@ -193,20 +193,21 @@ def _bump(self) -> None: async def _run(command: str) -> str: """``command``'s stdout then its stderr, as verifiers' bash tool returns them. It runs in a process group of its - own, stopped with whatever it started if it outlives ``COMMAND_TIMEOUT_SECONDS``.""" + own, stopped with whatever it started if its output isn't done within ``COMMAND_TIMEOUT_SECONDS``, or if the call + is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" process = await asyncio.create_subprocess_exec( "bash", "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True, ) try: stdout, stderr = await asyncio.wait_for(process.communicate(), COMMAND_TIMEOUT_SECONDS) - except asyncio.TimeoutError: - return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" - finally: - if process.returncode is None: - with contextlib.suppress(ProcessLookupError): - os.killpg(process.pid, signal.SIGKILL) - await process.wait() + except BaseException as e: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + if isinstance(e, asyncio.TimeoutError): + return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" + raise return (stdout + stderr).decode(errors="replace") diff --git a/packages/agentenv-protocol/tests/test_a2a_stand_in.py b/packages/agentenv-protocol/tests/test_a2a_stand_in.py index 1908c6859..50b7189d9 100644 --- a/packages/agentenv-protocol/tests/test_a2a_stand_in.py +++ b/packages/agentenv-protocol/tests/test_a2a_stand_in.py @@ -161,6 +161,16 @@ def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(cli os.kill(int((tmp_path / "child").read_text()), 0) +def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(client, monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + + output = _ran(client, "sleep 60 & echo $! > child") + + assert output == "error: the command timed out after 0.5s" + with pytest.raises(ProcessLookupError): + os.kill(int((tmp_path / "child").read_text()), 0) + + def test_output_names_a_command_it_started(client): assert _driver(client, "output", exec_id="nope").status_code == 404 diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index d54a57946..b65b34107 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -63,7 +63,10 @@ _DIND_CONTAINER = "agent-dind" _DIND_NETWORK = "agent-docker-net" _DIND_PORT = 2375 -SHARE_STORE_CREDENTIALS = "share_store_credentials" +# Metadata an agent whose commands its player runs sets False: it is given no model endpoint, key or headers, whatever +# a deploy passes, and none of the object store's shared credentials. +TAKES_CREDENTIALS = "takes_credentials" +MODEL_CREDENTIALS_ENV = ("LITELLM_API_KEY", "LITELLM_BASE_URL", MODEL_HEADERS_ENV) def unwritten_local_agent(agent_id: str, field: str) -> str: @@ -455,16 +458,7 @@ async def deploy( provider = None if sandbox is not None else ( build_sandbox_provider(sandbox_type) if sandbox_type else get_agent_sandbox_provider() ) - resolved_env = dict(env_vars) if env_vars else {} - if "LITELLM_API_KEY" not in resolved_env and "LITELLM_API_KEY" not in self.default_env_vars: - resolved_env["LITELLM_API_KEY"] = config.get_litellm_api_key() - if "LITELLM_BASE_URL" not in resolved_env and "LITELLM_BASE_URL" not in self.default_env_vars: - resolved_env["LITELLM_BASE_URL"] = config.get_litellm_base_url() - if MODEL_HEADERS_ENV not in resolved_env and MODEL_HEADERS_ENV not in self.default_env_vars: - headers = config.get_model_headers( - resolved_env.get("LITELLM_BASE_URL", self.default_env_vars.get("LITELLM_BASE_URL"))) - if headers: - resolved_env[MODEL_HEADERS_ENV] = json.dumps(headers) + resolved_env = self._with_model_credentials(env_vars or {}, config) if "LITELLM_BASE_URL" in resolved_env: resolved_env["LITELLM_BASE_URL"] = (provider or sandbox).url_from_sandbox(resolved_env["LITELLM_BASE_URL"]) @@ -545,12 +539,29 @@ def _refusal(self, image: DockerImageArtifact, link: SandboxProvider) -> str | N return f"{how}: {problem}" return None + def _with_model_credentials(self, env_vars: dict[str, str], config) -> dict[str, str]: + """``env_vars`` with the model endpoint, key and headers each taken from config when neither they nor + ``default_env_vars`` set it; for an agent that takes no credentials, with none of them, whatever set them.""" + if not self.metadata.get(TAKES_CREDENTIALS, True): + return {name: value for name, value in env_vars.items() if name not in MODEL_CREDENTIALS_ENV} + resolved_env = dict(env_vars) + if "LITELLM_API_KEY" not in resolved_env and "LITELLM_API_KEY" not in self.default_env_vars: + resolved_env["LITELLM_API_KEY"] = config.get_litellm_api_key() + if "LITELLM_BASE_URL" not in resolved_env and "LITELLM_BASE_URL" not in self.default_env_vars: + resolved_env["LITELLM_BASE_URL"] = config.get_litellm_base_url() + if MODEL_HEADERS_ENV not in resolved_env and MODEL_HEADERS_ENV not in self.default_env_vars: + headers = config.get_model_headers( + resolved_env.get("LITELLM_BASE_URL", self.default_env_vars.get("LITELLM_BASE_URL"))) + if headers: + resolved_env[MODEL_HEADERS_ENV] = json.dumps(headers) + return resolved_env + def _build_merged_env(self, resolved_env: dict[str, str], a2a_port: int) -> dict[str, str]: merged_env = dict(self.default_env_vars) merged_env.update(resolved_env) merged_env["A2A_PORT"] = str(a2a_port) shared = get_config().get_object_store().shared_credentials_env() - if self.metadata.get(SHARE_STORE_CREDENTIALS, True) and shared.keys().isdisjoint(merged_env): + if self.metadata.get(TAKES_CREDENTIALS, True) and shared.keys().isdisjoint(merged_env): merged_env.update(shared) merged_env.update(all_sandbox_container_env()) diff --git a/src/agent_env/a2a_agent/stand_in/__init__.py b/src/agent_env/a2a_agent/stand_in/__init__.py index cbf38df10..40f014b9e 100644 --- a/src/agent_env/a2a_agent/stand_in/__init__.py +++ b/src/agent_env/a2a_agent/stand_in/__init__.py @@ -14,7 +14,7 @@ import agentenv_protocol from packaging.requirements import Requirement -from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS, A2AAgent +from agent_env.a2a_agent.a2a_agent import TAKES_CREDENTIALS, A2AAgent from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.store.base import NotFoundError @@ -41,11 +41,7 @@ def stand_in_agent() -> A2AAgent: _IMAGE_ID, description="The stand-in agent: a driver outside its sandbox answers its turns.", context_path=str(context), dockerfile_path=dockerfile, platform=None, ) - return A2AAgent.put( - id=STAND_IN_AGENT_ID, docker_image_artifact=image, - default_env_vars={"LITELLM_API_KEY": "unused", "LITELLM_BASE_URL": "http://unused.invalid"}, - metadata={SHARE_STORE_CREDENTIALS: False}, - ) + return A2AAgent.put(id=STAND_IN_AGENT_ID, docker_image_artifact=image, metadata={TAKES_CREDENTIALS: False}) def _stage(context: Path) -> Path: diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index c6e4a4641..050df47ba 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -144,16 +144,22 @@ async def call(self, name: str, arguments: dict[str, Any] | None = None) -> Tool if name not in self._offered_by: raise EpisodeError( f"no env of turn {self._turn.step_id!r} offers tool {name!r}; they offer {sorted(self._offered_by)}") - started = _now() - env_id = self._offered_by[name] - if env_id is None: - result = await self._bash(arguments or {}) - else: - result = await self._sessions[env_id].call_tool(name, arguments) - self._turn._spans.append(_span(name, "execute_tool", started, prompt={"input": arguments or {}}, - completion={"output": result.content})) + started, env_id = _now(), self._offered_by[name] + try: + if env_id is None: + result = await self._bash(arguments or {}) + else: + result = await self._sessions[env_id].call_tool(name, arguments) + except Exception as e: + self._record(name, arguments, started, {"error": f"{type(e).__name__}: {e}"}) + raise + self._record(name, arguments, started, {"output": result.content}) return result + def _record(self, name: str, arguments: dict[str, Any] | None, started: str, completion: dict[str, Any]) -> None: + self._turn._spans.append(_span(name, "execute_tool", started, prompt={"input": arguments or {}}, + completion=completion)) + async def _bash(self, arguments: dict[str, Any]) -> ToolResult: if not isinstance(command := arguments.get("command"), str): raise EpisodeError(f"tool {BASH.name!r} takes a command string, not {command!r}") @@ -181,6 +187,7 @@ def __init__(self, id: str, task: Task, context: TaskStepContext): self._abandoned: str | None = None self._stand_ins: list[StandInClient] = [] self._drains: list[asyncio.Task[None]] = [] + self._played: set[str] = set() self._in_process: list[InProcessStandIn] = [] async def next_turn(self) -> Turn | None: @@ -188,6 +195,7 @@ async def next_turn(self) -> Turn | None: failed first. End the turns you hold before waiting on one that depends on them.""" while True: changed = self._changed + self._close_turns_of_failed_steps() while self._waiting and self._abandoned is None: turn = self._waiting.popleft() if turn._open: @@ -202,6 +210,7 @@ async def result(self) -> dict[str, float]: over: what ``agent-env run`` reports. Making a reward of them is the caller's.""" while not self._run.done(): changed = self._changed + self._close_turns_of_failed_steps() if open_turn := next((turn for turn in self._turns if turn._open), None): raise EpisodeError( f"turn {open_turn.step_id!r} of agent {open_turn.agent!r} isn't ended, so the run can't finish" @@ -220,6 +229,8 @@ async def abandon(self, reason: str) -> None: await turn._fail(EpisodeError(f"episode abandoned: {reason}")) else: self._abandoned = reason + for turn in self._turns: + turn._open = False self._run.cancel(msg=reason) await asyncio.wait((self._run,)) @@ -260,6 +271,13 @@ async def _wait(self, changed: asyncio.Future[None]) -> None: """Until the run ends or what was ``changed`` when the caller last looked changes again.""" await asyncio.wait((self._run, changed), return_when=asyncio.FIRST_COMPLETED) + def _close_turns_of_failed_steps(self) -> None: + """A step that failed has given up on its turns, which its stand-in drops only a long-poll later.""" + failed = {failure.get("step_id") for failure in self._context.metadata.get("failed_steps") or []} + for turn in self._turns: + if turn.step_id in failed: + turn._open = False + def _raise_unless_completed(self) -> None: if self._abandoned is not None: raise EpisodeError(f"episode abandoned: {self._abandoned}") @@ -319,13 +337,15 @@ async def _drain( self._changed = asyncio.get_running_loop().create_future() def _step_of(self, agent: str, context_id: str) -> str: - """The ``prompt_agent`` step whose conversation with ``agent`` a message on ``context_id`` is in.""" + """The ``prompt_agent`` step a message to ``agent`` on ``context_id`` comes from: of its prompts in that + conversation, the first not yet played, since a prompt continuing a conversation runs after the one it + continues.""" conversations = self._context.metadata.get("a2a_conversations", {}) - for step in self._seeded_steps(): - if step.type == "prompt_agent" and step.agent_name == agent and ( - step.context_id or conversations.get(step.id)) == context_id: - return step.id - return context_id + steps = [step.id for step in self._seeded_steps() if step.type == "prompt_agent" and step.agent_name == agent + and (step.context_id or conversations.get(step.id)) == context_id] + step_id = next((step_id for step_id in steps if step_id not in self._played), steps[-1] if steps else context_id) + self._played.add(step_id) + return step_id def _seeded_steps(self) -> list: return substitute_seed_in_steps(self._task.steps, self._context.metadata.get("seed") or {}) diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index e15f4c1b5..f5907ac9e 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -307,8 +307,8 @@ def _items_artifact(uid: str) -> EnvironmentArtifact: @pytest.mark.asyncio async def test_an_agent_that_works_in_its_sandbox_is_played_through_a_stand_in_deployed_there(local_stack, tmp_path): """A task that loads files into its agent and checks what the agent left there: the agent is never deployed, - a stand-in is, in its place, and the loop's ``bash`` runs in it, as its user and in its folder, with no model key - of its own; the check scores what the commands left, and nothing is left running.""" + a stand-in is, in its place, and the loop's ``bash`` runs in it, as its user and in its folder, with no model + credentials; the check scores what the commands left, and nothing is left running.""" uid = uuid.uuid4().hex[:8] greeting = FileArtifact.put_bytes(id=f"greeting-{uid}", description="stand-in e2e", filename="greeting.txt", content=b"hello", content_type="text/plain") @@ -329,7 +329,7 @@ async def test_an_agent_that_works_in_its_sandbox_is_played_through_a_stand_in_d assert [tool.name for tool in tools.list()] == ["bash"] assert (await tools.call("bash", {"command": "pwd; whoami; cat inputs/greeting.txt"})).text == ( "/app\nagent\nhello") - assert (await tools.call("bash", {"command": "printenv LITELLM_API_KEY"})).text == "unused\n" + assert (await tools.call("bash", {"command": "env | grep -c ^LITELLM_ || true"})).text == "0\n" await tools.call("bash", {"command": "mkdir -p outputs && cp inputs/greeting.txt outputs/answer.txt"}) await turn.end("Copied.") diff --git a/tst/unit/a2a_agent/stand_in_test.py b/tst/unit/a2a_agent/stand_in_test.py index 11eaebb88..41d7c3ca4 100644 --- a/tst/unit/a2a_agent/stand_in_test.py +++ b/tst/unit/a2a_agent/stand_in_test.py @@ -4,17 +4,29 @@ from importlib.metadata import version from agent_env.a2a_agent import stand_in -from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS +from agent_env.a2a_agent.a2a_agent import TAKES_CREDENTIALS, A2AAgent from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent -def test_the_stand_in_is_a_build_context_given_no_model_key_or_store_credentials(local_stores, cli_routing): +def test_the_stand_in_is_a_build_context_that_takes_no_credentials(local_stores, cli_routing): agent = stand_in_agent() assert agent.id == STAND_IN_AGENT_ID assert agent.docker_image_artifact.context_only - assert agent.default_env_vars == {"LITELLM_API_KEY": "unused", "LITELLM_BASE_URL": "http://unused.invalid"} - assert agent.metadata == {SHARE_STORE_CREDENTIALS: False} + assert (agent.default_env_vars, agent.metadata) == ({}, {TAKES_CREDENTIALS: False}) + + +class _NoModelConfig: + def __getattr__(self, name): + raise AssertionError(f"an agent that takes no credentials asked config for {name}") + + +def test_an_agent_that_takes_no_credentials_gets_no_model_endpoint_key_or_headers_whoever_passes_them(): + agent = A2AAgent(id="a", version=1, docker_image_artifact=None, metadata={TAKES_CREDENTIALS: False}) + passed = {"LITELLM_API_KEY": "sk-from-the-step", "LITELLM_BASE_URL": "http://model", "LITELLM_EXTRA_HEADERS": "{}", + "GITHUB_TOKEN": "the task's own"} + + assert agent._with_model_credentials(passed, _NoModelConfig()) == {"GITHUB_TOKEN": "the task's own"} def test_an_unchanged_context_reuses_the_stand_in_and_a_changed_one_writes_it_again( diff --git a/tst/unit/store/aws/test_shared_credentials.py b/tst/unit/store/aws/test_shared_credentials.py index 65efc4602..2b50bcab5 100644 --- a/tst/unit/store/aws/test_shared_credentials.py +++ b/tst/unit/store/aws/test_shared_credentials.py @@ -5,7 +5,7 @@ import pytest from agentenv_protocol import client as protocol_v1 -from agent_env.a2a_agent.a2a_agent import SHARE_STORE_CREDENTIALS, A2AAgent +from agent_env.a2a_agent.a2a_agent import TAKES_CREDENTIALS, A2AAgent from agent_env.config import set_object_store from agent_env.store import routing from agent_env.store.object_store import s3_object_store @@ -85,4 +85,4 @@ def test_a_caller_supplied_key_suppresses_the_whole_shared_set(host_credentials) def test_an_agent_that_opts_out_gets_none_of_the_shared_set(host_credentials): set_object_store(S3ObjectStore(object(), "bucket", region="us-west-2", share_credentials=True)) - assert not any(key.startswith("AWS_") for key in _merged_env(metadata={SHARE_STORE_CREDENTIALS: False})) + assert not any(key.startswith("AWS_") for key in _merged_env(metadata={TAKES_CREDENTIALS: False})) diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index 35035f812..ded6a205d 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -110,8 +110,9 @@ class _ToolSessions: """Stands in for ``ToolSession``: the session on each URL serves the tools ``offered`` names for it, and every session opened, call made and session closed is recorded.""" - def __init__(self, offered): + def __init__(self, offered, raises=None): self.offered = offered + self.raises = raises or {} self.opened = [] self.calls = [] self.closed = [] @@ -129,6 +130,8 @@ async def list_tools(self): async def call_tool(self, name, arguments=None): sessions.calls.append((mcp_url, name, arguments)) + if name in sessions.raises: + raise sessions.raises[name] return ToolResult(content=[{"type": "text", "text": f"{name} on {mcp_url}"}]) async def close(self): @@ -216,8 +219,8 @@ async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch) STOCKED = {**ALICE, "env_ids": ["items", "pantry"]} -def _tool_sessions(monkeypatch, offered): - sessions = _ToolSessions(offered) +def _tool_sessions(monkeypatch, offered, raises=None): + sessions = _ToolSessions(offered, raises) monkeypatch.setattr(episode_run, "ToolSession", sessions) monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) return sessions @@ -572,9 +575,10 @@ async def test_cancelling_the_caller_of_abandon_still_abandons_the_run(teardowns abandoning.cancel() with pytest.raises(asyncio.CancelledError): await abandoning - await _until(lambda: Task.get_instance(episode.id).status == "cancelled") + with pytest.raises(EpisodeError, match="episode abandoned: stop"): + await episode.result() - assert (Task.get_instance(episode.id).status,Task.get_instance(episode.id).error) == ("cancelled", "stop") + assert (Task.get_instance(episode.id).status, Task.get_instance(episode.id).error) == ("cancelled", "stop") assert teardowns.finished == [episode.id] @@ -660,6 +664,40 @@ async def test_a_turns_calls_and_its_reply_are_handed_back_as_its_trajectory(mon assert all(span["start_time"] <= span["end_time"] for span in spans) +@pytest.mark.asyncio +async def test_a_call_that_raises_is_handed_back_in_the_trajectory_too(monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}, + raises={"items_add_item": TimeoutError("the env didn't answer")}) + stand_in = _SandboxedStandIn() + turn = _sandboxed_turn(stand_in) + + async with turn.tools() as tools: + with pytest.raises(TimeoutError, match="the env didn't answer"): + await tools.call("items_add_item", {"name": "crate"}) + await turn.end("Gave up.") + + [(_, reply)] = stand_in.replies + assert reply["tool_call_count"] == 1 + assert json.loads(reply["trajectory"][0]["attributes"]["gen_ai.completion"]) == { + "error": "TimeoutError: the env didn't answer"} + + +@pytest.mark.asyncio +async def test_a_prompt_continuing_a_conversation_gets_its_own_turn(teardowns): + first = {**TURN, "context_id": "chat"} + again = {**TURN, "id": "again", "prompt_id": "again", "prompt": "And another.", "depends_on": ["turn"], + "context_id": "chat"} + episode = await _spec(ITEMS, ALICE, first, again, REPLIES).start() + + turn = await episode.next_turn() + assert (turn.step_id, turn.prompt) == ("turn", "Add a crate.") + await turn.end("Added.") + turn = await episode.next_turn() + assert (turn.step_id, turn.prompt) == ("again", "And another.") + await turn.end("Added another.") + assert await episode.result() == {"replies": 1.0} + + @pytest.mark.asyncio async def test_the_turns_trajectory_is_stored_as_any_agents_is(teardowns, monkeypatch): _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) From ebb3a4442c4c32e43a85a933cf88e829e40dfb2e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Tue, 20 Oct 2026 15:48:11 -0700 Subject: [PATCH 07/39] test(stand-in): count a killed child as stopped once it's a zombie A killed command's background child outlives its parent as a zombie until whatever adopted it reaps it, and signal 0 still reaches a zombie, so checking for its absence right after the timeout raced on Linux. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_a2a_stand_in.py | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/packages/agentenv-protocol/tests/test_a2a_stand_in.py b/packages/agentenv-protocol/tests/test_a2a_stand_in.py index 50b7189d9..4fac062ad 100644 --- a/packages/agentenv-protocol/tests/test_a2a_stand_in.py +++ b/packages/agentenv-protocol/tests/test_a2a_stand_in.py @@ -1,6 +1,6 @@ from __future__ import annotations -import os +import subprocess import time from typing import Any @@ -55,6 +55,18 @@ def _ran(client: TestClient, command: str) -> str: return result["output"] +def _stopped(pid: int) -> None: + """Until ``pid`` is gone, or a zombie: a killed child whose parent has gone stays one until what adopted it reaps it.""" + deadline = time.monotonic() + 10 + while (state := _state(pid)) and not state.startswith("Z"): + assert time.monotonic() < deadline, state + time.sleep(0.02) + + +def _state(pid: int) -> str: + return subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip() + + @pytest.fixture def client(monkeypatch, tmp_path): monkeypatch.setattr(stand_in, "SKILLS_DIR", tmp_path / "skills") @@ -157,8 +169,7 @@ def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(cli output = _ran(client, "sleep 60 & echo $! > child; wait") assert output == "error: the command timed out after 0.5s" - with pytest.raises(ProcessLookupError): - os.kill(int((tmp_path / "child").read_text()), 0) + _stopped(int((tmp_path / "child").read_text())) def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(client, monkeypatch, tmp_path): @@ -167,8 +178,7 @@ def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(c output = _ran(client, "sleep 60 & echo $! > child") assert output == "error: the command timed out after 0.5s" - with pytest.raises(ProcessLookupError): - os.kill(int((tmp_path / "child").read_text()), 0) + _stopped(int((tmp_path / "child").read_text())) def test_output_names_a_command_it_started(client): From 1196dd13ec4a6663ec001275f97adeb10028a065 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 09:20:50 -0700 Subject: [PATCH 08/39] feat(protocol): StandInAgent, an agent with no model that a driver plays through driver/v1 urn:agentenv:driver/v1 hands an agent's turns to a driver outside its sandbox: claim sets the token, next long-polls for the open turns, reply ends one with its text or error and trajectory, and exec and output run a command where the agent works. StandInAgent serves it: each message is an open turn until the driver replies, a cancelled one drops out, skills go to ~/skills//, and commands run as bash -c in its folder, as its user, with a 3600 s timeout that kills their process group. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 25 ++ .../agentenv_protocol/a2a_agent/__init__.py | 22 ++ .../agentenv_protocol/a2a_agent/extensions.py | 88 +++++++ .../agentenv_protocol/a2a_agent/stand_in.py | 215 ++++++++++++++++++ .../tests/test_a2a_stand_in.py | 205 +++++++++++++++++ 5 files changed, 555 insertions(+) create mode 100644 packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py create mode 100644 packages/agentenv-protocol/tests/test_a2a_stand_in.py diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index cc0b804b9..366341a7e 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -697,6 +697,31 @@ The answer is agent-wide rather than per caller, and like every extension route probe carries no access check of its own, so expose an agent only to the control plane that drives it. +`DRIVER_V1` hands an agent's turns to a driver outside its sandbox, and `StandInAgent` +(`agentenv_protocol.a2a_agent.stand_in`) serves it. The stand-in has no model: each message +it gets stays an open turn, its task working, until the driver replies. The driver's first +`POST /ext/driver/claim` (`{"token"}`, at least 22 characters) sets the token every later +call carries in its body. A claim with another token gets 409, and any other call with a +wrong one gets 401. The token is held in memory only, never in the environment a command +inherits. The other operations: + +- `POST /ext/driver/next` (`{"token", "after"}`) waits up to 25 s for the turns' version to + pass `after`, then returns `{"version", "turns": [{"task_id", "context_id", "text", + "system_prompt", "skills"}]}`. A turn its sender cancels drops out of the list. +- `POST /ext/driver/reply` (`{"token", "task_id", "text"}` or `"error"` in place of + `"text"`, plus `"trajectory"` and `"tool_call_count"`) ends a turn. A text is the task's + result, an error fails it with `stand_in.driver_error`, and the trajectory's spans become + the task's trajectory. +- `POST /ext/driver/exec` (`{"token", "command"}`) starts `bash -c` in the stand-in's working + folder, as its user, and returns `{"exec_id"}`. `POST /ext/driver/output` (`{"token", + "exec_id"}`) waits up to 25 s and returns `{"done", "output"}`. The output is stdout then + stderr, or an error line after 3600 s, when the command's process group is killed. + +Skills the stand-in is given are written to `~/skills//`, and a turn names each one's +folder. `StandInAgent(shell=False)` stands in for an agent that works in no sandbox: it keeps +no skills and refuses `exec`. Commands inherit the stand-in's environment, so run it with no +credentials in it. + Runnable, self-contained reference agents live in [`examples/`](https://github.com/scaleapi/agentenv-framework/tree/main/packages/agentenv-protocol/examples): the normal, streaming, and multimodal `run(request)` paths, single- and multi-operation custom extensions, and advanced ASGI-lifespan plus diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index 93141797a..605fe39d1 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -5,6 +5,7 @@ from .extensions import ( AGENT_CONFIG_V1, ATTRIBUTION_PROBE_V1, + DRIVER_V1, INSTALL_V1, MCP_CONFIG_V1, PEER_AGENTS_V1, @@ -19,6 +20,16 @@ ContextEventsTrajectoryRequest, ContextObjectTrajectoryRequest, ContextTrajectoryRequest, + DriverClaimRequest, + DriverExecRequest, + DriverExecResponse, + DriverNextRequest, + DriverNextResponse, + DriverOutputRequest, + DriverOutputResponse, + DriverReplyRequest, + DriverSkill, + DriverTurn, ExtensionActivation, ExtensionConfiguration, ExtensionDefinition, @@ -117,6 +128,7 @@ __all__ = [ "AGENT_CONFIG_V1", "ATTRIBUTION_PROBE_V1", + "DRIVER_V1", "INSTALL_V1", "MCP_CONFIG_V1", "PEER_AGENTS_V1", @@ -139,6 +151,16 @@ "ContextEventsTrajectoryRequest", "ContextObjectTrajectoryRequest", "ContextTrajectoryRequest", + "DriverClaimRequest", + "DriverExecRequest", + "DriverExecResponse", + "DriverNextRequest", + "DriverNextResponse", + "DriverOutputRequest", + "DriverOutputResponse", + "DriverReplyRequest", + "DriverSkill", + "DriverTurn", "DataPart", "DefaultExtensionHandlers", "ExtensionActivation", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index a35c19bbf..82df2bd90 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -287,6 +287,71 @@ class TriggerDecideRequest(ExtensionRequest): env_triggers: dict[str, dict[str, Any]] | None = None +DRIVER_TOKEN_MIN_LENGTH = 22 + + +class DriverClaimRequest(ExtensionRequest): + token: str = Field(min_length=DRIVER_TOKEN_MIN_LENGTH) + + +class DriverNextRequest(ExtensionRequest): + token: str + after: int = Field(default=0, ge=0) + + +class DriverSkill(ExtensionResponse): + name: str + description: str + path: str | None = None + + +class DriverTurn(ExtensionResponse): + task_id: str + context_id: str + text: str + system_prompt: str | None = None + skills: list[DriverSkill] = Field(default_factory=list) + + +class DriverNextResponse(ExtensionResponse): + version: int + turns: list[DriverTurn] + + +class DriverReplyRequest(ExtensionRequest): + token: str + task_id: str + text: str | None = None + error: str | None = None + trajectory: list[dict[str, Any]] = Field(default_factory=list) + tool_call_count: int = Field(default=0, ge=0) + + @model_validator(mode="after") + def _text_or_error(self) -> DriverReplyRequest: + if (self.text is None) == (self.error is None): + raise ValueError("a reply carries exactly one of text and error") + return self + + +class DriverExecRequest(ExtensionRequest): + token: str + command: str + + +class DriverExecResponse(ExtensionResponse): + exec_id: str + + +class DriverOutputRequest(ExtensionRequest): + token: str + exec_id: str + + +class DriverOutputResponse(ExtensionResponse): + done: bool + output: str | None = None + + def _tuple(values: Iterable[str]) -> tuple[str, ...]: return tuple(dict.fromkeys(values)) @@ -1167,6 +1232,28 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: configuration_validator=_configure_install, ) +DRIVER_V1 = ExtensionDefinition( + uri="urn:agentenv:driver/v1", + description="Hand the agent's turns to a driver outside its sandbox, and run its commands where the agent works.", + endpoint="/ext/driver", + core_operations={ + "claim": _op("claim", "POST", "/ext/driver/claim", RUNTIME, request_model=DriverClaimRequest), + "next": _op( + "next", "POST", "/ext/driver/next", RUNTIME, + request_model=DriverNextRequest, response_model=DriverNextResponse, + ), + "reply": _op("reply", "POST", "/ext/driver/reply", RUNTIME, request_model=DriverReplyRequest), + "exec": _op( + "exec", "POST", "/ext/driver/exec", RUNTIME, + request_model=DriverExecRequest, response_model=DriverExecResponse, + ), + "output": _op( + "output", "POST", "/ext/driver/output", RUNTIME, + request_model=DriverOutputRequest, response_model=DriverOutputResponse, + ), + }, +) + STANDARD_EXTENSIONS: Mapping[str, ExtensionDefinition] = MappingProxyType( { extension.uri: extension @@ -1180,6 +1267,7 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: ATTRIBUTION_PROBE_V1, TRIGGERS_V1, INSTALL_V1, + DRIVER_V1, ) } ) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py new file mode 100644 index 000000000..00e915be2 --- /dev/null +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py @@ -0,0 +1,215 @@ +"""The stand-in: an agent with no model, played by a driver outside its sandbox. + +Each message it gets is an open turn until the driver replies to it. Through ``urn:agentenv:driver/v1`` the driver +reads the open turns, replies to them and runs commands where the stand-in works, as it; its first ``claim`` sets the +token every later call carries. A turn its sender cancels closes without a reply. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import hmac +import os +import signal +import uuid +from dataclasses import dataclass, field +from pathlib import Path + +from starlette.exceptions import HTTPException + +from .extensions import ( + DRIVER_V1, + MCP_CONFIG_V1, + SKILL_CONFIG_V1, + TRAJECTORY_V1, + BundleSkillRequest, + DriverClaimRequest, + DriverExecRequest, + DriverExecResponse, + DriverNextRequest, + DriverNextResponse, + DriverOutputRequest, + DriverOutputResponse, + DriverReplyRequest, + DriverSkill, + DriverTurn, + InlineSkillRequest, + extension, +) +from .framework import AgentEnvAgent, AgentIdentity, a2a_agent +from .tasks.v1 import AgentConfig, TaskRequest, TaskResult, TextPart, Usage +from ..transfers import download + +LONG_POLL_SECONDS = 25.0 +COMMAND_TIMEOUT_SECONDS = 3600 +SKILLS_DIR = Path.home() / "skills" +TRAJECTORY_FORMAT = "agentenv-stand-in/v1" + + +class StandInConfig(AgentConfig): + system_prompt: str | None = None + + +@dataclass +class _Turn: + task_id: str + context_id: str + text: str + system_prompt: str | None + skills: list[DriverSkill] + reply: asyncio.Future[DriverReplyRequest] = field( + default_factory=lambda: asyncio.get_running_loop().create_future() + ) + + +IDENTITY = AgentIdentity( + name="agentenv-stand-in", + description="Stands in for an agent played from outside AgentEnv: a driver answers its turns.", + version="1.0.0", +) +EXTENSIONS = (MCP_CONFIG_V1, TRAJECTORY_V1, DRIVER_V1) + + +@a2a_agent(identity=IDENTITY, config=StandInConfig, extensions=EXTENSIONS) +class StandInAgent(AgentEnvAgent): + """``shell=False`` is a stand-in for an agent that works in no sandbox: it keeps no skills and runs no commands. + ``long_poll_seconds`` bounds how long ``next`` and ``output`` wait for something to report.""" + + def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): + self._shell = shell + self._long_poll_seconds = long_poll_seconds + self._token: str | None = None + self._turns: dict[str, _Turn] = {} + self._version = 0 + self._changed: asyncio.Future[None] | None = None + self._commands: dict[str, asyncio.Task[str]] = {} + + async def run(self, request: TaskRequest[StandInConfig]) -> TaskResult: + turn = _Turn( + task_id=request.task_id, + context_id=request.context_id, + text="\n".join(part.text for part in request.parts if isinstance(part, TextPart)), + system_prompt=request.config.system_prompt, + skills=[DriverSkill(name=skill["name"], description=skill.get("description") or "", + path=str(SKILLS_DIR / skill["name"]) if self._shell else None) + for skill in request.skills], + ) + self._turns[turn.task_id] = turn + self._bump() + try: + reply = await turn.reply + finally: + del self._turns[turn.task_id] + self._bump() + request.trajectory.set_format(TRAJECTORY_FORMAT) + for span in reply.trajectory: + request.trajectory.append(span) + if reply.error is not None: + return TaskResult.failure("stand_in.driver_error", reply.error) + return TaskResult.text(reply.text, usage=Usage(tool_call_count=reply.tool_call_count)) + + @extension(DRIVER_V1.claim) + async def claim(self, request: DriverClaimRequest) -> None: + if self._token is None: + self._token = request.token + elif not hmac.compare_digest(self._token, request.token): + raise HTTPException(status_code=409, detail="this stand-in is already claimed") + + @extension(DRIVER_V1.next) + async def next_turns(self, request: DriverNextRequest) -> DriverNextResponse: + self._check(request.token) + if self._version <= request.after: + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(self._next_change()), self._long_poll_seconds) + return DriverNextResponse(version=self._version, turns=[ + DriverTurn(task_id=turn.task_id, context_id=turn.context_id, text=turn.text, + system_prompt=turn.system_prompt, skills=turn.skills) + for turn in self._turns.values() if not turn.reply.done() + ]) + + @extension(DRIVER_V1.reply) + async def reply(self, request: DriverReplyRequest) -> None: + self._check(request.token) + turn = self._turns.get(request.task_id) + if turn is None or turn.reply.done(): + raise HTTPException(status_code=404, detail=f"no open turn {request.task_id!r}") + turn.reply.set_result(request) + + @extension(DRIVER_V1.exec) + async def exec_command(self, request: DriverExecRequest) -> DriverExecResponse: + self._check(request.token) + if not self._shell: + raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it runs no commands") + command_id = uuid.uuid4().hex + self._commands[command_id] = asyncio.create_task(_run(request.command)) + return DriverExecResponse(exec_id=command_id) + + @extension(DRIVER_V1.output) + async def output(self, request: DriverOutputRequest) -> DriverOutputResponse: + self._check(request.token) + if (command := self._commands.get(request.exec_id)) is None: + raise HTTPException(status_code=404, detail=f"no command {request.exec_id!r}") + await asyncio.wait((command,), timeout=self._long_poll_seconds) + if not command.done(): + return DriverOutputResponse(done=False) + self._commands.pop(request.exec_id, None) + return DriverOutputResponse(done=True, output=command.result()) + + @extension(SKILL_CONFIG_V1.add.inline) + async def add_inline_skill(self, request: InlineSkillRequest) -> dict[str, str]: + folder = self._skill_folder(request.name) + folder.mkdir(parents=True) + (folder / "SKILL.md").write_text(request.skill_md) + return {"name": request.name} + + @extension(SKILL_CONFIG_V1.add.bundle) + async def add_bundle_skill(self, request: BundleSkillRequest) -> dict[str, str]: + folder = self._skill_folder(request.name) + for item in request.skill_bundle.files: + await download(item.object, folder / item.path) + return {"name": request.name} + + def _check(self, token: str) -> None: + if self._token is None or not hmac.compare_digest(self._token, token): + raise HTTPException(status_code=401, detail="the driver's token is missing or wrong") + + def _skill_folder(self, name: str) -> Path: + if not self._shell: + raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it keeps no skills") + return SKILLS_DIR / name + + def _next_change(self) -> asyncio.Future[None]: + if self._changed is None: + self._changed = asyncio.get_running_loop().create_future() + return self._changed + + def _bump(self) -> None: + self._version += 1 + if self._changed is not None: + self._changed.set_result(None) + self._changed = None + + +async def _run(command: str) -> str: + """``command``'s stdout then its stderr, as verifiers' bash tool returns them. It runs in a process group of its + own, stopped with whatever it started if its output isn't done within ``COMMAND_TIMEOUT_SECONDS``, or if the call + is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" + process = await asyncio.create_subprocess_exec( + "bash", "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True, + ) + try: + stdout, stderr = await asyncio.wait_for(process.communicate(), COMMAND_TIMEOUT_SECONDS) + except BaseException as e: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + if isinstance(e, asyncio.TimeoutError): + return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" + raise + return (stdout + stderr).decode(errors="replace") + + +if __name__ == "__main__": + StandInAgent().serve() diff --git a/packages/agentenv-protocol/tests/test_a2a_stand_in.py b/packages/agentenv-protocol/tests/test_a2a_stand_in.py new file mode 100644 index 000000000..4fac062ad --- /dev/null +++ b/packages/agentenv-protocol/tests/test_a2a_stand_in.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import subprocess +import time +from typing import Any + +import pytest +from starlette.testclient import TestClient + +from agentenv_protocol.a2a_agent import stand_in +from agentenv_protocol.a2a_agent.stand_in import StandInAgent + +TOKEN = "t" * 32 + + +def _send(client: TestClient, text: str, *, context_id: str = "context-1") -> str: + response = client.post("/a2a", json={ + "jsonrpc": "2.0", "id": "1", "method": "message/send", + "params": { + "message": {"kind": "message", "messageId": f"message-{text}", "role": "user", "contextId": context_id, + "parts": [{"kind": "text", "text": text}]}, + "configuration": {"blocking": False}, + }, + }) + return response.json()["result"]["id"] + + +def _task(client: TestClient, task_id: str) -> dict[str, Any]: + return client.post("/a2a", json={"jsonrpc": "2.0", "id": "1", "method": "tasks/get", + "params": {"id": task_id}}).json()["result"] + + +def _settled(client: TestClient, task_id: str) -> dict[str, Any]: + deadline = time.monotonic() + 10 + while (task := _task(client, task_id))["status"]["state"] not in ("completed", "failed", "canceled"): + assert time.monotonic() < deadline, task + time.sleep(0.02) + return task + + +def _driver(client: TestClient, operation: str, **body: Any): + return client.post(f"/ext/driver/{operation}", json={"token": TOKEN, **body}) + + +def _turns(client: TestClient, after: int = 0) -> dict[str, Any]: + response = _driver(client, "next", after=after) + assert response.status_code == 200, response.text + return response.json() + + +def _ran(client: TestClient, command: str) -> str: + exec_id = _driver(client, "exec", command=command).json()["exec_id"] + while not (result := _driver(client, "output", exec_id=exec_id).json())["done"]: + pass + return result["output"] + + +def _stopped(pid: int) -> None: + """Until ``pid`` is gone, or a zombie: a killed child whose parent has gone stays one until what adopted it reaps it.""" + deadline = time.monotonic() + 10 + while (state := _state(pid)) and not state.startswith("Z"): + assert time.monotonic() < deadline, state + time.sleep(0.02) + + +def _state(pid: int) -> str: + return subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip() + + +@pytest.fixture +def client(monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "SKILLS_DIR", tmp_path / "skills") + monkeypatch.chdir(tmp_path) + with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + yield client + + +def test_the_card_advertises_the_driver_and_takes_a_system_prompt(): + card = StandInAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json") + extensions = {extension["uri"]: extension for extension in card["capabilities"]["extensions"]} + + assert sorted(extensions["urn:agentenv:driver/v1"]["params"]["methods"]) == [ + "claim", "exec", "next", "output", "reply"] + assert "system_prompt" in extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"][ + "supported"] + + +def test_the_first_claim_sets_the_token_that_every_driver_call_needs(): + with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: + assert _driver(client, "next").status_code == 401 + assert client.post("/ext/driver/claim", json={"token": "short"}).status_code == 400 + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/driver/claim", json={"token": "u" * 32}).status_code == 409 + assert client.post("/ext/driver/next", json={"token": "u" * 32}).status_code == 401 + assert _driver(client, "next").status_code == 200 + + +def test_a_message_is_an_open_turn_until_the_driver_replies(client): + client.post("/ext/agent-config", json={"system_prompt": "be brief"}) + task_id = _send(client, "fix the build") + + seen = _turns(client) + assert seen["turns"] == [{"task_id": task_id, "context_id": "context-1", "text": "fix the build", + "system_prompt": "be brief", "skills": []}] + span = {"name": "bash", "attributes": {"gen_ai.operation.name": "execute_tool"}} + assert _driver(client, "reply", task_id=task_id, text="fixed", trajectory=[span], tool_call_count=1).status_code \ + == 200 + + task = _settled(client, task_id) + assert task["status"]["state"] == "completed" + parts = task["status"]["message"]["parts"] + assert parts[0] == {"kind": "text", "text": "fixed"} + assert parts[1]["data"]["usage"]["tool_call_count"] == 1 + assert client.post("/ext/trajectory", json={"task_id": task_id}).json()["trajectory"] == [span] + assert _turns(client, after=seen["version"])["turns"] == [] + + +def test_an_error_reply_fails_the_turn(client): + task_id = _send(client, "do it") + + _driver(client, "reply", task_id=task_id, error="the driver gave up") + + task = _settled(client, task_id) + assert task["status"]["state"] == "failed" + assert "the driver gave up" in task["status"]["message"]["parts"][0]["text"] + + +def test_a_reply_names_exactly_one_open_turn(client): + task_id = _send(client, "do it") + + assert _driver(client, "reply", task_id=task_id).status_code == 400 + assert _driver(client, "reply", task_id="nope", text="done").status_code == 404 + _driver(client, "reply", task_id=task_id, text="done") + _settled(client, task_id) + assert _driver(client, "reply", task_id=task_id, text="again").status_code == 404 + + +def test_a_turn_its_sender_cancels_is_withdrawn(client): + task_id = _send(client, "do it") + assert [turn["task_id"] for turn in _turns(client)["turns"]] == [task_id] + + client.post("/a2a", json={"jsonrpc": "2.0", "id": "1", "method": "tasks/cancel", "params": {"id": task_id}}) + + assert _settled(client, task_id)["status"]["state"] == "canceled" + assert _turns(client)["turns"] == [] + + +def test_next_waits_for_a_change_after_the_version_it_was_given(client): + version = _turns(client)["version"] + + started = time.monotonic() + unchanged = _turns(client, after=version) + + assert unchanged["version"] == version + assert time.monotonic() - started >= 0.2 + + +def test_exec_returns_stdout_then_stderr_from_the_folder_it_works_in(client, tmp_path): + (tmp_path / "inputs").mkdir() + + assert _ran(client, "ls; echo out; echo err >&2; exit 3") == "inputs\nout\nerr\n" + + +def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(client, monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + + output = _ran(client, "sleep 60 & echo $! > child; wait") + + assert output == "error: the command timed out after 0.5s" + _stopped(int((tmp_path / "child").read_text())) + + +def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(client, monkeypatch, tmp_path): + monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + + output = _ran(client, "sleep 60 & echo $! > child") + + assert output == "error: the command timed out after 0.5s" + _stopped(int((tmp_path / "child").read_text())) + + +def test_output_names_a_command_it_started(client): + assert _driver(client, "output", exec_id="nope").status_code == 404 + + +def test_a_skill_is_written_where_its_turns_say_it_is(client, tmp_path): + response = client.post("/ext/skill-config", json={"name": "deploy", "description": "How to deploy.", + "skill_md": "# Deploy\n"}) + assert response.status_code == 200, response.text + _send(client, "go") + + skill = _turns(client)["turns"][0]["skills"] + assert skill == [{"name": "deploy", "description": "How to deploy.", "path": str(tmp_path / "skills" / "deploy")}] + assert (tmp_path / "skills" / "deploy" / "SKILL.md").read_text() == "# Deploy\n" + + +def test_a_stand_in_without_a_shell_keeps_no_skills_and_runs_no_commands(): + with TestClient(StandInAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: + client.post("/ext/driver/claim", json={"token": TOKEN}) + + assert _driver(client, "exec", command="true").status_code == 409 + assert client.post("/ext/skill-config", json={"name": "x", "description": "d", "skill_md": "#"}).status_code \ + == 409 From 24c6f76938f64d4392f35dd90f95f0dd3ea719b2 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 09:32:54 -0700 Subject: [PATCH 09/39] refactor(agents): the stand-in is an agent of the agents bundle, written by the episode as a run would The stand-in moves from a2a_agent/stand_in/, which staged the installed protocol package into its own build context, to agents/stand-in/, an agent folder like claude-code's and pi's: a Dockerfile, an a2a_server.py declaring its install/v1, an install.sh for a VM's host, the protocol from PyPI in requirements.txt, which a release pins, and an agent.toml setting takes_credentials = false, which [metadata] now takes. Its id is @local/agentenv-framework/agents/stand-in. An episode writes it with write_installed, which plans one installed bundle's entity with what it references (plan_entities, the checks plan_bundle makes, with nothing run) and materializes it, so the ledger reuses it until its folder changes. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 3 +- src/agent_env/a2a_agent/a2a_agent.py | 25 +++++--- src/agent_env/a2a_agent/stand_in/Dockerfile | 13 ---- src/agent_env/a2a_agent/stand_in/__init__.py | 63 ------------------- .../a2a_agent/stand_in/a2a_server.py | 37 ----------- src/agent_env/agents/README.md | 7 +++ .../agents/agents/stand-in/Dockerfile | 17 +++++ .../agents/agents/stand-in/a2a_server.py | 31 +++++++++ .../agents/agents/stand-in/agent.toml | 3 + .../agents/agents/stand-in/install.sh | 27 ++++++++ .../agents/agents/stand-in/requirements.txt | 3 + src/agent_env/bundle/installed.py | 22 ++++++- src/agent_env/bundle/plan.py | 26 +++++++- src/agent_env/task/__init__.py | 2 +- src/agent_env/task/episode/rules.py | 4 +- src/agent_env/task/episode/steps.py | 7 ++- tst/unit/a2a_agent/stand_in_test.py | 52 --------------- tst/unit/agents/stand_in/__init__.py | 0 tst/unit/agents/stand_in/agent.py | 12 ++++ .../agents/stand_in/test_stand_in_agent.py | 56 +++++++++++++++++ tst/unit/bundle/installed_entities_test.py | 48 +++++++++++++- tst/unit/bundle/materialize_test.py | 14 +++-- tst/unit/bundle/to_toml_round_trip_test.py | 13 ++-- tst/unit/task_step/install_agent_test.py | 2 +- 24 files changed, 294 insertions(+), 193 deletions(-) delete mode 100644 src/agent_env/a2a_agent/stand_in/Dockerfile delete mode 100644 src/agent_env/a2a_agent/stand_in/__init__.py delete mode 100644 src/agent_env/a2a_agent/stand_in/a2a_server.py create mode 100644 src/agent_env/agents/agents/stand-in/Dockerfile create mode 100644 src/agent_env/agents/agents/stand-in/a2a_server.py create mode 100644 src/agent_env/agents/agents/stand-in/agent.toml create mode 100644 src/agent_env/agents/agents/stand-in/install.sh create mode 100644 src/agent_env/agents/agents/stand-in/requirements.txt delete mode 100644 tst/unit/a2a_agent/stand_in_test.py create mode 100644 tst/unit/agents/stand_in/__init__.py create mode 100644 tst/unit/agents/stand_in/agent.py create mode 100644 tst/unit/agents/stand_in/test_stand_in_agent.py diff --git a/AGENTS.md b/AGENTS.md index ba2cf4314..6f3f57d41 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -204,7 +204,8 @@ orchestrators call it the same way. What callers rely on: `context.heartbeat()`, a no-op without one. The key is in `_REDACTED_KEYS`, so it is never persisted. - An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task, with namespace routing on, in which each of its `trainer_agents` is played through a stand-in agent - (`agentenv_protocol.a2a_agent.stand_in`, written as `@local/agentenv-stand-in` from its build context). + (the `agents` bundle's `stand-in`, `@local/agentenv-framework/agents/stand-in`, which the episode writes + from its folder as a bundle run would). The agent's `deploy_agent` or `install_agent`, wrapped as `episode_agent`, deploys the stand-in in the agent's place when the task uses the agent's sandbox and serves it in this process when it doesn't; the agent's `prompt_agent` steps message it as they would the agent, and the caller plays each diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index b65b34107..b5a748cc9 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -120,6 +120,7 @@ class A2AAgent: # What an agent.toml's [metadata] may set: the types each key takes, and how a problem names them. toml_metadata: ClassVar[dict[str, tuple[tuple[type, ...], str]]] = { "default_model": ((str,), "a string"), "min_disk_size_gb": ((int, float), "a number"), + TAKES_CREDENTIALS: ((bool,), "true or false"), } EXT_MCP_CONFIG: ClassVar[str] = "urn:agentenv:mcp-config/v1" @@ -370,10 +371,10 @@ def from_toml(cls, data: dict[str, Any], ctx: AuthoringContext) -> A2AAgent: def to_toml(self, out: ExportContext) -> dict[str, Any]: """The agent.toml of a folder this agent is built from, which ``from_toml`` reads back as it.""" - metadata = {key: value for key, value in self.metadata.items() if key in self.toml_metadata - and isinstance(value, self.toml_metadata[key][0]) and not isinstance(value, bool)} + metadata = {key: value for key, value in self.metadata.items() + if key in self.toml_metadata and self._fits_toml_metadata(key, value)} if dropped := sorted(set(self.metadata) - set(metadata)): - out.left_out("metadata", f"an agent.toml's [metadata] takes only {' and '.join(self.toml_metadata)}, not " + out.left_out("metadata", f"an agent.toml's [metadata] takes only {self._toml_metadata_keys()}, not " f"{', '.join(dropped)}") toml: dict[str, Any] = {"image": out.image(self.docker_image_artifact)} if self.default_env_vars: @@ -390,16 +391,26 @@ def accept_toml(cls, data: dict[str, Any], ctx: AuthoringContext) -> dict[str, A for name, value in fields.get("metadata", {}).items(): if name not in cls.toml_metadata: values.append(f"unknown [metadata] key {name!r}; an agent's [metadata] takes " - f"{' and '.join(cls.toml_metadata)}") + f"{cls._toml_metadata_keys()}") continue - kinds, described = cls.toml_metadata[name] - if isinstance(value, bool) or not isinstance(value, kinds): - values.append(f"metadata.{name} must be {described}, not {value!r}") + if not cls._fits_toml_metadata(name, value): + values.append(f"metadata.{name} must be {cls.toml_metadata[name][1]}, not {value!r}") problems += [ctx.config_problem(problem) for problem in values] if problems: ctx.refuse(problems) return fields + @classmethod + def _toml_metadata_keys(cls) -> str: + *rest, last = cls.toml_metadata + return f"{', '.join(rest)} and {last}" + + @classmethod + def _fits_toml_metadata(cls, name: str, value: Any) -> bool: + """A bool is an int to ``isinstance``, so it fits only a key that takes one.""" + kinds = cls.toml_metadata[name][0] + return isinstance(value, kinds) and (bool in kinds or not isinstance(value, bool)) + @classmethod def get(cls, id: str, version: Optional[int] = None) -> A2AAgent: from .store import get_a2a_agent_store diff --git a/src/agent_env/a2a_agent/stand_in/Dockerfile b/src/agent_env/a2a_agent/stand_in/Dockerfile deleted file mode 100644 index 67edcb2ac..000000000 --- a/src/agent_env/a2a_agent/stand_in/Dockerfile +++ /dev/null @@ -1,13 +0,0 @@ -FROM mirror.gcr.io/library/python:3.12-slim@sha256:423ed6ab25b1921a477529254bfeeabf5855151dc2c3141699a1bfc852199fbf -RUN apt-get update \ - && apt-get install -y --no-install-recommends bash ca-certificates coreutils curl findutils git procps \ - && rm -rf /var/lib/apt/lists/* -COPY requirements.txt /opt/agentenv-stand-in/ -RUN pip install --no-cache-dir -r /opt/agentenv-stand-in/requirements.txt && useradd --create-home agent -COPY agentenv_protocol/ /opt/agentenv-stand-in/agentenv_protocol/ -ENV PYTHONPATH=/opt/agentenv-stand-in -WORKDIR /app -COPY a2a_server.py . -RUN chown agent:agent /app -USER agent -CMD ["python", "/app/a2a_server.py"] diff --git a/src/agent_env/a2a_agent/stand_in/__init__.py b/src/agent_env/a2a_agent/stand_in/__init__.py deleted file mode 100644 index 40f014b9e..000000000 --- a/src/agent_env/a2a_agent/stand_in/__init__.py +++ /dev/null @@ -1,63 +0,0 @@ -"""The stand-in agent (``agentenv_protocol.a2a_agent.stand_in``) as the Episode API deploys it, in the place of an agent -a loop plays: ``@local/agentenv-stand-in``, an image registered by its build context alone and built where each deploy -runs. Its context is this folder, the installed protocol package, and that package's requirements for serving an -agent, pinned to the versions installed here.""" - -from __future__ import annotations - -import shutil -import tempfile -import threading -from importlib.metadata import requires, version -from pathlib import Path - -import agentenv_protocol -from packaging.requirements import Requirement - -from agent_env.a2a_agent.a2a_agent import TAKES_CREDENTIALS, A2AAgent -from agent_env.artifact.artifacts.docker_image import DockerImageArtifact -from agent_env.store.base import NotFoundError - -STAND_IN_AGENT_ID = "@local/agentenv-stand-in" -_IMAGE_ID = f"{STAND_IN_AGENT_ID}/image" -_PROTOCOL_DISTRIBUTION = "agentenv-framework-protocol" -_HERE = Path(__file__).parent -_writing = threading.Lock() - - -def stand_in_agent() -> A2AAgent: - """The stand-in, written to the store first when it isn't there or its build context has changed.""" - with _writing, tempfile.TemporaryDirectory() as staged: - context = _stage(Path(staged)) - dockerfile = str(context / "Dockerfile") - digest, _ = DockerImageArtifact.context_identity(_IMAGE_ID, str(context), dockerfile, platform=None) - try: - agent = A2AAgent.get(STAND_IN_AGENT_ID) - if agent.docker_image_artifact.source_digest == digest: - return agent - except NotFoundError: - pass - image = DockerImageArtifact.put_context( - _IMAGE_ID, description="The stand-in agent: a driver outside its sandbox answers its turns.", - context_path=str(context), dockerfile_path=dockerfile, platform=None, - ) - return A2AAgent.put(id=STAND_IN_AGENT_ID, docker_image_artifact=image, metadata={TAKES_CREDENTIALS: False}) - - -def _stage(context: Path) -> Path: - for name in ("Dockerfile", "a2a_server.py"): - shutil.copy2(_HERE / name, context / name) - shutil.copytree(Path(agentenv_protocol.__file__).parent, context / "agentenv_protocol", - ignore=shutil.ignore_patterns("__pycache__", "*.pyc")) - (context / "requirements.txt").write_text("".join(f"{line}\n" for line in _requirements())) - return context - - -def _requirements() -> list[str]: - pinned = [] - for spec in requires(_PROTOCOL_DISTRIBUTION) or []: - requirement = Requirement(spec) - if requirement.marker is None or requirement.marker.evaluate({"extra": "agent"}): - extras = f"[{','.join(sorted(requirement.extras))}]" if requirement.extras else "" - pinned.append(f"{requirement.name}{extras}=={version(requirement.name)}") - return sorted(pinned) diff --git a/src/agent_env/a2a_agent/stand_in/a2a_server.py b/src/agent_env/a2a_agent/stand_in/a2a_server.py deleted file mode 100644 index dadb0b388..000000000 --- a/src/agent_env/a2a_agent/stand_in/a2a_server.py +++ /dev/null @@ -1,37 +0,0 @@ -"""The stand-in agent as AgentEnv deploys it: served from ``/app`` in its image, or installed on a VM host as -``install_agent`` installs it, in the place of the agent it replaces (its ``{workdir}``, as its ``{user}``). -``AGENT_CARD`` is what ``install_agent`` reads; importing this module serves nothing. ``install_agent`` formats each -command with ``str.format``, so a literal brace is doubled.""" - -from agentenv_protocol.a2a_agent import INSTALL_V1, a2a_agent, enable -from agentenv_protocol.a2a_agent.stand_in import EXTENSIONS, IDENTITY, StandInAgent, StandInConfig - -HOME = "/opt/agentenv-stand-in-{agent_name}" -INSTALL_COMMANDS_HOST = [ - "bash -c 'set -e; python3 -c \"import ensurepip, venv\" 2>/dev/null || {{ export DEBIAN_FRONTEND=noninteractive; " - "apt-get update -qq; apt-get install -y -qq --no-install-recommends python3 python3-venv; }}'", - f"rm -rf {HOME} && mkdir -p {HOME} && tar -xzf {{agent_ctx_tar}} -C {HOME}", - f"python3 -m venv {HOME}/venv && {HOME}/venv/bin/pip install -q -r {HOME}/requirements.txt", - "id -u {user} >/dev/null 2>&1 || useradd -m {user}; mkdir -p {workdir} && chown -R {user}: {workdir}", - f"nohup runuser -u {{user}} -- env -C {{workdir}} A2A_PORT={{a2a_port}} PYTHONPATH={HOME} {HOME}/venv/bin/python " - f"{HOME}/a2a_server.py > /tmp/agentenv-stand-in-{{agent_name}}.log 2>&1 < /dev/null &", -] - - -@a2a_agent( - identity=IDENTITY, - config=StandInConfig, - extensions=( - *EXTENSIONS, - enable(INSTALL_V1, install_commands_host=INSTALL_COMMANDS_HOST, - required_params_host=["agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], a2a_port=8000), - ), -) -class InstalledStandIn(StandInAgent): - pass - - -AGENT_CARD = InstalledStandIn().create_app().state.agentenv_a2a.card - -if __name__ == "__main__": - InstalledStandIn().serve() diff --git a/src/agent_env/agents/README.md b/src/agent_env/agents/README.md index dec562ce9..401345ac6 100644 --- a/src/agent_env/agents/README.md +++ b/src/agent_env/agents/README.md @@ -26,5 +26,12 @@ the `[model]` endpoint as an OpenAI-compatible one and takes its model from the `[model] default` (or the `agent` role) set. It declares `install/v1` too, for Debian or Ubuntu targets. Its folder's `README.md` covers the extensions it implements. +`agents/stand-in/` is the Episode API's, not an agent to pick for a task. It's the protocol's `StandInAgent`, which +has no model: an episode deploys it in the place of the agent its caller plays, answers its turns through +`urn:agentenv:driver/v1`, and runs the caller's `bash` in it. Its image works in `/app` as user `agent`. On a `vm` +sandbox's host, `install.sh` installs it where the agent it replaces says it works, with the host's own Python in a +venv. Its `agent.toml` sets `takes_credentials = false`, so it's given no model key or store credentials. It has no +smoke task: with nobody answering its turns, a prompt only waits out its deadline. + `agent-env run agents` runs every agent's smoke task, building each image; `--task claude-code-smoke` runs Claude Code's alone, `--task pi-smoke` pi's. diff --git a/src/agent_env/agents/agents/stand-in/Dockerfile b/src/agent_env/agents/agents/stand-in/Dockerfile new file mode 100644 index 000000000..4e088a623 --- /dev/null +++ b/src/agent_env/agents/agents/stand-in/Dockerfile @@ -0,0 +1,17 @@ +FROM mirror.gcr.io/library/python:3.12-slim@sha256:423ed6ab25b1921a477529254bfeeabf5855151dc2c3141699a1bfc852199fbf +# The tools the steps that act on an agent's sandbox run: collect_artifacts, verify_sandbox, run_code, load_artifact. +RUN apt-get update \ + && apt-get install -y --no-install-recommends bash ca-certificates coreutils curl findutils git procps \ + && rm -rf /var/lib/apt/lists/* +WORKDIR /opt/agent +COPY requirements.txt ./ +RUN pip install --no-cache-dir -r requirements.txt +# install.sh is install/v1's: in the image, so it stays in its build context. +COPY a2a_server.py install.sh ./ +RUN useradd --create-home agent && mkdir -p /app && chown agent:agent /app +ENV A2A_PORT=8000 +USER agent +# The commands the driver runs work where the stand-in does: in /app, as agent. +WORKDIR /app +EXPOSE 8000 +CMD ["python", "/opt/agent/a2a_server.py"] diff --git a/src/agent_env/agents/agents/stand-in/a2a_server.py b/src/agent_env/agents/agents/stand-in/a2a_server.py new file mode 100644 index 000000000..3b7c29706 --- /dev/null +++ b/src/agent_env/agents/agents/stand-in/a2a_server.py @@ -0,0 +1,31 @@ +"""Serves the stand-in agent. ``AGENT_CARD`` is what agent-env's install/v1 reads. ``install_agent`` formats each +command with ``str.format``.""" + +from agentenv_protocol.a2a_agent import INSTALL_V1, a2a_agent, enable +from agentenv_protocol.a2a_agent.stand_in import EXTENSIONS, IDENTITY, StandInAgent, StandInConfig + +_UNPACK = "mkdir -p {work_dir}/agent && tar -xzf {agent_ctx_tar} -C {work_dir}/agent" +# In the place of the agent it replaces: its {workdir}, as its {user}. +_INSTALL_COMMANDS_HOST = [_UNPACK, "sh {work_dir}/agent/install.sh {agent_name} {a2a_port} {workdir} {user}"] + + +@a2a_agent( + identity=IDENTITY, + config=StandInConfig, + extensions=( + *EXTENSIONS, + enable(INSTALL_V1, install_commands_host=_INSTALL_COMMANDS_HOST, + required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], + a2a_port=8000), + ), +) +class InstalledStandIn(StandInAgent): + pass + + +agent = InstalledStandIn() +app = agent.create_app() +AGENT_CARD = app.state.agentenv_a2a.card + +if __name__ == "__main__": + agent.serve() diff --git a/src/agent_env/agents/agents/stand-in/agent.toml b/src/agent_env/agents/agents/stand-in/agent.toml new file mode 100644 index 000000000..fce9bb754 --- /dev/null +++ b/src/agent_env/agents/agents/stand-in/agent.toml @@ -0,0 +1,3 @@ +[metadata] +# The commands a driver runs in its sandbox can read its environment, so it's given no model credentials. +takes_credentials = false diff --git a/src/agent_env/agents/agents/stand-in/install.sh b/src/agent_env/agents/agents/stand-in/install.sh new file mode 100644 index 000000000..7e152dde4 --- /dev/null +++ b/src/agent_env/agents/agents/stand-in/install.sh @@ -0,0 +1,27 @@ +#!/bin/sh +# Installs the stand-in onto this Linux host from the folder this script is in, and starts it on PORT, as USER in +# WORKDIR, the place of the agent it replaces: +# +# install.sh NAME PORT WORKDIR USER +# +# It runs in a venv of the host's own Python, which apt installs when the host has none. +set -eu +name=$1 port=$2 workdir=$3 user=$4 +here=$(cd "$(dirname "$0")" && pwd) +root=/opt/agentenv-stand-in-$name + +if ! python3 -c "import ensurepip, venv" 2>/dev/null; then + export DEBIAN_FRONTEND=noninteractive + apt-get update -qq + apt-get install -y -qq --no-install-recommends python3 python3-venv +fi +rm -rf "$root" +mkdir -p "$root" +cp "$here/a2a_server.py" "$here/requirements.txt" "$root/" +python3 -m venv "$root/venv" +"$root/venv/bin/pip" install -q -r "$root/requirements.txt" +id -u "$user" >/dev/null 2>&1 || useradd -m "$user" +mkdir -p "$workdir" +chown -R "$user": "$workdir" +nohup runuser -u "$user" -- env -C "$workdir" A2A_PORT="$port" "$root/venv/bin/python" "$root/a2a_server.py" \ + >"/tmp/agentenv-stand-in-$name.log" 2>&1 =0.1.356 diff --git a/src/agent_env/bundle/installed.py b/src/agent_env/bundle/installed.py index 72268a68a..a27a15bd1 100644 --- a/src/agent_env/bundle/installed.py +++ b/src/agent_env/bundle/installed.py @@ -24,8 +24,9 @@ from agent_env.store.ids import LOCAL_PREFIX from ._fs import os_reason -from .parse import Bundle, BundleError, parse_bundle -from .plan import check_bundle +from .materialize import materialize +from .parse import Bundle, BundleError, BundleKind, parse_bundle +from .plan import check_bundle, plan_entities from .resolve import resolve_bundle BUNDLES = "agent_env.bundles" @@ -106,6 +107,23 @@ def checked(bundle: InstalledBundle, *, load: bool = True) -> Bundle: return parsed +def write_installed(entity_id: str) -> int: + """Write ``entity_id``, an env, agent or artifact an installed bundle defines, with what it references, as a run of + a bundle naming it does, and return its version: the one already written when nothing it's made from changed. + Needs namespace routing. Raises BundleError.""" + bundles = installed_bundles() + bundle = next((bundle for bundle in bundles if entity_id.startswith(f"{bundle.id_root}/")), None) + if bundle is None: + raise BundleError([f"{entity_id!r} isn't an entity of an installed bundle"]) + if bundle.problem: + raise BundleError([f"{bundle.qualified}: {bundle.problem}"]) + plan = plan_entities(resolve_bundle(parse_bundle(bundle.root, id_root=bundle.id_root)), [entity_id], + installed=lambda: bundles) + [write] = [write for write in plan.writes + if write.id == entity_id and write.kind not in (BundleKind.TASK, BundleKind.EVAL)] + return materialize(plan).version_of(write.kind.store, entity_id) + + def _installed(name: str, plugin: Plugin, ep: EntryPoint, *, repeated: bool) -> InstalledBundle: package, version = identity(ep) diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index 045d3363b..2bae5ba68 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -117,6 +117,13 @@ def plan_bundle(resolved: ResolvedBundle, *, tasks: Sequence[str] = (), evals: S return _Planner(resolved, installed).plan(tasks, evals) +def plan_entities(resolved: ResolvedBundle, ids: Sequence[str], *, installed: Installed | None = None) -> Plan: + """Plan writing the envs, agents and artifacts ``ids`` names in this bundle, with what they reference, and + running nothing: what a run writes for one a task of another bundle names. Raises BundleError listing every + problem found.""" + return _Planner(resolved, installed).plan_entities(ids) + + def check_bundle(resolved: ResolvedBundle, *, installed: Installed | None = None) -> None: """The checks ``plan_bundle`` makes without reading a store, over every task and eval in the bundle. Raises BundleError listing every problem found.""" @@ -284,7 +291,16 @@ def plan(self, tasks: Sequence[str], evals: Sequence[str]) -> Plan: self._check_writers() self._check_loops() selected_tasks, selected_evals = self._selection(tasks, evals) - closure = self._closure(selected_tasks + selected_evals) + return self._planned(selected_tasks + selected_evals, selected_tasks, selected_evals) + + def plan_entities(self, ids: Sequence[str]) -> Plan: + self._check_writers() + self._check_loops() + return self._planned(self._select_entities(ids), [], []) + + def _planned(self, roots: list[ResolvedEntry], selected_tasks: list[ResolvedEntry], + selected_evals: list[ResolvedEntry]) -> Plan: + closure = self._closure(roots) store_refs, store_latest = self._check_store_refs(closure) self._build_tasks(closure) self._check_files(closure) @@ -415,6 +431,14 @@ def _select(self, kind: BundleKind, wanted: Sequence[str]) -> list[ResolvedEntry self.problems.append(f"{flag}: this bundle has no {kind.store} with that name or id{listed}") return list(selected.values()) + def _select_entities(self, ids: Sequence[str]) -> list[ResolvedEntry]: + entities = {_nfc(entry.entry.id): entry for entry in self.entries + if entry.entry.kind not in (BundleKind.TASK, BundleKind.EVAL)} + for entity_id in ids: + if _nfc(entity_id) not in entities: + self.problems.append(f"this bundle has no env, agent or artifact {entity_id!r}") + return [entities[key] for key in dict.fromkeys(map(_nfc, ids)) if key in entities] + def _closure(self, roots: list[ResolvedEntry]) -> set[_Key]: closure: set[_Key] = set() frontier = [_key(root.entry) for root in roots] diff --git a/src/agent_env/task/__init__.py b/src/agent_env/task/__init__.py index fbaffc5f4..a16aa8032 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,13 +1,13 @@ """Task module for AgentEnv.""" from .task import Task -from .episode import Episode, EpisodeError, EpisodeSpec, Skill, Turn from .store import ( TaskStore, TaskQuery, get_task_store, set_task_store, reset_task_store, TaskInstance, TaskInstanceStore, get_task_instance_store, set_task_instance_store, reset_task_instance_store, TaskStepResult, TaskStepStatus, SupersededAttemptError, record_step_complete, record_task_cancelled, record_task_failure, ) +from .episode import Episode, EpisodeError, EpisodeSpec, Skill, Turn __all__ = [ "Task", diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 176ce491e..e2a4819ca 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -1,7 +1,7 @@ """Which steps a task may have when some of its agents are played from outside AgentEnv, and why the others are refused. -An outside agent is played through a stand-in agent deployed in its place (``agentenv_protocol.a2a_agent.stand_in``), -so steps that act on the agent, or in its sandbox, act on the stand-in. Every registered step type is either checked +An outside agent is played through a stand-in agent deployed in its place (the ``agents`` bundle's ``stand-in``), so +steps that act on the agent, or in its sandbox, act on the stand-in. Every registered step type is either checked here (``CHECKS``) or listed as untouched by an outside agent (``UNAFFECTED``). A test holds the registry to the two, so a new step type is classified before an episode runs it. Plugin step types aren't inspected. """ diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py index f6df0f33c..09cb9b1e6 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -5,13 +5,16 @@ from agent_env.a2a_agent import protocol from agent_env.a2a_agent.a2a_agent import A2AAgent -from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent +from agent_env.bundle.installed import write_installed from agent_env.task.episode.run import _episode_of from agent_env.task_step.context import TaskStepContext from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url +# The agents bundle's stand-in, deployed in the place of an agent an episode plays. +STAND_IN_AGENT_ID = "@local/agentenv-framework/agents/stand-in" + class _EpisodeAgentStep(TaskStep): """An outside agent's ``deploy_agent`` or ``install_agent``. In a task that uses the agent's sandbox it runs that @@ -52,7 +55,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: config = A2AAgent.negotiate_agent_config(deployed.a2a_card, {"system_prompt": system_prompt}) await protocol.post_agent_config(deployed.a2a_url + config.endpoint, config.fields) return context - await asyncio.to_thread(stand_in_agent) + await asyncio.to_thread(write_installed, STAND_IN_AGENT_ID) override = {"a2a_agent_id": STAND_IN_AGENT_ID} if self.step.type == "deploy_agent" else { "stand_in": STAND_IN_AGENT_ID} step_params = context.metadata.setdefault("user_overrides", {}).setdefault("step_params", {}) diff --git a/tst/unit/a2a_agent/stand_in_test.py b/tst/unit/a2a_agent/stand_in_test.py deleted file mode 100644 index 41d7c3ca4..000000000 --- a/tst/unit/a2a_agent/stand_in_test.py +++ /dev/null @@ -1,52 +0,0 @@ -"""The stand-in agent the Episode API deploys in place of an agent played from outside: written to the store as an -image's build context, and written again only when that context changes.""" - -from importlib.metadata import version - -from agent_env.a2a_agent import stand_in -from agent_env.a2a_agent.a2a_agent import TAKES_CREDENTIALS, A2AAgent -from agent_env.a2a_agent.stand_in import STAND_IN_AGENT_ID, stand_in_agent - - -def test_the_stand_in_is_a_build_context_that_takes_no_credentials(local_stores, cli_routing): - agent = stand_in_agent() - - assert agent.id == STAND_IN_AGENT_ID - assert agent.docker_image_artifact.context_only - assert (agent.default_env_vars, agent.metadata) == ({}, {TAKES_CREDENTIALS: False}) - - -class _NoModelConfig: - def __getattr__(self, name): - raise AssertionError(f"an agent that takes no credentials asked config for {name}") - - -def test_an_agent_that_takes_no_credentials_gets_no_model_endpoint_key_or_headers_whoever_passes_them(): - agent = A2AAgent(id="a", version=1, docker_image_artifact=None, metadata={TAKES_CREDENTIALS: False}) - passed = {"LITELLM_API_KEY": "sk-from-the-step", "LITELLM_BASE_URL": "http://model", "LITELLM_EXTRA_HEADERS": "{}", - "GITHUB_TOKEN": "the task's own"} - - assert agent._with_model_credentials(passed, _NoModelConfig()) == {"GITHUB_TOKEN": "the task's own"} - - -def test_an_unchanged_context_reuses_the_stand_in_and_a_changed_one_writes_it_again( - local_stores, cli_routing, monkeypatch): - first = stand_in_agent() - assert stand_in_agent().version == first.version - - monkeypatch.setattr(stand_in, "_requirements", lambda: ["httpx==0.0.1"]) - - assert stand_in_agent().version == first.version + 1 - - -def test_its_context_holds_the_server_the_protocol_package_and_the_pinned_requirements_to_serve(tmp_path): - context = stand_in._stage(tmp_path) - - assert sorted(path.name for path in context.iterdir()) == [ - "Dockerfile", "a2a_server.py", "agentenv_protocol", "requirements.txt"] - assert (context / "agentenv_protocol" / "a2a_agent" / "stand_in.py").is_file() - assert not list(context.rglob("__pycache__")) - pinned = dict(line.split("==") for line in (context / "requirements.txt").read_text().splitlines()) - assert {"a2a-sdk[http-server]", "httpx", "pydantic", "starlette", "uvicorn"} <= pinned.keys() - assert not pinned.keys() & {"fastmcp", "pytest", "pytest-asyncio"} - assert all(pinned[name] == version(name.split("[")[0]) for name in pinned) diff --git a/tst/unit/agents/stand_in/__init__.py b/tst/unit/agents/stand_in/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/unit/agents/stand_in/agent.py b/tst/unit/agents/stand_in/agent.py new file mode 100644 index 000000000..7bf5d791a --- /dev/null +++ b/tst/unit/agents/stand_in/agent.py @@ -0,0 +1,12 @@ +"""The stand-in's a2a_server, imported from its folder by its path: Claude Code's, imported by the same name, may be +loaded already.""" + +import importlib.util +from pathlib import Path + +import agent_env + +AGENT_DIR = Path(agent_env.__file__).parent / "agents" / "agents" / "stand-in" +_spec = importlib.util.spec_from_file_location("stand_in_a2a_server", AGENT_DIR / "a2a_server.py") +a2a_server = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(a2a_server) diff --git a/tst/unit/agents/stand_in/test_stand_in_agent.py b/tst/unit/agents/stand_in/test_stand_in_agent.py new file mode 100644 index 000000000..d01cd7c31 --- /dev/null +++ b/tst/unit/agents/stand_in/test_stand_in_agent.py @@ -0,0 +1,56 @@ +"""The agents bundle's stand-in: its card, how it installs on a VM's host, and how the episode writes it.""" + +from __future__ import annotations + +import re +import subprocess + +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.bundle.installed import write_installed +from agent_env.task.episode.steps import STAND_IN_AGENT_ID +from agent_env.task_step.task_steps.install_agent import _build_param_resolvers + +from tst.unit.agents.stand_in.agent import AGENT_DIR, a2a_server + + +def _extensions() -> dict: + card = a2a_server.AGENT_CARD.model_dump(mode="json", exclude_none=True, by_alias=True) + return {ext["uri"]: ext for ext in card["capabilities"]["extensions"]} + + +def test_the_card_offers_the_driver_and_a_host_install_only(): + extensions = _extensions() + assert {"urn:agentenv:driver/v1", "urn:agentenv:mcp-config/v1", "urn:agentenv:trajectory/v1", + "urn:agentenv:install/v1"} <= set(extensions) + params = extensions["urn:agentenv:install/v1"]["params"] + assert params["install_commands_host"] and not params.get("install_commands") + + +def test_the_host_install_asks_only_for_params_install_agent_supplies_and_runs_the_folders_script(): + params = _extensions()["urn:agentenv:install/v1"]["params"] + supplied = set(_build_param_resolvers(container_name=None, work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, + config=None, agent_name="a")) + placeholders = {name for command in params["install_commands_host"] for name in re.findall(r"\{(\w+)\}", command)} + assert placeholders == set(params["required_params_host"]) and placeholders <= supplied + assert params["install_commands_host"][-1] == "sh {work_dir}/agent/install.sh {agent_name} {a2a_port} {workdir} {user}" + + +def test_the_install_script_parses_and_the_image_keeps_it_in_its_build_context(): + subprocess.run(["sh", "-n", str(AGENT_DIR / "install.sh")], check=True) + assert "COPY a2a_server.py install.sh ./" in (AGENT_DIR / "Dockerfile").read_text() + + +def test_the_image_works_in_app_as_agent(): + lines = (AGENT_DIR / "Dockerfile").read_text().splitlines() + last = {key: next(line.split(maxsplit=1)[1] for line in reversed(lines) if line.startswith(f"{key} ")) + for key in ("USER", "WORKDIR")} + assert last == {"USER": "agent", "WORKDIR": "/app"} + + +def test_the_episode_writes_it_from_the_bundle_taking_no_credentials(local_stores, cli_routing): + version = write_installed(STAND_IN_AGENT_ID) + agent = A2AAgent.get(STAND_IN_AGENT_ID) + + assert (agent.version, agent.metadata) == (version, {"takes_credentials": False}) + assert agent.docker_image_artifact.context_only + assert write_installed(STAND_IN_AGENT_ID) == version diff --git a/tst/unit/bundle/installed_entities_test.py b/tst/unit/bundle/installed_entities_test.py index 26af5b58c..317228ebb 100644 --- a/tst/unit/bundle/installed_entities_test.py +++ b/tst/unit/bundle/installed_entities_test.py @@ -12,10 +12,12 @@ from agent_env.bundle import BundleError, BundleKind, dry_run_bundle, parse_bundle from agent_env.bundle import preflight as preflight_module from agent_env.bundle import run as run_module -from agent_env.bundle.installed import InstalledBundle +from agent_env.a2a_agent.a2a_agent import A2AAgent +from agent_env.bundle import installed as installed_module +from agent_env.bundle.installed import InstalledBundle, write_installed from agent_env.bundle.ledger import locked_ids from agent_env.bundle.materialize import materialize -from agent_env.bundle.plan import check_bundle, plan_bundle +from agent_env.bundle.plan import check_bundle, plan_bundle, plan_entities from agent_env.bundle.resolve import resolve_bundle from tst.unit.bundle._support import RefusingStore, layout, local_builds @@ -184,6 +186,48 @@ def test_a_pulled_entity_naming_an_image_folder_of_its_bundle_pulls_the_image(bu (BundleKind.TASK, f"{ROOT}/t")] + +def test_an_entity_planned_alone_is_written_with_what_it_references_and_nothing_runs(tmp_path): + kit = _kit(tmp_path, {"artifacts/base/Dockerfile": "FROM scratch\n", "agents/coder/agent.toml": 'image = "base"\n'}) + plan = plan_entities(resolve_bundle(parse_bundle(kit.root, id_root=KIT)), [f"{KIT}/coder"], installed=lambda: (kit,)) + + assert _writes(plan) == [(BundleKind.ARTIFACT, f"{KIT}/base"), (BundleKind.AGENT, f"{KIT}/coder")] + assert (plan.tasks, plan.evals) == ((), ()) + + +def test_planning_an_entity_the_bundle_doesnt_define_is_refused(tmp_path): + kit = _kit(tmp_path, {"tasks/t.json": json.dumps([{"id": "agent", "type": "deploy_agent", "a2a_agent_id": SOLVER}])}) + + with pytest.raises(BundleError) as caught: + plan_entities(resolve_bundle(parse_bundle(kit.root, id_root=KIT)), [f"{KIT}/coder", f"{KIT}/t"]) + assert caught.value.problems == (f"this bundle has no env, agent or artifact '{KIT}/coder'", + f"this bundle has no env, agent or artifact '{KIT}/t'") + + +def test_an_installed_bundles_agent_is_written_by_its_id_and_anew_only_once_it_changes( + tmp_path, monkeypatch, local_stores, cli_routing): + kit = _kit(tmp_path) + monkeypatch.setattr(installed_module, "installed_bundles", lambda: (kit,)) + + version = write_installed(SOLVER) + assert A2AAgent.get(SOLVER).docker_image_artifact.id == f"{SOLVER}__agent_image" + assert write_installed(SOLVER) == version + (kit.root / "agents/solver/Dockerfile").write_text("FROM scratch\nENV CHANGED=1\n") + assert write_installed(SOLVER) == version + 1 + + +@pytest.mark.parametrize(("installed", "problem"), [ + ((), "'@local/demo-dist/kit/solver' isn't an entity of an installed bundle"), + (("broken",), "demo-dist/kit: its folder is gone"), +]) +def test_an_entity_no_installed_bundle_can_write_is_refused(tmp_path, monkeypatch, installed, problem): + bundles = tuple(_kit(tmp_path, problem="its folder is gone") for _ in installed) + monkeypatch.setattr(installed_module, "installed_bundles", lambda: bundles) + + with pytest.raises(BundleError) as caught: + write_installed(SOLVER) + assert caught.value.problems == (problem,) + def test_a_pulled_entity_naming_the_running_bundles_own_entity_uses_it(tmp_path, local_stores, cli_routing): base = _installed(tmp_path / "site" / "base", { "artifacts/base/Dockerfile": "FROM scratch\n", diff --git a/tst/unit/bundle/materialize_test.py b/tst/unit/bundle/materialize_test.py index 102c04148..ebed67c49 100644 --- a/tst/unit/bundle/materialize_test.py +++ b/tst/unit/bundle/materialize_test.py @@ -347,6 +347,7 @@ def refuse(ref, auth): [metadata] default_model = "claude-sonnet-4-6" min_disk_size_gb = 20 +takes_credentials = false """ @@ -359,7 +360,7 @@ def test_an_agent_is_written_from_its_agent_toml_over_a_store_image(bundle_dir): agent = A2AAgent.get(f"{ROOT}/solver") assert (agent.docker_image_artifact.id, agent.docker_image_artifact.version) == ("claude-image", 1) assert (agent.default_env_vars, agent.metadata) == ( - {"LOG_LEVEL": "debug"}, {"default_model": "claude-sonnet-4-6", "min_disk_size_gb": 20}) + {"LOG_LEVEL": "debug"}, {"default_model": "claude-sonnet-4-6", "min_disk_size_gb": 20, "takes_credentials": False}) assert Task.get(f"{ROOT}/t").steps[0].a2a_agent_id == f"{ROOT}/solver" assert _summary(_run(bundle_dir))[f"{ROOT}/solver"] == (1, True, ()) @@ -372,11 +373,16 @@ def test_an_agent_is_written_from_its_agent_toml_over_a_store_image(bundle_dir): ('image = "claude-image"\ndefault_env_vars = { PORT = 8080 }\n', "agent.toml: default_env_vars.PORT must be a string, not 8080"), ('image = "claude-image"\n[metadata]\nowner = "me"\n', - "agent.toml: unknown [metadata] key 'owner'; an agent's [metadata] takes default_model and min_disk_size_gb"), + "agent.toml: unknown [metadata] key 'owner'; an agent's [metadata] takes default_model, min_disk_size_gb and " + "takes_credentials"), ('image = "claude-image"\n[metadata]\nmin_disk_size_gb = "20"\n', "agent.toml: metadata.min_disk_size_gb must be a number, not '20'"), ('image = "claude-image"\n[metadata]\nmin_disk_size_gb = true\n', "agent.toml: metadata.min_disk_size_gb must be a number, not True"), + ('image = "claude-image"\n[metadata]\ntakes_credentials = "no"\n', + "agent.toml: metadata.takes_credentials must be true or false, not 'no'"), + ('image = "claude-image"\n[metadata]\ntakes_credentials = 0\n', + "agent.toml: metadata.takes_credentials must be true or false, not 0"), ('image = "no-such-image"\n', "image: there is no artifact 'no-such-image' in the store"), ('image = "greeting"\n', "image: 'greeting' is this bundle's file, but this field takes docker_image"), ]) @@ -397,8 +403,8 @@ def test_every_problem_in_an_agent_toml_is_refused_together(bundle_dir): assert sorted(_problems(lambda: _run(bundle_dir))) == [ "agents/solver: agent.toml: default_env_vars.PORT must be a string, not 8080", - "agents/solver: agent.toml: unknown [metadata] key 'color'; an agent's [metadata] takes default_model and " - "min_disk_size_gb", + "agents/solver: agent.toml: unknown [metadata] key 'color'; an agent's [metadata] takes default_model, " + "min_disk_size_gb and takes_credentials", "agents/solver: agent.toml: unknown key 'owner'; an a2a_agent agent takes image, default_env_vars, " "metadata, type and id", ] diff --git a/tst/unit/bundle/to_toml_round_trip_test.py b/tst/unit/bundle/to_toml_round_trip_test.py index 6a5fd41a0..bd47bb90c 100644 --- a/tst/unit/bundle/to_toml_round_trip_test.py +++ b/tst/unit/bundle/to_toml_round_trip_test.py @@ -112,17 +112,20 @@ def test_a_website_env_reads_back_with_both_images_built_from_its_one_folder(hom def test_an_agent_reads_back_with_its_env_vars_and_the_metadata_its_toml_carries(home): A2AAgent.put(id="solver", docker_image_artifact=context_image(home, "solver-image"), default_env_vars={"MODE": "fast", "OPENAI_API_KEY": "sk-not-real"}, - metadata={"default_model": "claude", "min_disk_size_gb": 20, "agent_card": {"name": "x"}}) + metadata={"default_model": "claude", "min_disk_size_gb": 20, "takes_credentials": False, + "agent_card": {"name": "x"}}) _env(home, "crm") steps = [{"id": "env", "type": "deploy_env", "env_id": "crm"}, {"id": "agent", "type": "deploy_agent", "env_ids": ["crm"], "a2a_agent_id": "solver"}] with _read_back(home, steps, "agent", "solver", user=True) as (read, imported): assert read.default_env_vars == {"MODE": "fast"} - assert read.metadata == {"default_model": "claude", "min_disk_size_gb": 20} - assert {(drop.where, drop.why.split(",")[0]) for drop in imported.dropped} >= { - ("agents/solver: metadata", "an agent.toml's [metadata] takes only default_model and min_disk_size_gb"), - ("agents/solver/agent.toml: default_env_vars.OPENAI_API_KEY", "dropped: its name looks like a secret")} + assert read.metadata == {"default_model": "claude", "min_disk_size_gb": 20, "takes_credentials": False} + whys = {drop.where: drop.why for drop in imported.dropped} + assert whys["agents/solver: metadata"] == ( + "an agent.toml's [metadata] takes only default_model, min_disk_size_gb and takes_credentials, not agent_card") + assert whys["agents/solver/agent.toml: default_env_vars.OPENAI_API_KEY"].startswith( + "dropped: its name looks like a secret") @pytest.mark.parametrize("description", ["rows.csv", "the crm's rows"]) diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index c2525076e..6cb064023 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -312,7 +312,7 @@ def test_the_trust_lines_resolve_as_the_trust_env_param(): resolvers = _build_param_resolvers(container_name="c", work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, config=None, agent_name="a", trust_lines="SSL_CERT_FILE=/x\n") assert resolvers["trust_env"]() == "SSL_CERT_FILE=/x\n" -STAND_IN = "@local/agentenv-stand-in" +STAND_IN = "@local/agentenv-framework/agents/stand-in" CODEX_INSTALL = {"params": {"workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", "install_commands_host": ["install codex"], "a2a_port": 8000}} STAND_IN_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], From 83a636eff2ff0883e4b818d94376666f83d9621e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 13:43:58 -0700 Subject: [PATCH 10/39] refactor(protocol): ProxyAgent and urn:agentenv:proxy/v1, without what nothing uses StandInAgent is ProxyAgent and driver/v1 is proxy/v1, with routes under /ext/proxy and its models named Proxy*. Gone: the shell=False mode, whose refusals guarded calls no caller makes; mcp-config, whose registrations nothing read, with role taken through agent-config instead so a deploy naming one still hands it over; and the module's own serve entry point. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/{stand_in.py => proxy.py} | 0 .../tests/{test_a2a_stand_in.py => test_a2a_proxy.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/{stand_in.py => proxy.py} (100%) rename packages/agentenv-protocol/tests/{test_a2a_stand_in.py => test_a2a_proxy.py} (100%) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py similarity index 100% rename from packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/stand_in.py rename to packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py diff --git a/packages/agentenv-protocol/tests/test_a2a_stand_in.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py similarity index 100% rename from packages/agentenv-protocol/tests/test_a2a_stand_in.py rename to packages/agentenv-protocol/tests/test_a2a_proxy.py From c9cbae752acd22aa15de285252a0b8759ec36fd5 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Wed, 21 Oct 2026 13:44:15 -0700 Subject: [PATCH 11/39] refactor(protocol): the content of the ProxyAgent rename the previous commit left out The previous commit moved stand_in.py and its tests to their new names but not their new content, nor the extension's, the exports' or the README's. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 36 +++--- .../agentenv_protocol/a2a_agent/__init__.py | 44 +++---- .../agentenv_protocol/a2a_agent/extensions.py | 58 ++++----- .../src/agentenv_protocol/a2a_agent/proxy.py | 113 ++++++++---------- .../agentenv-protocol/tests/test_a2a_proxy.py | 47 ++++---- 5 files changed, 139 insertions(+), 159 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 366341a7e..66561ca02 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -697,30 +697,30 @@ The answer is agent-wide rather than per caller, and like every extension route probe carries no access check of its own, so expose an agent only to the control plane that drives it. -`DRIVER_V1` hands an agent's turns to a driver outside its sandbox, and `StandInAgent` -(`agentenv_protocol.a2a_agent.stand_in`) serves it. The stand-in has no model: each message -it gets stays an open turn, its task working, until the driver replies. The driver's first -`POST /ext/driver/claim` (`{"token"}`, at least 22 characters) sets the token every later -call carries in its body. A claim with another token gets 409, and any other call with a -wrong one gets 401. The token is held in memory only, never in the environment a command -inherits. The other operations: - -- `POST /ext/driver/next` (`{"token", "after"}`) waits up to 25 s for the turns' version to +`PROXY_V1` hands an agent's turns to a driver outside its sandbox, and `ProxyAgent` +(`agentenv_protocol.a2a_agent.proxy`) serves it. The proxy has no model: each message it +gets stays an open turn, its task working, until the driver replies. Its agent-config takes +`system_prompt` and `role`, and it offers no mcp-config: the driver reaches the envs itself, +as the role. The driver's first `POST /ext/proxy/claim` (`{"token"}`, at least 22 +characters) sets the token every later call carries in its body. A claim with another token +gets 409, and any other call with a wrong one gets 401. The token is held in memory only, +never in the environment a command inherits. The other operations: + +- `POST /ext/proxy/next` (`{"token", "after"}`) waits up to 25 s for the turns' version to pass `after`, then returns `{"version", "turns": [{"task_id", "context_id", "text", "system_prompt", "skills"}]}`. A turn its sender cancels drops out of the list. -- `POST /ext/driver/reply` (`{"token", "task_id", "text"}` or `"error"` in place of +- `POST /ext/proxy/reply` (`{"token", "task_id", "text"}` or `"error"` in place of `"text"`, plus `"trajectory"` and `"tool_call_count"`) ends a turn. A text is the task's - result, an error fails it with `stand_in.driver_error`, and the trajectory's spans become - the task's trajectory. -- `POST /ext/driver/exec` (`{"token", "command"}`) starts `bash -c` in the stand-in's working - folder, as its user, and returns `{"exec_id"}`. `POST /ext/driver/output` (`{"token", + result, an error fails it with `proxy.driver_error`, and the trajectory's spans become the + task's trajectory, as format `agentenv-proxy/v1`. +- `POST /ext/proxy/exec` (`{"token", "command"}`) starts `bash -c` in the proxy's working + folder, as its user, and returns `{"exec_id"}`. `POST /ext/proxy/output` (`{"token", "exec_id"}`) waits up to 25 s and returns `{"done", "output"}`. The output is stdout then stderr, or an error line after 3600 s, when the command's process group is killed. -Skills the stand-in is given are written to `~/skills//`, and a turn names each one's -folder. `StandInAgent(shell=False)` stands in for an agent that works in no sandbox: it keeps -no skills and refuses `exec`. Commands inherit the stand-in's environment, so run it with no -credentials in it. +Skills the proxy is given are written to `~/skills//`, and a turn names each one's +folder. Commands inherit the proxy's environment, so run it with no credentials in it, and +expose it only to the control plane that drives it. Runnable, self-contained reference agents live in [`examples/`](https://github.com/scaleapi/agentenv-framework/tree/main/packages/agentenv-protocol/examples): the normal, streaming, and multimodal `run(request)` paths, single- and diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index 605fe39d1..8961bd91b 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -5,7 +5,7 @@ from .extensions import ( AGENT_CONFIG_V1, ATTRIBUTION_PROBE_V1, - DRIVER_V1, + PROXY_V1, INSTALL_V1, MCP_CONFIG_V1, PEER_AGENTS_V1, @@ -20,16 +20,16 @@ ContextEventsTrajectoryRequest, ContextObjectTrajectoryRequest, ContextTrajectoryRequest, - DriverClaimRequest, - DriverExecRequest, - DriverExecResponse, - DriverNextRequest, - DriverNextResponse, - DriverOutputRequest, - DriverOutputResponse, - DriverReplyRequest, - DriverSkill, - DriverTurn, + ProxyClaimRequest, + ProxyExecRequest, + ProxyExecResponse, + ProxyNextRequest, + ProxyNextResponse, + ProxyOutputRequest, + ProxyOutputResponse, + ProxyReplyRequest, + ProxySkill, + ProxyTurn, ExtensionActivation, ExtensionConfiguration, ExtensionDefinition, @@ -128,7 +128,7 @@ __all__ = [ "AGENT_CONFIG_V1", "ATTRIBUTION_PROBE_V1", - "DRIVER_V1", + "PROXY_V1", "INSTALL_V1", "MCP_CONFIG_V1", "PEER_AGENTS_V1", @@ -151,16 +151,16 @@ "ContextEventsTrajectoryRequest", "ContextObjectTrajectoryRequest", "ContextTrajectoryRequest", - "DriverClaimRequest", - "DriverExecRequest", - "DriverExecResponse", - "DriverNextRequest", - "DriverNextResponse", - "DriverOutputRequest", - "DriverOutputResponse", - "DriverReplyRequest", - "DriverSkill", - "DriverTurn", + "ProxyClaimRequest", + "ProxyExecRequest", + "ProxyExecResponse", + "ProxyNextRequest", + "ProxyNextResponse", + "ProxyOutputRequest", + "ProxyOutputResponse", + "ProxyReplyRequest", + "ProxySkill", + "ProxyTurn", "DataPart", "DefaultExtensionHandlers", "ExtensionActivation", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index 82df2bd90..77223b2c3 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -287,38 +287,38 @@ class TriggerDecideRequest(ExtensionRequest): env_triggers: dict[str, dict[str, Any]] | None = None -DRIVER_TOKEN_MIN_LENGTH = 22 +PROXY_TOKEN_MIN_LENGTH = 22 -class DriverClaimRequest(ExtensionRequest): - token: str = Field(min_length=DRIVER_TOKEN_MIN_LENGTH) +class ProxyClaimRequest(ExtensionRequest): + token: str = Field(min_length=PROXY_TOKEN_MIN_LENGTH) -class DriverNextRequest(ExtensionRequest): +class ProxyNextRequest(ExtensionRequest): token: str after: int = Field(default=0, ge=0) -class DriverSkill(ExtensionResponse): +class ProxySkill(ExtensionResponse): name: str description: str - path: str | None = None + path: str -class DriverTurn(ExtensionResponse): +class ProxyTurn(ExtensionResponse): task_id: str context_id: str text: str system_prompt: str | None = None - skills: list[DriverSkill] = Field(default_factory=list) + skills: list[ProxySkill] = Field(default_factory=list) -class DriverNextResponse(ExtensionResponse): +class ProxyNextResponse(ExtensionResponse): version: int - turns: list[DriverTurn] + turns: list[ProxyTurn] -class DriverReplyRequest(ExtensionRequest): +class ProxyReplyRequest(ExtensionRequest): token: str task_id: str text: str | None = None @@ -327,27 +327,27 @@ class DriverReplyRequest(ExtensionRequest): tool_call_count: int = Field(default=0, ge=0) @model_validator(mode="after") - def _text_or_error(self) -> DriverReplyRequest: + def _text_or_error(self) -> ProxyReplyRequest: if (self.text is None) == (self.error is None): raise ValueError("a reply carries exactly one of text and error") return self -class DriverExecRequest(ExtensionRequest): +class ProxyExecRequest(ExtensionRequest): token: str command: str -class DriverExecResponse(ExtensionResponse): +class ProxyExecResponse(ExtensionResponse): exec_id: str -class DriverOutputRequest(ExtensionRequest): +class ProxyOutputRequest(ExtensionRequest): token: str exec_id: str -class DriverOutputResponse(ExtensionResponse): +class ProxyOutputResponse(ExtensionResponse): done: bool output: str | None = None @@ -1232,24 +1232,24 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: configuration_validator=_configure_install, ) -DRIVER_V1 = ExtensionDefinition( - uri="urn:agentenv:driver/v1", - description="Hand the agent's turns to a driver outside its sandbox, and run its commands where the agent works.", - endpoint="/ext/driver", +PROXY_V1 = ExtensionDefinition( + uri="urn:agentenv:proxy/v1", + description="Hand the agent's turns to a driver outside its sandbox, and run the driver's commands where the agent works.", + endpoint="/ext/proxy", core_operations={ - "claim": _op("claim", "POST", "/ext/driver/claim", RUNTIME, request_model=DriverClaimRequest), + "claim": _op("claim", "POST", "/ext/proxy/claim", RUNTIME, request_model=ProxyClaimRequest), "next": _op( - "next", "POST", "/ext/driver/next", RUNTIME, - request_model=DriverNextRequest, response_model=DriverNextResponse, + "next", "POST", "/ext/proxy/next", RUNTIME, + request_model=ProxyNextRequest, response_model=ProxyNextResponse, ), - "reply": _op("reply", "POST", "/ext/driver/reply", RUNTIME, request_model=DriverReplyRequest), + "reply": _op("reply", "POST", "/ext/proxy/reply", RUNTIME, request_model=ProxyReplyRequest), "exec": _op( - "exec", "POST", "/ext/driver/exec", RUNTIME, - request_model=DriverExecRequest, response_model=DriverExecResponse, + "exec", "POST", "/ext/proxy/exec", RUNTIME, + request_model=ProxyExecRequest, response_model=ProxyExecResponse, ), "output": _op( - "output", "POST", "/ext/driver/output", RUNTIME, - request_model=DriverOutputRequest, response_model=DriverOutputResponse, + "output", "POST", "/ext/proxy/output", RUNTIME, + request_model=ProxyOutputRequest, response_model=ProxyOutputResponse, ), }, ) @@ -1267,7 +1267,7 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: ATTRIBUTION_PROBE_V1, TRIGGERS_V1, INSTALL_V1, - DRIVER_V1, + PROXY_V1, ) } ) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 00e915be2..97cd950e2 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -1,7 +1,7 @@ -"""The stand-in: an agent with no model, played by a driver outside its sandbox. +"""The proxy: an agent with no model, whose turns a driver outside its sandbox answers. -Each message it gets is an open turn until the driver replies to it. Through ``urn:agentenv:driver/v1`` the driver -reads the open turns, replies to them and runs commands where the stand-in works, as it; its first ``claim`` sets the +Each message it gets is an open turn until the driver replies to it. Through ``urn:agentenv:proxy/v1`` the driver +reads the open turns, replies to them and runs commands where the proxy works, as it; its first ``claim`` sets the token every later call carries. A turn its sender cancels closes without a reply. """ @@ -19,22 +19,21 @@ from starlette.exceptions import HTTPException from .extensions import ( - DRIVER_V1, - MCP_CONFIG_V1, + PROXY_V1, SKILL_CONFIG_V1, TRAJECTORY_V1, BundleSkillRequest, - DriverClaimRequest, - DriverExecRequest, - DriverExecResponse, - DriverNextRequest, - DriverNextResponse, - DriverOutputRequest, - DriverOutputResponse, - DriverReplyRequest, - DriverSkill, - DriverTurn, InlineSkillRequest, + ProxyClaimRequest, + ProxyExecRequest, + ProxyExecResponse, + ProxyNextRequest, + ProxyNextResponse, + ProxyOutputRequest, + ProxyOutputResponse, + ProxyReplyRequest, + ProxySkill, + ProxyTurn, extension, ) from .framework import AgentEnvAgent, AgentIdentity, a2a_agent @@ -44,11 +43,13 @@ LONG_POLL_SECONDS = 25.0 COMMAND_TIMEOUT_SECONDS = 3600 SKILLS_DIR = Path.home() / "skills" -TRAJECTORY_FORMAT = "agentenv-stand-in/v1" +TRAJECTORY_FORMAT = "agentenv-proxy/v1" -class StandInConfig(AgentConfig): +class ProxyConfig(AgentConfig): system_prompt: str | None = None + # The role a deploy names, which the driver acts as in the envs it reaches itself. + role: str | None = None @dataclass @@ -57,27 +58,25 @@ class _Turn: context_id: str text: str system_prompt: str | None - skills: list[DriverSkill] - reply: asyncio.Future[DriverReplyRequest] = field( + skills: list[ProxySkill] + reply: asyncio.Future[ProxyReplyRequest] = field( default_factory=lambda: asyncio.get_running_loop().create_future() ) IDENTITY = AgentIdentity( - name="agentenv-stand-in", - description="Stands in for an agent played from outside AgentEnv: a driver answers its turns.", + name="agentenv-proxy", + description="An agent played from outside AgentEnv: a driver answers its turns.", version="1.0.0", ) -EXTENSIONS = (MCP_CONFIG_V1, TRAJECTORY_V1, DRIVER_V1) +EXTENSIONS = (TRAJECTORY_V1, PROXY_V1) -@a2a_agent(identity=IDENTITY, config=StandInConfig, extensions=EXTENSIONS) -class StandInAgent(AgentEnvAgent): - """``shell=False`` is a stand-in for an agent that works in no sandbox: it keeps no skills and runs no commands. - ``long_poll_seconds`` bounds how long ``next`` and ``output`` wait for something to report.""" +@a2a_agent(identity=IDENTITY, config=ProxyConfig, extensions=EXTENSIONS) +class ProxyAgent(AgentEnvAgent): + """``long_poll_seconds`` bounds how long ``next`` and ``output`` wait for something to report.""" - def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): - self._shell = shell + def __init__(self, *, long_poll_seconds: float = LONG_POLL_SECONDS): self._long_poll_seconds = long_poll_seconds self._token: str | None = None self._turns: dict[str, _Turn] = {} @@ -85,14 +84,14 @@ def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_S self._changed: asyncio.Future[None] | None = None self._commands: dict[str, asyncio.Task[str]] = {} - async def run(self, request: TaskRequest[StandInConfig]) -> TaskResult: + async def run(self, request: TaskRequest[ProxyConfig]) -> TaskResult: turn = _Turn( task_id=request.task_id, context_id=request.context_id, text="\n".join(part.text for part in request.parts if isinstance(part, TextPart)), system_prompt=request.config.system_prompt, - skills=[DriverSkill(name=skill["name"], description=skill.get("description") or "", - path=str(SKILLS_DIR / skill["name"]) if self._shell else None) + skills=[ProxySkill(name=skill["name"], description=skill.get("description") or "", + path=str(SKILLS_DIR / skill["name"])) for skill in request.skills], ) self._turns[turn.task_id] = turn @@ -106,79 +105,71 @@ async def run(self, request: TaskRequest[StandInConfig]) -> TaskResult: for span in reply.trajectory: request.trajectory.append(span) if reply.error is not None: - return TaskResult.failure("stand_in.driver_error", reply.error) + return TaskResult.failure("proxy.driver_error", reply.error) return TaskResult.text(reply.text, usage=Usage(tool_call_count=reply.tool_call_count)) - @extension(DRIVER_V1.claim) - async def claim(self, request: DriverClaimRequest) -> None: + @extension(PROXY_V1.claim) + async def claim(self, request: ProxyClaimRequest) -> None: if self._token is None: self._token = request.token elif not hmac.compare_digest(self._token, request.token): - raise HTTPException(status_code=409, detail="this stand-in is already claimed") + raise HTTPException(status_code=409, detail="this proxy is already claimed") - @extension(DRIVER_V1.next) - async def next_turns(self, request: DriverNextRequest) -> DriverNextResponse: + @extension(PROXY_V1.next) + async def next_turns(self, request: ProxyNextRequest) -> ProxyNextResponse: self._check(request.token) if self._version <= request.after: with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(asyncio.shield(self._next_change()), self._long_poll_seconds) - return DriverNextResponse(version=self._version, turns=[ - DriverTurn(task_id=turn.task_id, context_id=turn.context_id, text=turn.text, - system_prompt=turn.system_prompt, skills=turn.skills) + return ProxyNextResponse(version=self._version, turns=[ + ProxyTurn(task_id=turn.task_id, context_id=turn.context_id, text=turn.text, + system_prompt=turn.system_prompt, skills=turn.skills) for turn in self._turns.values() if not turn.reply.done() ]) - @extension(DRIVER_V1.reply) - async def reply(self, request: DriverReplyRequest) -> None: + @extension(PROXY_V1.reply) + async def reply(self, request: ProxyReplyRequest) -> None: self._check(request.token) turn = self._turns.get(request.task_id) if turn is None or turn.reply.done(): raise HTTPException(status_code=404, detail=f"no open turn {request.task_id!r}") turn.reply.set_result(request) - @extension(DRIVER_V1.exec) - async def exec_command(self, request: DriverExecRequest) -> DriverExecResponse: + @extension(PROXY_V1.exec) + async def exec_command(self, request: ProxyExecRequest) -> ProxyExecResponse: self._check(request.token) - if not self._shell: - raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it runs no commands") command_id = uuid.uuid4().hex self._commands[command_id] = asyncio.create_task(_run(request.command)) - return DriverExecResponse(exec_id=command_id) + return ProxyExecResponse(exec_id=command_id) - @extension(DRIVER_V1.output) - async def output(self, request: DriverOutputRequest) -> DriverOutputResponse: + @extension(PROXY_V1.output) + async def output(self, request: ProxyOutputRequest) -> ProxyOutputResponse: self._check(request.token) if (command := self._commands.get(request.exec_id)) is None: raise HTTPException(status_code=404, detail=f"no command {request.exec_id!r}") await asyncio.wait((command,), timeout=self._long_poll_seconds) if not command.done(): - return DriverOutputResponse(done=False) + return ProxyOutputResponse(done=False) self._commands.pop(request.exec_id, None) - return DriverOutputResponse(done=True, output=command.result()) + return ProxyOutputResponse(done=True, output=command.result()) @extension(SKILL_CONFIG_V1.add.inline) async def add_inline_skill(self, request: InlineSkillRequest) -> dict[str, str]: - folder = self._skill_folder(request.name) + folder = SKILLS_DIR / request.name folder.mkdir(parents=True) (folder / "SKILL.md").write_text(request.skill_md) return {"name": request.name} @extension(SKILL_CONFIG_V1.add.bundle) async def add_bundle_skill(self, request: BundleSkillRequest) -> dict[str, str]: - folder = self._skill_folder(request.name) for item in request.skill_bundle.files: - await download(item.object, folder / item.path) + await download(item.object, SKILLS_DIR / request.name / item.path) return {"name": request.name} def _check(self, token: str) -> None: if self._token is None or not hmac.compare_digest(self._token, token): raise HTTPException(status_code=401, detail="the driver's token is missing or wrong") - def _skill_folder(self, name: str) -> Path: - if not self._shell: - raise HTTPException(status_code=409, detail="this stand-in works in no sandbox, so it keeps no skills") - return SKILLS_DIR / name - def _next_change(self) -> asyncio.Future[None]: if self._changed is None: self._changed = asyncio.get_running_loop().create_future() @@ -209,7 +200,3 @@ async def _run(command: str) -> str: return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" raise return (stdout + stderr).decode(errors="replace") - - -if __name__ == "__main__": - StandInAgent().serve() diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 4fac062ad..34605b90f 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -7,8 +7,8 @@ import pytest from starlette.testclient import TestClient -from agentenv_protocol.a2a_agent import stand_in -from agentenv_protocol.a2a_agent.stand_in import StandInAgent +from agentenv_protocol.a2a_agent import proxy +from agentenv_protocol.a2a_agent.proxy import ProxyAgent TOKEN = "t" * 32 @@ -39,7 +39,7 @@ def _settled(client: TestClient, task_id: str) -> dict[str, Any]: def _driver(client: TestClient, operation: str, **body: Any): - return client.post(f"/ext/driver/{operation}", json={"token": TOKEN, **body}) + return client.post(f"/ext/proxy/{operation}", json={"token": TOKEN, **body}) def _turns(client: TestClient, after: int = 0) -> dict[str, Any]: @@ -69,31 +69,32 @@ def _state(pid: int) -> str: @pytest.fixture def client(monkeypatch, tmp_path): - monkeypatch.setattr(stand_in, "SKILLS_DIR", tmp_path / "skills") + monkeypatch.setattr(proxy, "SKILLS_DIR", tmp_path / "skills") monkeypatch.chdir(tmp_path) - with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: - assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 + with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 yield client -def test_the_card_advertises_the_driver_and_takes_a_system_prompt(): - card = StandInAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json") +def test_the_card_advertises_the_proxy_and_takes_a_system_prompt_and_a_role_but_no_mcp_servers(): + card = ProxyAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json") extensions = {extension["uri"]: extension for extension in card["capabilities"]["extensions"]} - assert sorted(extensions["urn:agentenv:driver/v1"]["params"]["methods"]) == [ + assert sorted(extensions["urn:agentenv:proxy/v1"]["params"]["methods"]) == [ "claim", "exec", "next", "output", "reply"] - assert "system_prompt" in extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"][ - "supported"] + assert {"system_prompt", "role"} <= set( + extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"]["supported"]) + assert "urn:agentenv:mcp-config/v1" not in extensions def test_the_first_claim_sets_the_token_that_every_driver_call_needs(): - with TestClient(StandInAgent(long_poll_seconds=0.2).create_app()) as client: + with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: assert _driver(client, "next").status_code == 401 - assert client.post("/ext/driver/claim", json={"token": "short"}).status_code == 400 - assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 - assert client.post("/ext/driver/claim", json={"token": TOKEN}).status_code == 200 - assert client.post("/ext/driver/claim", json={"token": "u" * 32}).status_code == 409 - assert client.post("/ext/driver/next", json={"token": "u" * 32}).status_code == 401 + assert client.post("/ext/proxy/claim", json={"token": "short"}).status_code == 400 + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 + assert client.post("/ext/proxy/claim", json={"token": "u" * 32}).status_code == 409 + assert client.post("/ext/proxy/next", json={"token": "u" * 32}).status_code == 401 assert _driver(client, "next").status_code == 200 @@ -164,7 +165,7 @@ def test_exec_returns_stdout_then_stderr_from_the_folder_it_works_in(client, tmp def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(client, monkeypatch, tmp_path): - monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + monkeypatch.setattr(proxy, "COMMAND_TIMEOUT_SECONDS", 0.5) output = _ran(client, "sleep 60 & echo $! > child; wait") @@ -173,7 +174,7 @@ def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(cli def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(client, monkeypatch, tmp_path): - monkeypatch.setattr(stand_in, "COMMAND_TIMEOUT_SECONDS", 0.5) + monkeypatch.setattr(proxy, "COMMAND_TIMEOUT_SECONDS", 0.5) output = _ran(client, "sleep 60 & echo $! > child") @@ -195,11 +196,3 @@ def test_a_skill_is_written_where_its_turns_say_it_is(client, tmp_path): assert skill == [{"name": "deploy", "description": "How to deploy.", "path": str(tmp_path / "skills" / "deploy")}] assert (tmp_path / "skills" / "deploy" / "SKILL.md").read_text() == "# Deploy\n" - -def test_a_stand_in_without_a_shell_keeps_no_skills_and_runs_no_commands(): - with TestClient(StandInAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: - client.post("/ext/driver/claim", json={"token": TOKEN}) - - assert _driver(client, "exec", command="true").status_code == 409 - assert client.post("/ext/skill-config", json={"name": "x", "description": "d", "skill_md": "#"}).status_code \ - == 409 From 639da220896acc2f29e5eb7a8653a70cc2ea8199 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 09:49:17 -0700 Subject: [PATCH 12/39] feat(protocol): every action the driver's model takes runs in the proxy, which records the turn proxy/v1 gains tools, which lists what a turn can call (the env servers registered through mcp-config, whose tools the proxy now calls with ToolSession, and bash), and act and result, which run one action, a tool call or a bash command, by a driver-picked id, answer at once when it ends within the long-poll, and return a structured result. record takes each model call, as an OpenAI chat message, an Anthropic message or a string, and once a turn records one, every action must be one the last recorded call made and all must run before the next record or the reply. The proxy writes the turn's trajectory itself, live: a chat span per recorded call, an execute_tool span per action and a chain span for the turn, so reply carries only parts or an error. Content travels as parts, and a message with a part that isn't text fails before the driver sees it. A turn's actions stop when it closes, and the app's lifespan stops what still runs. Tokens are compared as bytes, a skill's folder is removed when writing it fails and a skill name is taken once, and shell=False offers no bash and keeps no skills. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 64 ++- .../agentenv_protocol/a2a_agent/__init__.py | 32 +- .../agentenv_protocol/a2a_agent/extensions.py | 114 +++-- .../src/agentenv_protocol/a2a_agent/proxy.py | 414 +++++++++++++++--- .../agentenv-protocol/tests/test_a2a_proxy.py | 368 +++++++++++++--- 5 files changed, 810 insertions(+), 182 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 66561ca02..76ad0b97d 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -699,28 +699,48 @@ that drives it. `PROXY_V1` hands an agent's turns to a driver outside its sandbox, and `ProxyAgent` (`agentenv_protocol.a2a_agent.proxy`) serves it. The proxy has no model: each message it -gets stays an open turn, its task working, until the driver replies. Its agent-config takes -`system_prompt` and `role`, and it offers no mcp-config: the driver reaches the envs itself, -as the role. The driver's first `POST /ext/proxy/claim` (`{"token"}`, at least 22 -characters) sets the token every later call carries in its body. A claim with another token -gets 409, and any other call with a wrong one gets 401. The token is held in memory only, -never in the environment a command inherits. The other operations: - -- `POST /ext/proxy/next` (`{"token", "after"}`) waits up to 25 s for the turns' version to - pass `after`, then returns `{"version", "turns": [{"task_id", "context_id", "text", - "system_prompt", "skills"}]}`. A turn its sender cancels drops out of the list. -- `POST /ext/proxy/reply` (`{"token", "task_id", "text"}` or `"error"` in place of - `"text"`, plus `"trajectory"` and `"tool_call_count"`) ends a turn. A text is the task's - result, an error fails it with `proxy.driver_error`, and the trajectory's spans become the - task's trajectory, as format `agentenv-proxy/v1`. -- `POST /ext/proxy/exec` (`{"token", "command"}`) starts `bash -c` in the proxy's working - folder, as its user, and returns `{"exec_id"}`. `POST /ext/proxy/output` (`{"token", - "exec_id"}`) waits up to 25 s and returns `{"done", "output"}`. The output is stdout then - stderr, or an error line after 3600 s, when the command's process group is killed. - -Skills the proxy is given are written to `~/skills//`, and a turn names each one's -folder. Commands inherit the proxy's environment, so run it with no credentials in it, and -expose it only to the control plane that drives it. +gets stays an open turn, its task working, until the driver replies, and every action the +driver's model takes runs in the proxy, where the agent would work. The driver's first +`POST /ext/proxy/claim` (`{"token"}`, at least 22 characters) sets the token every later +call carries in its body. A claim with another token gets 409, and any other call with a +wrong one 401. The token is held in memory only, never in the environment a command +inherits. Content travels as parts, `{"kind": "text", "text": ...}`; a message with any +other part fails as `proxy.unsupported_part` before the driver sees it. The other +operations: + +- `next` (`{"token", "after"}`) waits up to 25 s for the turns' version to pass `after`, + then returns `{"version", "turns": [{"task_id", "context_id", "parts", "system_prompt", + "skills"}]}`. A turn its sender cancels drops out of the list. +- `tools` (`{"token", "task_id"}`) lists what the turn can call, with input schemas: the + tools of each env server registered through mcp-config, and `bash`. Two servers offering + one name are refused with 409. +- `act` (`{"token", "task_id", "action_id", "action"}`) runs one action: `{"kind": "tool", + "server", "name", "arguments"}` calls an env server through `ToolSession`, with the URL + and headers registered for it, and `{"kind": "bash", "command"}` runs `bash -c` in the + proxy's folder, as its user, in a process group of its own; either takes + `"timeout_seconds"` (3600 for `bash`, 600 for a tool by default). The driver picks + `action_id`, so a retried `act` returns the action it already started. It answers + `{"done", "result"}` at once when the action ends within the long-poll; `result` + (`{"token", "action_id"}`) long-polls one that doesn't. A `bash` result is + `{"exit_code", "stdout", "stderr", "timed_out"}`, a tool's `{"content", "is_error"}`. +- `record` (`{"token", "task_id", "message", "model", "usage"}`) records one model call: an + OpenAI chat message, an Anthropic message or a string; anything else gets 400. Once a turn + records a call, each action must be one the last recorded call made, by name and + arguments, and all of them must run before the next `record` or the `reply`, else 409. + Calls to tools the turn doesn't offer aren't checked. +- `reply` (`{"token", "task_id", "parts"}`, or `"error"` in place of `"parts"`) ends the + turn: parts complete the task, and an error fails it as `proxy.driver_error`. + +A turn's actions belong to it: when it closes, by reply or withdrawal, those still running +are stopped, their whole process groups killed, and its results forgotten; the app's +lifespan stops what runs when it shuts down. The proxy writes the turn's trajectory as it +goes, readable live, as OTel GenAI spans in format `agentenv-proxy/v1`: `chat` for each +recorded call, `execute_tool` for each action, and `chain` for the turn. Skills it is given +are written to `~/skills//`, once each, and a turn names each one's folder. +`ProxyAgent(shell=False)` is a proxy for an agent that works in no sandbox: it offers no +`bash` and keeps no skills. It holds the env servers' headers as any agent does, so run it +with no credentials in its environment and expose it only to the control plane that drives +it. Runnable, self-contained reference agents live in [`examples/`](https://github.com/scaleapi/agentenv-framework/tree/main/packages/agentenv-protocol/examples): the normal, streaming, and multimodal `run(request)` paths, single- and diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py index 8961bd91b..67a1b904c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/__init__.py @@ -20,15 +20,23 @@ ContextEventsTrajectoryRequest, ContextObjectTrajectoryRequest, ContextTrajectoryRequest, + ProxyActRequest, + ProxyActionState, + ProxyBashAction, + ProxyBashResult, ProxyClaimRequest, - ProxyExecRequest, - ProxyExecResponse, ProxyNextRequest, ProxyNextResponse, - ProxyOutputRequest, - ProxyOutputResponse, + ProxyRecordRequest, ProxyReplyRequest, + ProxyResultRequest, ProxySkill, + ProxyTextPart, + ProxyTool, + ProxyToolAction, + ProxyToolResult, + ProxyToolsRequest, + ProxyToolsResponse, ProxyTurn, ExtensionActivation, ExtensionConfiguration, @@ -151,15 +159,23 @@ "ContextEventsTrajectoryRequest", "ContextObjectTrajectoryRequest", "ContextTrajectoryRequest", + "ProxyActRequest", + "ProxyActionState", + "ProxyBashAction", + "ProxyBashResult", "ProxyClaimRequest", - "ProxyExecRequest", - "ProxyExecResponse", "ProxyNextRequest", "ProxyNextResponse", - "ProxyOutputRequest", - "ProxyOutputResponse", + "ProxyRecordRequest", "ProxyReplyRequest", + "ProxyResultRequest", "ProxySkill", + "ProxyTextPart", + "ProxyTool", + "ProxyToolAction", + "ProxyToolResult", + "ProxyToolsRequest", + "ProxyToolsResponse", "ProxyTurn", "DataPart", "DefaultExtensionHandlers", diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index 77223b2c3..8e272050f 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -10,7 +10,7 @@ from dataclasses import dataclass, field from enum import Enum from types import MappingProxyType -from typing import Annotated, Any +from typing import Annotated, Any, Literal from pydantic import ( BaseModel, @@ -299,6 +299,11 @@ class ProxyNextRequest(ExtensionRequest): after: int = Field(default=0, ge=0) +class ProxyTextPart(ExtensionRequest): + kind: Literal["text"] = "text" + text: str + + class ProxySkill(ExtensionResponse): name: str description: str @@ -308,7 +313,7 @@ class ProxySkill(ExtensionResponse): class ProxyTurn(ExtensionResponse): task_id: str context_id: str - text: str + parts: list[ProxyTextPart] system_prompt: str | None = None skills: list[ProxySkill] = Field(default_factory=list) @@ -318,38 +323,86 @@ class ProxyNextResponse(ExtensionResponse): turns: list[ProxyTurn] -class ProxyReplyRequest(ExtensionRequest): +class ProxyToolsRequest(ExtensionRequest): token: str task_id: str - text: str | None = None - error: str | None = None - trajectory: list[dict[str, Any]] = Field(default_factory=list) - tool_call_count: int = Field(default=0, ge=0) - @model_validator(mode="after") - def _text_or_error(self) -> ProxyReplyRequest: - if (self.text is None) == (self.error is None): - raise ValueError("a reply carries exactly one of text and error") - return self + +class ProxyTool(ExtensionResponse): + name: str + description: str | None = None + input_schema: dict[str, Any] | None = None + server: str | None = None # the env server that offers it; none for bash -class ProxyExecRequest(ExtensionRequest): - token: str +class ProxyToolsResponse(ExtensionResponse): + tools: list[ProxyTool] + + +class ProxyToolAction(ExtensionRequest): + kind: Literal["tool"] + server: str + name: str + arguments: dict[str, Any] = Field(default_factory=dict) + timeout_seconds: float | None = Field(default=None, gt=0) + + +class ProxyBashAction(ExtensionRequest): + kind: Literal["bash"] command: str + timeout_seconds: float | None = Field(default=None, gt=0) -class ProxyExecResponse(ExtensionResponse): - exec_id: str +class ProxyActRequest(ExtensionRequest): + token: str + task_id: str + action_id: str = Field(min_length=1, max_length=128) + action: Annotated[ProxyToolAction | ProxyBashAction, Field(discriminator="kind")] -class ProxyOutputRequest(ExtensionRequest): +class ProxyResultRequest(ExtensionRequest): token: str - exec_id: str + action_id: str + +class ProxyBashResult(ExtensionResponse): + kind: Literal["bash"] = "bash" + exit_code: int | None # none when the command was stopped + stdout: str + stderr: str + timed_out: bool -class ProxyOutputResponse(ExtensionResponse): + +class ProxyToolResult(ExtensionResponse): + kind: Literal["tool"] = "tool" + content: list[dict[str, Any]] + is_error: bool + + +class ProxyActionState(ExtensionResponse): done: bool - output: str | None = None + result: Annotated[ProxyBashResult | ProxyToolResult, Field(discriminator="kind")] | None = None + + +class ProxyRecordRequest(ExtensionRequest): + token: str + task_id: str + message: str | dict[str, Any] + model: str | None = None + usage: dict[str, Any] | None = None + + +class ProxyReplyRequest(ExtensionRequest): + token: str + task_id: str + parts: list[ProxyTextPart] | None = Field(default=None, min_length=1) + error: str | None = None + + @model_validator(mode="after") + def _parts_or_error(self) -> ProxyReplyRequest: + if (self.parts is None) == (self.error is None): + raise ValueError("a reply carries exactly one of parts and error") + return self def _tuple(values: Iterable[str]) -> tuple[str, ...]: @@ -1234,7 +1287,7 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: PROXY_V1 = ExtensionDefinition( uri="urn:agentenv:proxy/v1", - description="Hand the agent's turns to a driver outside its sandbox, and run the driver's commands where the agent works.", + description="Hand the agent's turns to a driver outside its sandbox, and run the driver's actions where the agent works.", endpoint="/ext/proxy", core_operations={ "claim": _op("claim", "POST", "/ext/proxy/claim", RUNTIME, request_model=ProxyClaimRequest), @@ -1242,15 +1295,20 @@ def _configure_trajectory(*, live: bool = False) -> ExtensionConfiguration: "next", "POST", "/ext/proxy/next", RUNTIME, request_model=ProxyNextRequest, response_model=ProxyNextResponse, ), - "reply": _op("reply", "POST", "/ext/proxy/reply", RUNTIME, request_model=ProxyReplyRequest), - "exec": _op( - "exec", "POST", "/ext/proxy/exec", RUNTIME, - request_model=ProxyExecRequest, response_model=ProxyExecResponse, + "tools": _op( + "tools", "POST", "/ext/proxy/tools", RUNTIME, + request_model=ProxyToolsRequest, response_model=ProxyToolsResponse, ), - "output": _op( - "output", "POST", "/ext/proxy/output", RUNTIME, - request_model=ProxyOutputRequest, response_model=ProxyOutputResponse, + "act": _op( + "act", "POST", "/ext/proxy/act", RUNTIME, + request_model=ProxyActRequest, response_model=ProxyActionState, ), + "result": _op( + "result", "POST", "/ext/proxy/result", RUNTIME, + request_model=ProxyResultRequest, response_model=ProxyActionState, + ), + "record": _op("record", "POST", "/ext/proxy/record", RUNTIME, request_model=ProxyRecordRequest), + "reply": _op("reply", "POST", "/ext/proxy/reply", RUNTIME, request_model=ProxyReplyRequest), }, ) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 97cd950e2..5798a3a09 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -1,8 +1,10 @@ """The proxy: an agent with no model, whose turns a driver outside its sandbox answers. -Each message it gets is an open turn until the driver replies to it. Through ``urn:agentenv:proxy/v1`` the driver -reads the open turns, replies to them and runs commands where the proxy works, as it; its first ``claim`` sets the -token every later call carries. A turn its sender cancels closes without a reply. +Each message it gets is an open turn until the driver replies. Through ``urn:agentenv:proxy/v1`` the driver reads the +open turns, lists what a turn can call, has the proxy run each action where it works (a call to an env's MCP server +registered with it, or a ``bash`` command as its user, in its folder), reports the model's calls, and replies. The proxy +records the turn as it goes. Its first ``claim`` sets the token every later call carries. A turn its sender cancels +closes without a reply, stopping its running actions. """ from __future__ import annotations @@ -10,58 +12,103 @@ import asyncio import contextlib import hmac +import json import os +import shutil import signal -import uuid +from collections.abc import AsyncIterator from dataclasses import dataclass, field +from datetime import datetime, timezone from pathlib import Path +from typing import Any +import httpx +from starlette.applications import Starlette from starlette.exceptions import HTTPException from .extensions import ( + MCP_CONFIG_V1, PROXY_V1, SKILL_CONFIG_V1, TRAJECTORY_V1, BundleSkillRequest, InlineSkillRequest, + ProxyActionState, + ProxyActRequest, + ProxyBashAction, + ProxyBashResult, ProxyClaimRequest, - ProxyExecRequest, - ProxyExecResponse, ProxyNextRequest, ProxyNextResponse, - ProxyOutputRequest, - ProxyOutputResponse, + ProxyRecordRequest, ProxyReplyRequest, + ProxyResultRequest, ProxySkill, + ProxyTextPart, + ProxyTool, + ProxyToolAction, + ProxyToolResult, + ProxyToolsRequest, + ProxyToolsResponse, ProxyTurn, + enable, extension, ) from .framework import AgentEnvAgent, AgentIdentity, a2a_agent -from .tasks.v1 import AgentConfig, TaskRequest, TaskResult, TextPart, Usage +from .tasks.v1 import AgentConfig, TaskRequest, TaskResult, TextPart, TrajectoryLog, Usage +from ..tool_client import ToolSession, ToolSessionError from ..transfers import download LONG_POLL_SECONDS = 25.0 -COMMAND_TIMEOUT_SECONDS = 3600 +BASH_TIMEOUT_SECONDS = 3600 +TOOL_TIMEOUT_SECONDS = 600 SKILLS_DIR = Path.home() / "skills" TRAJECTORY_FORMAT = "agentenv-proxy/v1" +BASH = "bash" +BASH_TOOL = ProxyTool( + name=BASH, + description="Run a bash command where the agent works, as its user, in its folder; returns its stdout, then its " + "stderr.", + input_schema={"type": "object", "properties": {"command": {"type": "string", "description": "The command."}}, + "required": ["command"]}, +) class ProxyConfig(AgentConfig): system_prompt: str | None = None - # The role a deploy names, which the driver acts as in the envs it reaches itself. - role: str | None = None @dataclass class _Turn: task_id: str context_id: str - text: str + parts: list[ProxyTextPart] system_prompt: str | None skills: list[ProxySkill] + servers: dict[str, dict[str, Any]] + trajectory: TrajectoryLog + started: str = field(default_factory=lambda: _now()) reply: asyncio.Future[ProxyReplyRequest] = field( default_factory=lambda: asyncio.get_running_loop().create_future() ) + actions: dict[str, asyncio.Task[ProxyBashResult | ProxyToolResult]] = field(default_factory=dict) + sessions: dict[str, ToolSession] = field(default_factory=dict) + tools: dict[str, ProxyTool] | None = None + listing: asyncio.Lock = field(default_factory=asyncio.Lock) + recorded: bool = False + pending: list[tuple[str, str]] = field(default_factory=list) # (tool, its canonical arguments) not yet run + last_event: str = field(default_factory=lambda: _now()) + model: str | None = None + + +@contextlib.asynccontextmanager +async def stopping_actions(app: Starlette) -> AsyncIterator[None]: + """The proxy's lifespan: what its turns still run when the app stops is stopped with it. A subclass declaring its + own definition passes it on.""" + try: + yield + finally: + await app.state.agentenv_a2a.agent._stop() IDENTITY = AgentIdentity( @@ -69,30 +116,37 @@ class _Turn: description="An agent played from outside AgentEnv: a driver answers its turns.", version="1.0.0", ) -EXTENSIONS = (TRAJECTORY_V1, PROXY_V1) +EXTENSIONS = (MCP_CONFIG_V1, enable(TRAJECTORY_V1, live=True), PROXY_V1) -@a2a_agent(identity=IDENTITY, config=ProxyConfig, extensions=EXTENSIONS) +@a2a_agent(identity=IDENTITY, config=ProxyConfig, extensions=EXTENSIONS, lifespan=stopping_actions) class ProxyAgent(AgentEnvAgent): - """``long_poll_seconds`` bounds how long ``next`` and ``output`` wait for something to report.""" + """``shell=False`` is a proxy for an agent that works in no sandbox: it offers no ``bash`` and keeps no skills. + ``long_poll_seconds`` bounds how long ``next``, ``act`` and ``result`` wait for something to report.""" - def __init__(self, *, long_poll_seconds: float = LONG_POLL_SECONDS): + def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): + self._shell = shell self._long_poll_seconds = long_poll_seconds - self._token: str | None = None + self._token: bytes | None = None self._turns: dict[str, _Turn] = {} self._version = 0 self._changed: asyncio.Future[None] | None = None - self._commands: dict[str, asyncio.Task[str]] = {} + self._actions: dict[str, tuple[_Turn, asyncio.Task[ProxyBashResult | ProxyToolResult]]] = {} async def run(self, request: TaskRequest[ProxyConfig]) -> TaskResult: + if any(not isinstance(part, TextPart) for part in request.parts): + return TaskResult.failure("proxy.unsupported_part", "the proxy passes on text parts only") + request.trajectory.set_format(TRAJECTORY_FORMAT) turn = _Turn( task_id=request.task_id, context_id=request.context_id, - text="\n".join(part.text for part in request.parts if isinstance(part, TextPart)), + parts=[ProxyTextPart(text=part.text) for part in request.parts], system_prompt=request.config.system_prompt, skills=[ProxySkill(name=skill["name"], description=skill.get("description") or "", path=str(SKILLS_DIR / skill["name"])) for skill in request.skills], + servers=dict(request.mcp_servers), + trajectory=request.trajectory, ) self._turns[turn.task_id] = turn self._bump() @@ -101,18 +155,23 @@ async def run(self, request: TaskRequest[ProxyConfig]) -> TaskResult: finally: del self._turns[turn.task_id] self._bump() - request.trajectory.set_format(TRAJECTORY_FORMAT) - for span in reply.trajectory: - request.trajectory.append(span) + await self._close(turn) + prompt = {"prompt": "\n".join(part.text for part in turn.parts), "system_prompt": turn.system_prompt} if reply.error is not None: + turn.trajectory.append(_span("turn", "chain", turn.started, turn.model, prompt=prompt, + completion={"error": reply.error})) return TaskResult.failure("proxy.driver_error", reply.error) - return TaskResult.text(reply.text, usage=Usage(tool_call_count=reply.tool_call_count)) + turn.trajectory.append(_span("turn", "chain", turn.started, turn.model, prompt=prompt, completion={ + "content": [{"type": "text", "text": part.text} for part in reply.parts]})) + return TaskResult.success(parts=tuple(TextPart(text=part.text) for part in reply.parts), + usage=Usage(tool_call_count=len(turn.actions))) @extension(PROXY_V1.claim) async def claim(self, request: ProxyClaimRequest) -> None: + token = request.token.encode() if self._token is None: - self._token = request.token - elif not hmac.compare_digest(self._token, request.token): + self._token = token + elif not hmac.compare_digest(self._token, token): raise HTTPException(status_code=409, detail="this proxy is already claimed") @extension(PROXY_V1.next) @@ -122,54 +181,186 @@ async def next_turns(self, request: ProxyNextRequest) -> ProxyNextResponse: with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(asyncio.shield(self._next_change()), self._long_poll_seconds) return ProxyNextResponse(version=self._version, turns=[ - ProxyTurn(task_id=turn.task_id, context_id=turn.context_id, text=turn.text, + ProxyTurn(task_id=turn.task_id, context_id=turn.context_id, parts=turn.parts, system_prompt=turn.system_prompt, skills=turn.skills) for turn in self._turns.values() if not turn.reply.done() ]) - @extension(PROXY_V1.reply) - async def reply(self, request: ProxyReplyRequest) -> None: - self._check(request.token) - turn = self._turns.get(request.task_id) - if turn is None or turn.reply.done(): - raise HTTPException(status_code=404, detail=f"no open turn {request.task_id!r}") - turn.reply.set_result(request) + @extension(PROXY_V1.tools) + async def list_tools(self, request: ProxyToolsRequest) -> ProxyToolsResponse: + turn = self._open_turn(request.token, request.task_id) + return ProxyToolsResponse(tools=list((await self._tools(turn)).values())) - @extension(PROXY_V1.exec) - async def exec_command(self, request: ProxyExecRequest) -> ProxyExecResponse: + @extension(PROXY_V1.act) + async def act(self, request: ProxyActRequest) -> ProxyActionState: self._check(request.token) - command_id = uuid.uuid4().hex - self._commands[command_id] = asyncio.create_task(_run(request.command)) - return ProxyExecResponse(exec_id=command_id) + if (known := self._actions.get(request.action_id)) is not None: + if known[0].task_id != request.task_id: + raise HTTPException(status_code=409, detail=f"action {request.action_id!r} is another turn's") + return await self._state(known[1]) + turn = self._open_turn(request.token, request.task_id) + action = request.action + if isinstance(action, ProxyBashAction) and not self._shell: + raise HTTPException(status_code=409, detail="this proxy works in no sandbox, so it runs no commands") + if turn.recorded: + call = (action.name, _canonical(action.arguments)) if isinstance(action, ProxyToolAction) else ( + BASH, _canonical({"command": action.command})) + if call not in turn.pending: + raise HTTPException(status_code=409, detail=f"the last recorded model call didn't make this call to " + f"{call[0]!r}") + turn.pending.remove(call) + task = asyncio.create_task(self._perform(turn, action)) + turn.actions[request.action_id] = task + self._actions[request.action_id] = (turn, task) + return await self._state(task) - @extension(PROXY_V1.output) - async def output(self, request: ProxyOutputRequest) -> ProxyOutputResponse: + @extension(PROXY_V1.result) + async def result(self, request: ProxyResultRequest) -> ProxyActionState: self._check(request.token) - if (command := self._commands.get(request.exec_id)) is None: - raise HTTPException(status_code=404, detail=f"no command {request.exec_id!r}") - await asyncio.wait((command,), timeout=self._long_poll_seconds) - if not command.done(): - return ProxyOutputResponse(done=False) - self._commands.pop(request.exec_id, None) - return ProxyOutputResponse(done=True, output=command.result()) + if (known := self._actions.get(request.action_id)) is None: + raise HTTPException(status_code=404, detail=f"no action {request.action_id!r}") + return await self._state(known[1]) + + @extension(PROXY_V1.record) + async def record(self, request: ProxyRecordRequest) -> None: + turn = self._open_turn(request.token, request.task_id) + content, calls = _read_message(request.message) + if turn.pending: + raise HTTPException(status_code=409, detail=_unrun(turn)) + offered = await self._tools(turn) + turn.pending = [(name, _canonical(arguments)) for name, arguments in calls if name in offered] + turn.recorded = True + turn.model = request.model or turn.model + span = _span("chat", "chat", turn.last_event, turn.model, completion={"content": content}) + usage = request.usage or {} + for key, names in (("input_tokens", ("input_tokens", "prompt_tokens")), + ("output_tokens", ("output_tokens", "completion_tokens"))): + if (count := next((usage[name] for name in names if name in usage), None)) is not None: + span["attributes"][f"gen_ai.usage.{key}"] = count + turn.trajectory.append(span) + turn.last_event = span["end_time"] + + @extension(PROXY_V1.reply) + async def reply(self, request: ProxyReplyRequest) -> None: + turn = self._open_turn(request.token, request.task_id) + if request.parts is not None and turn.pending: + raise HTTPException(status_code=409, detail=_unrun(turn)) + turn.reply.set_result(request) @extension(SKILL_CONFIG_V1.add.inline) async def add_inline_skill(self, request: InlineSkillRequest) -> dict[str, str]: - folder = SKILLS_DIR / request.name - folder.mkdir(parents=True) - (folder / "SKILL.md").write_text(request.skill_md) + async with self._skill_folder(request.name) as folder: + (folder / "SKILL.md").write_text(request.skill_md) return {"name": request.name} @extension(SKILL_CONFIG_V1.add.bundle) async def add_bundle_skill(self, request: BundleSkillRequest) -> dict[str, str]: - for item in request.skill_bundle.files: - await download(item.object, SKILLS_DIR / request.name / item.path) + async with self._skill_folder(request.name) as folder: + for item in request.skill_bundle.files: + await download(item.object, folder / item.path) return {"name": request.name} def _check(self, token: str) -> None: - if self._token is None or not hmac.compare_digest(self._token, token): + if self._token is None or not hmac.compare_digest(self._token, token.encode()): raise HTTPException(status_code=401, detail="the driver's token is missing or wrong") + def _open_turn(self, token: str, task_id: str) -> _Turn: + self._check(token) + turn = self._turns.get(task_id) + if turn is None or turn.reply.done(): + raise HTTPException(status_code=404, detail=f"no open turn {task_id!r}") + return turn + + async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: + """What the turn can call, listed once: each registered env server's tools, and ``bash``. Two offering one + name are refused.""" + async with turn.listing: + if turn.tools is None: + offered = {BASH: BASH_TOOL} if self._shell else {} + for server in turn.servers: + try: + listed = await (await self._session(turn, server)).list_tools() + except (ToolSessionError, httpx.HTTPError) as e: + raise HTTPException(status_code=502, detail=f"env server {server!r}: {e}") from None + for tool in listed: + if tool.name in offered: + raise HTTPException(status_code=409, detail=( + f"tool {tool.name!r} is offered by both {offered[tool.name].server or 'the proxy'!r} " + f"and {server!r}")) + offered[tool.name] = ProxyTool(name=tool.name, description=tool.description, + input_schema=tool.inputSchema, server=server) + turn.tools = offered + return turn.tools + + async def _session(self, turn: _Turn, server: str) -> ToolSession: + if server not in turn.sessions: + registration = turn.servers.get(server) + if registration is None: + raise ToolSessionError(f"no env server {server!r} is registered") + turn.sessions[server] = await ToolSession(mcp_url=registration["url"], + headers=registration.get("headers") or None, + timeout=TOOL_TIMEOUT_SECONDS).open() + return turn.sessions[server] + + async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) -> ProxyBashResult | ProxyToolResult: + started = _now() + if isinstance(action, ProxyBashAction): + timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS + result = await _run(action.command, timeout) + name, arguments = BASH, {"command": action.command} + output = (f"error: the command timed out after {timeout:g}s" if result.timed_out + else result.stdout + result.stderr) + completion = {"output": output, "exit_code": result.exit_code} + else: + timeout = action.timeout_seconds or TOOL_TIMEOUT_SECONDS + name, arguments = action.name, action.arguments + try: + session = await self._session(turn, action.server) + called = await asyncio.wait_for(session.call_tool(action.name, action.arguments), timeout) + result = ProxyToolResult(content=called.content, is_error=called.isError) + except asyncio.TimeoutError: + result = _tool_error(f"the tool call timed out after {timeout:g}s") + except (ToolSessionError, httpx.HTTPError) as e: + result = _tool_error(f"{type(e).__name__}: {e}") + text = "\n".join(block["text"] for block in result.content if block.get("type") == "text") + completion = {"error": text} if result.is_error else {"output": text} + span = _span(name, "execute_tool", started, None, prompt={"input": arguments}, completion=completion) + turn.trajectory.append(span) + turn.last_event = span["end_time"] + return result + + async def _state(self, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> ProxyActionState: + await asyncio.wait((task,), timeout=self._long_poll_seconds) + if not task.done(): + return ProxyActionState(done=False) + return ProxyActionState(done=True, result=task.result()) + + async def _close(self, turn: _Turn) -> None: + """Stop what the turn still runs and forget its actions, so what the driver never collected goes with it.""" + for action_id, task in turn.actions.items(): + task.cancel() + self._actions.pop(action_id, None) + await asyncio.gather(*turn.actions.values(), return_exceptions=True) + await asyncio.gather(*(session.close() for session in turn.sessions.values()), return_exceptions=True) + + async def _stop(self) -> None: + await asyncio.gather(*(self._close(turn) for turn in list(self._turns.values())), return_exceptions=True) + + @contextlib.asynccontextmanager + async def _skill_folder(self, name: str) -> AsyncIterator[Path]: + """``name``'s folder, made for the skill and removed again if writing it fails.""" + if not self._shell: + raise HTTPException(status_code=409, detail="this proxy works in no sandbox, so it keeps no skills") + folder = SKILLS_DIR / name + if folder.exists(): + raise HTTPException(status_code=409, detail=f"skill {name!r} is already added") + folder.mkdir(parents=True) + try: + yield folder + except BaseException: + shutil.rmtree(folder, ignore_errors=True) + raise + def _next_change(self) -> asyncio.Future[None]: if self._changed is None: self._changed = asyncio.get_running_loop().create_future() @@ -182,21 +373,122 @@ def _bump(self) -> None: self._changed = None -async def _run(command: str) -> str: - """``command``'s stdout then its stderr, as verifiers' bash tool returns them. It runs in a process group of its - own, stopped with whatever it started if its output isn't done within ``COMMAND_TIMEOUT_SECONDS``, or if the call - is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" +def _read_message(message: str | dict[str, Any]) -> tuple[list[dict[str, Any]], list[tuple[str, dict[str, Any]]]]: + """A model call's content blocks (thinking, text, tool_use) and the tool calls it made, from an OpenAI chat message, + an Anthropic message or a string. Raises a 400 for anything else.""" + if isinstance(message, str): + return [{"type": "text", "text": message}], [] + if message.get("role", "assistant") != "assistant": + raise HTTPException(status_code=400, detail="a recorded model call is the assistant's message") + content = message.get("content") + if isinstance(content, list) and not any(key in message for key in ("tool_calls", "reasoning_content", + "reasoning")): + return _anthropic(content) + if content is None or isinstance(content, (str, list)): + return _openai(message) + raise HTTPException(status_code=400, detail="a recorded model call is an OpenAI chat message, an Anthropic " + "message or a string") + + +def _openai(message: dict[str, Any]) -> tuple[list[dict[str, Any]], list[tuple[str, dict[str, Any]]]]: + content: list[dict[str, Any]] = [] + if isinstance(reasoning := message.get("reasoning_content") or message.get("reasoning"), str) and reasoning: + content.append({"type": "thinking", "thinking": reasoning}) + text = message.get("content") + if isinstance(text, str) and text: + content.append({"type": "text", "text": text}) + elif isinstance(text, list): + content += [_text_block(part) for part in text] + calls = [] + for call in message.get("tool_calls") or []: + function = call.get("function") if isinstance(call, dict) else None + if not isinstance(function, dict) or not isinstance(function.get("name"), str): + raise HTTPException(status_code=400, detail="a tool call names no function") + arguments = function.get("arguments") or {} + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + raise HTTPException(status_code=400, detail=f"tool call {function['name']!r}: its arguments aren't " + "JSON") from None + if not isinstance(arguments, dict): + raise HTTPException(status_code=400, detail=f"tool call {function['name']!r}: its arguments aren't an " + "object") + content.append({"type": "tool_use", "id": call.get("id"), "name": function["name"], "input": arguments}) + calls.append((function["name"], arguments)) + return content, calls + + +def _anthropic(blocks: list[Any]) -> tuple[list[dict[str, Any]], list[tuple[str, dict[str, Any]]]]: + content: list[dict[str, Any]] = [] + calls = [] + for block in blocks: + kind = block.get("type") if isinstance(block, dict) else None + if kind == "tool_use": + if not isinstance(block.get("name"), str) or not isinstance(block.get("input", {}), dict): + raise HTTPException(status_code=400, detail="a tool_use block needs a name and an input object") + content.append({"type": "tool_use", "id": block.get("id"), "name": block["name"], + "input": block.get("input", {})}) + calls.append((block["name"], block.get("input", {}))) + elif kind == "thinking" and isinstance(block.get("thinking"), str): + content.append({"type": "thinking", "thinking": block["thinking"]}) + elif kind == "redacted_thinking": + content.append({"type": "redacted_thinking"}) + else: + content.append(_text_block(block)) + return content, calls + + +def _text_block(block: Any) -> dict[str, Any]: + if not isinstance(block, dict) or block.get("type") != "text" or not isinstance(block.get("text"), str): + raise HTTPException(status_code=400, detail=f"unreadable content block: {block!r:.200}") + return {"type": "text", "text": block["text"]} + + +def _unrun(turn: _Turn) -> str: + return f"the last recorded model call's calls to {sorted({name for name, _ in turn.pending})} haven't run" + + +def _canonical(arguments: dict[str, Any]) -> str: + return json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str) + + +def _tool_error(message: str) -> ProxyToolResult: + return ProxyToolResult(content=[{"type": "text", "text": f"error: {message}"}], is_error=True) + + +def _span(name: str, operation: str, started: str, model: str | None, *, prompt: Any = None, + completion: Any = None) -> dict[str, Any]: + """An OTel GenAI span, the trajectory format rubrics and the explorer read.""" + attributes: dict[str, Any] = {"gen_ai.operation.name": operation} + if model: + attributes["gen_ai.request.model"] = model + if prompt is not None: + attributes["gen_ai.prompt"] = json.dumps(prompt, default=str) + if completion is not None: + attributes["gen_ai.completion"] = json.dumps(completion, default=str) + return {"name": name, "start_time": started, "end_time": _now(), "attributes": attributes} + + +def _now() -> str: + return datetime.now(timezone.utc).isoformat() + + +async def _run(command: str, timeout: float) -> ProxyBashResult: + """Run ``command`` in a process group of its own, stopped with whatever it started if it isn't done within + ``timeout``, or if the call is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" process = await asyncio.create_subprocess_exec( "bash", "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True, ) try: - stdout, stderr = await asyncio.wait_for(process.communicate(), COMMAND_TIMEOUT_SECONDS) + stdout, stderr = await asyncio.wait_for(process.communicate(), timeout) except BaseException as e: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) await process.wait() if isinstance(e, asyncio.TimeoutError): - return f"error: the command timed out after {COMMAND_TIMEOUT_SECONDS}s" + return ProxyBashResult(exit_code=None, stdout="", stderr="", timed_out=True) raise - return (stdout + stderr).decode(errors="replace") + return ProxyBashResult(exit_code=process.returncode, stdout=stdout.decode(errors="replace"), + stderr=stderr.decode(errors="replace"), timed_out=False) diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 34605b90f..dcc5ce9a7 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -1,24 +1,76 @@ from __future__ import annotations +import json +import socket import subprocess +import threading import time from typing import Any import pytest +import uvicorn from starlette.testclient import TestClient +from agentenv_protocol import EnvironmentCard, ToolContext, create_fastmcp_app, tool from agentenv_protocol.a2a_agent import proxy from agentenv_protocol.a2a_agent.proxy import ProxyAgent TOKEN = "t" * 32 -def _send(client: TestClient, text: str, *, context_id: str = "context-1") -> str: +class _Items: + def __init__(self) -> None: + self.items: list[str] = [] + + @tool() + async def add_item(self, item: str) -> str: + """Add an item.""" + self.items.append(item) + return json.dumps({"count": len(self.items)}) + + @tool() + def explode(self) -> str: + """Always fails.""" + raise RuntimeError("the tool blew up") + + @tool() + def whoami(self, ctx: ToolContext) -> str: + """Who is calling.""" + return json.dumps({"role": ctx.caller.role}) + + +@pytest.fixture(scope="module") +def env(): + """An items env's MCP server on loopback, as a task's env serves one.""" + handler = _Items() + app = create_fastmcp_app(handler, card=EnvironmentCard(name="items")).streamable_http_app() + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + server = uvicorn.Server(uvicorn.Config(app, log_level="warning")) + thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}, daemon=True) + thread.start() + while not server.started: + time.sleep(0.01) + yield handler, f"http://127.0.0.1:{sock.getsockname()[1]}/mcp" + server.should_exit = True + thread.join(5) + + +@pytest.fixture +def client(monkeypatch, tmp_path): + monkeypatch.setattr(proxy, "SKILLS_DIR", tmp_path / "skills") + monkeypatch.chdir(tmp_path) + with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 + yield client + + +def _send(client: TestClient, text: str, *, context_id: str = "context-1", parts: list | None = None) -> str: response = client.post("/a2a", json={ "jsonrpc": "2.0", "id": "1", "method": "message/send", "params": { "message": {"kind": "message", "messageId": f"message-{text}", "role": "user", "contextId": context_id, - "parts": [{"kind": "text", "text": text}]}, + "parts": parts or [{"kind": "text", "text": text}]}, "configuration": {"blocking": False}, }, }) @@ -38,21 +90,40 @@ def _settled(client: TestClient, task_id: str) -> dict[str, Any]: return task -def _driver(client: TestClient, operation: str, **body: Any): +def _proxy(client: TestClient, operation: str, **body: Any): return client.post(f"/ext/proxy/{operation}", json={"token": TOKEN, **body}) def _turns(client: TestClient, after: int = 0) -> dict[str, Any]: - response = _driver(client, "next", after=after) + response = _proxy(client, "next", after=after) assert response.status_code == 200, response.text return response.json() -def _ran(client: TestClient, command: str) -> str: - exec_id = _driver(client, "exec", command=command).json()["exec_id"] - while not (result := _driver(client, "output", exec_id=exec_id).json())["done"]: - pass - return result["output"] +def _done(client: TestClient, task_id: str, action_id: str, action: dict[str, Any]) -> dict[str, Any]: + response = _proxy(client, "act", task_id=task_id, action_id=action_id, action=action) + assert response.status_code == 200, response.text + state = response.json() + while not state["done"]: + state = _proxy(client, "result", action_id=action_id).json() + return state["result"] + + +def _bash(command: str, **options: Any) -> dict[str, Any]: + return {"kind": "bash", "command": command, **options} + + +def _register(client: TestClient, url: str, name: str = "items", role: str | None = None) -> None: + headers = {"AgentEnv-Role": role} if role else None + assert client.post("/ext/mcp-config", json={"url": url, "name": name, "headers": headers}).status_code == 200 + + +def _trajectory(client: TestClient, task_id: str) -> list[dict[str, Any]]: + return client.post("/ext/trajectory", json={"task_id": task_id}).json()["trajectory"] + + +def _attrs(span: dict[str, Any], key: str) -> Any: + return json.loads(span["attributes"][key]) def _stopped(pid: int) -> None: @@ -67,35 +138,26 @@ def _state(pid: int) -> str: return subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True).stdout.strip() -@pytest.fixture -def client(monkeypatch, tmp_path): - monkeypatch.setattr(proxy, "SKILLS_DIR", tmp_path / "skills") - monkeypatch.chdir(tmp_path) - with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: - assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 - yield client - - -def test_the_card_advertises_the_proxy_and_takes_a_system_prompt_and_a_role_but_no_mcp_servers(): +def test_the_card_offers_the_proxy_env_servers_and_a_system_prompt(): card = ProxyAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json") extensions = {extension["uri"]: extension for extension in card["capabilities"]["extensions"]} assert sorted(extensions["urn:agentenv:proxy/v1"]["params"]["methods"]) == [ - "claim", "exec", "next", "output", "reply"] - assert {"system_prompt", "role"} <= set( - extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"]["supported"]) - assert "urn:agentenv:mcp-config/v1" not in extensions + "act", "claim", "next", "record", "reply", "result", "tools"] + assert "urn:agentenv:mcp-config/v1" in extensions + assert "system_prompt" in extensions["urn:agentenv:agent-config/v1"]["params"]["methods"]["set"]["request"][ + "supported"] def test_the_first_claim_sets_the_token_that_every_driver_call_needs(): with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: - assert _driver(client, "next").status_code == 401 + assert _proxy(client, "next").status_code == 401 assert client.post("/ext/proxy/claim", json={"token": "short"}).status_code == 400 - assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 - assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 - assert client.post("/ext/proxy/claim", json={"token": "u" * 32}).status_code == 409 - assert client.post("/ext/proxy/next", json={"token": "u" * 32}).status_code == 401 - assert _driver(client, "next").status_code == 200 + assert client.post("/ext/proxy/claim", json={"token": "é" * 22}).status_code == 200 + assert client.post("/ext/proxy/claim", json={"token": "é" * 22}).status_code == 200 + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 409 + assert _proxy(client, "next").status_code == 401 + assert client.post("/ext/proxy/next", json={"token": "é" * 22}).status_code == 200 def test_a_message_is_an_open_turn_until_the_driver_replies(client): @@ -103,39 +165,49 @@ def test_a_message_is_an_open_turn_until_the_driver_replies(client): task_id = _send(client, "fix the build") seen = _turns(client) - assert seen["turns"] == [{"task_id": task_id, "context_id": "context-1", "text": "fix the build", - "system_prompt": "be brief", "skills": []}] - span = {"name": "bash", "attributes": {"gen_ai.operation.name": "execute_tool"}} - assert _driver(client, "reply", task_id=task_id, text="fixed", trajectory=[span], tool_call_count=1).status_code \ - == 200 + assert seen["turns"] == [{"task_id": task_id, "context_id": "context-1", + "parts": [{"kind": "text", "text": "fix the build"}], "system_prompt": "be brief", + "skills": []}] + assert _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "fixed"}]).status_code == 200 task = _settled(client, task_id) assert task["status"]["state"] == "completed" - parts = task["status"]["message"]["parts"] - assert parts[0] == {"kind": "text", "text": "fixed"} - assert parts[1]["data"]["usage"]["tool_call_count"] == 1 - assert client.post("/ext/trajectory", json={"task_id": task_id}).json()["trajectory"] == [span] + assert task["status"]["message"]["parts"][0] == {"kind": "text", "text": "fixed"} + [turn] = _trajectory(client, task_id) + assert (turn["name"], _attrs(turn, "gen_ai.prompt"), _attrs(turn, "gen_ai.completion")) == ( + "turn", {"prompt": "fix the build", "system_prompt": "be brief"}, + {"content": [{"type": "text", "text": "fixed"}]}) assert _turns(client, after=seen["version"])["turns"] == [] +def test_a_message_with_a_part_that_isnt_text_fails_before_the_driver_sees_it(client): + task_id = _send(client, "look", parts=[{"kind": "text", "text": "look"}, {"kind": "data", "data": {"a": 1}}]) + + task = _settled(client, task_id) + assert task["status"]["state"] == "failed" + assert "text parts only" in json.dumps(task["status"]["message"]) + assert _turns(client)["turns"] == [] + + def test_an_error_reply_fails_the_turn(client): task_id = _send(client, "do it") - _driver(client, "reply", task_id=task_id, error="the driver gave up") + _proxy(client, "reply", task_id=task_id, error="the driver gave up") task = _settled(client, task_id) assert task["status"]["state"] == "failed" assert "the driver gave up" in task["status"]["message"]["parts"][0]["text"] -def test_a_reply_names_exactly_one_open_turn(client): +def test_a_reply_names_exactly_one_open_turn_and_carries_parts_or_an_error(client): task_id = _send(client, "do it") - assert _driver(client, "reply", task_id=task_id).status_code == 400 - assert _driver(client, "reply", task_id="nope", text="done").status_code == 404 - _driver(client, "reply", task_id=task_id, text="done") + assert _proxy(client, "reply", task_id=task_id).status_code == 400 + assert _proxy(client, "reply", task_id=task_id, parts=[]).status_code == 400 + assert _proxy(client, "reply", task_id="nope", parts=[{"kind": "text", "text": "done"}]).status_code == 404 + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) _settled(client, task_id) - assert _driver(client, "reply", task_id=task_id, text="again").status_code == 404 + assert _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "again"}]).status_code == 404 def test_a_turn_its_sender_cancels_is_withdrawn(client): @@ -158,41 +230,211 @@ def test_next_waits_for_a_change_after_the_version_it_was_given(client): assert time.monotonic() - started >= 0.2 -def test_exec_returns_stdout_then_stderr_from_the_folder_it_works_in(client, tmp_path): +def test_a_turn_offers_its_env_servers_tools_and_bash(client, env): + _register(client, env[1]) + task_id = _send(client, "go") + + tools = {item["name"]: item for item in _proxy(client, "tools", task_id=task_id).json()["tools"]} + + assert set(tools) == {"add_item", "explode", "whoami", "bash"} + assert (tools["add_item"]["server"], "server" in tools["bash"]) == ("items", False) + assert tools["add_item"]["input_schema"]["properties"]["item"]["type"] == "string" + + +def test_two_env_servers_offering_one_tool_name_are_refused(client, env): + _register(client, env[1], name="items") + _register(client, env[1].replace("127.0.0.1", "localhost"), name="pantry") + task_id = _send(client, "go") + + response = _proxy(client, "tools", task_id=task_id) + + assert response.status_code == 409 + assert "offered by both 'items' and 'pantry'" in response.text + + +def test_a_tool_action_calls_the_env_with_the_headers_registered_for_the_proxy(client, env): + handler, url = env + _register(client, url, role="ops") + task_id = _send(client, "go") + + added = _done(client, task_id, "a1", {"kind": "tool", "server": "items", "name": "add_item", + "arguments": {"item": "crate"}}) + who = _done(client, task_id, "a2", {"kind": "tool", "server": "items", "name": "whoami"}) + failed = _done(client, task_id, "a3", {"kind": "tool", "server": "items", "name": "explode"}) + + assert (added["kind"], added["is_error"], "crate" in handler.items) == ("tool", False, True) + assert json.loads(who["content"][0]["text"]) == {"role": "ops"} + assert failed["is_error"] is True + + +def test_a_bash_action_runs_in_the_proxys_folder_and_returns_a_structured_result(client, tmp_path): (tmp_path / "inputs").mkdir() + task_id = _send(client, "go") - assert _ran(client, "ls; echo out; echo err >&2; exit 3") == "inputs\nout\nerr\n" + result = _done(client, task_id, "b1", _bash("ls; echo out; echo err >&2; exit 3")) + assert result == {"kind": "bash", "exit_code": 3, "stdout": "inputs\nout\n", "stderr": "err\n", "timed_out": False} -def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(client, monkeypatch, tmp_path): - monkeypatch.setattr(proxy, "COMMAND_TIMEOUT_SECONDS", 0.5) - output = _ran(client, "sleep 60 & echo $! > child; wait") +def test_an_action_done_within_the_long_poll_answers_at_once_and_a_slow_one_is_polled(client): + task_id = _send(client, "go") + + quick = _proxy(client, "act", task_id=task_id, action_id="q", action=_bash("echo hi")).json() + slow = _proxy(client, "act", task_id=task_id, action_id="s", action=_bash("sleep 0.5; echo late")).json() + + assert (quick["done"], quick["result"]["stdout"]) == (True, "hi\n") + assert slow == {"done": False} + assert _done(client, task_id, "s", _bash("sleep 0.5; echo late"))["stdout"] == "late\n" + + +def test_an_action_retried_with_its_id_runs_once(client, tmp_path): + task_id = _send(client, "go") + + first = _done(client, task_id, "once", _bash("echo x >> ran; wc -l < ran")) + again = _done(client, task_id, "once", _bash("echo x >> ran; wc -l < ran")) + + assert first == again + assert (tmp_path / "ran").read_text() == "x\n" - assert output == "error: the command timed out after 0.5s" - _stopped(int((tmp_path / "child").read_text())) +def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(client, tmp_path): + task_id = _send(client, "go") -def test_a_child_left_holding_the_output_is_stopped_when_the_command_times_out(client, monkeypatch, tmp_path): - monkeypatch.setattr(proxy, "COMMAND_TIMEOUT_SECONDS", 0.5) + waited = _done(client, task_id, "w", _bash("sleep 60 & echo $! > waited; wait", timeout_seconds=0.5)) + left = _done(client, task_id, "l", _bash("sleep 60 & echo $! > left", timeout_seconds=0.5)) - output = _ran(client, "sleep 60 & echo $! > child") + assert waited == left == {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True} + _stopped(int((tmp_path / "waited").read_text())) + _stopped(int((tmp_path / "left").read_text())) + + +def test_a_turn_that_closes_stops_its_running_actions_and_forgets_them(client, tmp_path): + task_id = _send(client, "go") + assert _proxy(client, "act", task_id=task_id, action_id="long", + action=_bash("sleep 60 & echo $! > child; wait")).json()["done"] is False + + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) + _settled(client, task_id) - assert output == "error: the command timed out after 0.5s" _stopped(int((tmp_path / "child").read_text())) + assert _proxy(client, "result", action_id="long").status_code == 404 + assert _proxy(client, "act", task_id=task_id, action_id="new", action=_bash("true")).status_code == 404 + + +def test_a_proxy_without_a_shell_offers_no_bash_and_keeps_no_skills(): + with TestClient(ProxyAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: + client.post("/ext/proxy/claim", json={"token": TOKEN}) + task_id = _send(client, "go") + + assert _proxy(client, "tools", task_id=task_id).json()["tools"] == [] + assert _proxy(client, "act", task_id=task_id, action_id="b", action=_bash("true")).status_code == 409 + assert client.post("/ext/skill-config", json={"name": "x", "description": "d", "skill_md": "#"}).status_code \ + == 409 + + +def test_recorded_model_calls_are_checked_against_the_actions_that_run(client, env): + _register(client, env[1]) + task_id = _send(client, "go") + message = {"role": "assistant", "content": "Adding it.", "reasoning_content": "The user wants a crate.", + "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "add_item", "arguments": '{"item": "crate"}'}}, + {"id": "c2", "type": "function", "function": {"name": "bash", "arguments": '{"command": "true"}'}}, + {"id": "c3", "type": "function", "function": {"name": "my_own_tool", "arguments": "{}"}}]} + + assert _proxy(client, "record", task_id=task_id, message=message, model="qwen", + usage={"prompt_tokens": 12, "completion_tokens": 7}).status_code == 200 + unmade = _proxy(client, "act", task_id=task_id, action_id="x", + action={"kind": "tool", "server": "items", "name": "add_item", "arguments": {"item": "box"}}) + assert unmade.status_code == 409 + assert _proxy(client, "record", task_id=task_id, message="Done.").status_code == 409 + assert _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "x"}]).status_code == 409 + _done(client, task_id, "b", _bash("true")) + _done(client, task_id, "a", {"kind": "tool", "server": "items", "name": "add_item", "arguments": {"item": "crate"}}) + assert _proxy(client, "record", task_id=task_id, message="Done.").status_code == 200 + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "Done."}]) + task = _settled(client, task_id) + spans = _trajectory(client, task_id) + assert [(span["name"], span["attributes"]["gen_ai.operation.name"]) for span in spans] == [ + ("chat", "chat"), ("bash", "execute_tool"), ("add_item", "execute_tool"), ("chat", "chat"), ("turn", "chain")] + assert _attrs(spans[0], "gen_ai.completion")["content"] == [ + {"type": "thinking", "thinking": "The user wants a crate."}, {"type": "text", "text": "Adding it."}, + {"type": "tool_use", "id": "c1", "name": "add_item", "input": {"item": "crate"}}, + {"type": "tool_use", "id": "c2", "name": "bash", "input": {"command": "true"}}, + {"type": "tool_use", "id": "c3", "name": "my_own_tool", "input": {}}] + assert (spans[0]["attributes"]["gen_ai.request.model"], spans[0]["attributes"]["gen_ai.usage.input_tokens"], + spans[0]["attributes"]["gen_ai.usage.output_tokens"]) == ("qwen", 12, 7) + assert _attrs(spans[2], "gen_ai.prompt") == {"input": {"item": "crate"}} + assert task["status"]["message"]["parts"][1]["data"]["usage"]["tool_call_count"] == 2 + + +@pytest.mark.parametrize(("message", "content"), [ + ("Plain words.", [{"type": "text", "text": "Plain words."}]), + ({"role": "assistant", "content": [{"type": "thinking", "thinking": "Hmm."}, {"type": "text", "text": "Hi."}, + {"type": "tool_use", "id": "t1", "name": "bash", "input": {"command": "ls"}}]}, + [{"type": "thinking", "thinking": "Hmm."}, {"type": "text", "text": "Hi."}, + {"type": "tool_use", "id": "t1", "name": "bash", "input": {"command": "ls"}}]), + ({"role": "assistant", "content": [{"type": "text", "text": "Parts."}]}, [{"type": "text", "text": "Parts."}]), +]) +def test_a_model_call_is_recorded_from_an_openai_or_anthropic_message_or_a_string(client, message, content): + task_id = _send(client, "go") + + assert _proxy(client, "record", task_id=task_id, message=message).status_code == 200 + + tools = [block for block in content if block["type"] == "tool_use"] + if tools: + _done(client, task_id, "b", _bash("ls")) + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "ok"}]) + _settled(client, task_id) + assert _attrs(_trajectory(client, task_id)[0], "gen_ai.completion") == {"content": content} -def test_output_names_a_command_it_started(client): - assert _driver(client, "output", exec_id="nope").status_code == 404 +@pytest.mark.parametrize("message", [ + {"role": "user", "content": "Not the model's."}, + {"role": "assistant", "content": {"text": "a dict"}}, + {"role": "assistant", "content": [{"type": "image", "source": {}}]}, + {"role": "assistant", "content": None, "tool_calls": [{"function": {"name": "bash", "arguments": "{broken"}}]}, +]) +def test_a_model_call_in_any_other_shape_is_refused(client, message): + task_id = _send(client, "go") -def test_a_skill_is_written_where_its_turns_say_it_is(client, tmp_path): - response = client.post("/ext/skill-config", json={"name": "deploy", "description": "How to deploy.", - "skill_md": "# Deploy\n"}) - assert response.status_code == 200, response.text + assert _proxy(client, "record", task_id=task_id, message=message).status_code == 400 + + +def test_a_skill_is_written_where_its_turns_say_it_is_and_only_once(client, tmp_path): + skill = {"name": "deploy", "description": "How to deploy.", "skill_md": "# Deploy\n"} + assert client.post("/ext/skill-config", json=skill).status_code == 200 + assert client.post("/ext/skill-config", json=skill).status_code == 409 _send(client, "go") - skill = _turns(client)["turns"][0]["skills"] - assert skill == [{"name": "deploy", "description": "How to deploy.", "path": str(tmp_path / "skills" / "deploy")}] + assert _turns(client)["turns"][0]["skills"] == [ + {"name": "deploy", "description": "How to deploy.", "path": str(tmp_path / "skills" / "deploy")}] assert (tmp_path / "skills" / "deploy" / "SKILL.md").read_text() == "# Deploy\n" + +def test_a_skill_whose_files_fail_to_download_leaves_no_folder(client, monkeypatch, tmp_path): + async def download(source, target): + if target.name == "notes.md": + raise OSError("the grant expired") + target.write_text("# Review\n") + + monkeypatch.setattr(proxy, "download", download) + read = {"media_type": "text/markdown", "max_bytes": 1024, + "read": {"kind": "http-get", "url": "https://objects.example.test/r", "expires_at": "2099-01-01T00:00:00Z"}} + bundle = {"name": "review", "description": "Review code", "skill_bundle": {"max_total_bytes": 2048, "files": [ + {"path": "SKILL.md", "object": read}, {"path": "notes.md", "object": read}]}} + + assert client.post("/ext/skill-config", json=bundle).status_code >= 400 + assert not (tmp_path / "skills" / "review").exists() + + +def test_what_a_turn_still_runs_is_stopped_when_the_app_stops(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: + client.post("/ext/proxy/claim", json={"token": TOKEN}) + task_id = _send(client, "go") + _proxy(client, "act", task_id=task_id, action_id="long", action=_bash("sleep 60 & echo $! > child; wait")) + while not (tmp_path / "child").exists(): + time.sleep(0.02) + + _stopped(int((tmp_path / "child").read_text())) From a52d7a2daa2c0a5d68f40a5ea9e52d1d0d5014de Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 10:07:20 -0700 Subject: [PATCH 13/39] feat(task): play an outside agent through a ProxyAgent in its place The episode deploys the agents bundle's proxy (agents/proxy, published as agentenv-proxy) where the played agent would run, or serves one in process for an agent that works in no sandbox, and claims it through proxy/v1. Every action the driver's model takes, env tool calls through the servers registered with mcp-config and bash, runs in the proxy, which records the turn: a chat span per recorded model call, an execute_tool span per action and a chain span per turn. Turn gains record_model_call (OpenAI chat messages, Anthropic messages or a string); a call the last recorded one didn't make, or an end while its calls haven't run, raises EpisodeError. Codex declares its place on a VM host. Co-Authored-By: Claude Opus 5.5 --- .github/dependabot.yml | 1 + .github/scripts/infra_build.py | 6 +- AGENTS.md | 13 +- .../agents/agents/claude-code/a2a_server.py | 2 +- .../agents/agents/codex/a2a_server.py | 3 + .../agents/{stand-in => proxy}/Dockerfile | 4 +- .../agents/{stand-in => proxy}/a2a_server.py | 13 +- .../agents/{stand-in => proxy}/agent.toml | 0 .../agents/{stand-in => proxy}/install.sh | 6 +- .../agents/agents/proxy/requirements.txt | 3 + .../agents/agents/stand-in/requirements.txt | 3 - src/agent_env/task/episode/proxy.py | 125 ++++++++++ src/agent_env/task/episode/rules.py | 10 +- src/agent_env/task/episode/run.py | 233 ++++++++---------- src/agent_env/task/episode/spec.py | 6 +- src/agent_env/task/episode/stand_in.py | 89 ------- src/agent_env/task/episode/steps.py | 26 +- .../task_step/task_steps/install_agent.py | 14 +- ...odal_vm_test.py => proxy_modal_vm_test.py} | 10 +- .../env/envs/server_env_local_test.py | 76 ++++-- .../env/envs/server_env_modal_test.py | 1 - tst/unit/agents/codex/test_card_contract.py | 2 + .../agents/{stand_in => proxy}/__init__.py | 0 tst/unit/agents/{stand_in => proxy}/agent.py | 6 +- .../test_proxy_agent.py} | 16 +- tst/unit/task/episode/conftest.py | 2 +- tst/unit/task/episode/run_test.py | 215 +++++++++------- tst/unit/task/episode/spec_test.py | 12 +- tst/unit/task_step/install_agent_test.py | 22 +- tst/unit/test_infra_images_script.py | 3 +- 30 files changed, 500 insertions(+), 422 deletions(-) rename src/agent_env/agents/agents/{stand-in => proxy}/Dockerfile (86%) rename src/agent_env/agents/agents/{stand-in => proxy}/a2a_server.py (67%) rename src/agent_env/agents/agents/{stand-in => proxy}/agent.toml (100%) rename src/agent_env/agents/agents/{stand-in => proxy}/install.sh (80%) create mode 100644 src/agent_env/agents/agents/proxy/requirements.txt delete mode 100644 src/agent_env/agents/agents/stand-in/requirements.txt create mode 100644 src/agent_env/task/episode/proxy.py delete mode 100644 src/agent_env/task/episode/stand_in.py rename tst/integration/agent/{stand_in_modal_vm_test.py => proxy_modal_vm_test.py} (93%) rename tst/unit/agents/{stand_in => proxy}/__init__.py (100%) rename tst/unit/agents/{stand_in => proxy}/agent.py (50%) rename tst/unit/agents/{stand_in/test_stand_in_agent.py => proxy/test_proxy_agent.py} (80%) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index d9a785ff0..37f5133e4 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -44,6 +44,7 @@ updates: - "/src/agent_env/agents/agents/codex" - "/src/agent_env/agents/agents/openhands" - "/src/agent_env/agents/agents/pi" + - "/src/agent_env/agents/agents/proxy" schedule: interval: "weekly" open-pull-requests-limit: 5 diff --git a/.github/scripts/infra_build.py b/.github/scripts/infra_build.py index 5d1e7ecf6..261b32e3e 100644 --- a/.github/scripts/infra_build.py +++ b/.github/scripts/infra_build.py @@ -35,6 +35,7 @@ "agentenv-codex": "src/agent_env/agents/agents/codex/Dockerfile", "agentenv-openhands": "src/agent_env/agents/agents/openhands/Dockerfile", "agentenv-pi": "src/agent_env/agents/agents/pi/Dockerfile", + "agentenv-proxy": "src/agent_env/agents/agents/proxy/Dockerfile", } _TAG = re.compile(r"^src-[0-9a-f]{64}$") _DIGEST = re.compile(r"^sha256:[0-9a-f]{64}$") @@ -52,8 +53,8 @@ _RETRY_WAITS_SECONDS = (30, 60, 120) # Each image's check, beyond its build: the gateway imports its server, Postgres turns healthy, the sidecars and the -# browser answer their own commands, Claude Code's, Codex's and pi's agents run their CLIs and import their servers, and -# OpenHands' imports its server, which loads the OpenHands SDK. +# browser answer their own commands, Claude Code's, Codex's and pi's agents run their CLIs and import their servers, +# OpenHands' imports its server, which loads the OpenHands SDK, and the proxy imports its server. _SMOKES = { "agentenv-gateway": ["docker", "run", "--rm", "--entrypoint", "python", "{image}", "-c", "import gateway.entrypoint"], "agentenv-db-web": ["docker", "run", "--rm", "--entrypoint", "/usr/bin/pgweb", "{image}", "--version"], @@ -68,6 +69,7 @@ "agentenv-openhands": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'python -c "import a2a_server"'], "agentenv-pi": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'pi --version && python3 -c "import a2a_server"'], + "agentenv-proxy": ["docker", "run", "--rm", "--entrypoint", "sh", "{image}", "-c", 'python -c "import a2a_server"'], } diff --git a/AGENTS.md b/AGENTS.md index 666967bb0..19cf1acce 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -211,15 +211,16 @@ orchestrators call it the same way. What callers rely on: `context.metadata[HEARTBEAT_KEY]` (`task_step/context.py`); a step that blocks for long calls `context.heartbeat()`, a no-op without one. The key is in `_REDACTED_KEYS`, so it is never persisted. - An episode (`EpisodeSpec`, `task/episode/`) is one `Task.run()` of a copy of the task, with - namespace routing on, in which each of its `trainer_agents` is played through a stand-in agent - (the `agents` bundle's `stand-in`, `@local/agentenv-framework/agents/stand-in`, which the episode writes + namespace routing on, in which each of its `trainer_agents` is played through a proxy agent + (the `agents` bundle's `proxy`, `@local/agentenv-framework/agents/proxy`, which the episode writes from its folder as a bundle run would). - The agent's `deploy_agent` or `install_agent`, wrapped as `episode_agent`, deploys the stand-in in the + The agent's `deploy_agent` or `install_agent`, wrapped as `episode_agent`, deploys the proxy in the agent's place when the task uses the agent's sandbox and serves it in this process when it doesn't; the agent's `prompt_agent` steps message it as they would the agent, and the caller plays each - message as a `Turn` through its `urn:agentenv:driver/v1`, where a turn's `bash` runs. Its record is - tagged `context.metadata["episode"]` and names the original task, so it isn't resumable with - `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded + message as a `Turn` through its `urn:agentenv:proxy/v1`: every env tool call and `bash` command runs + in the proxy, which records the turn and checks it against the model calls the caller records. Its + record is tagged `context.metadata["episode"]` and names the original task, so it isn't resumable + with `start_step`. The episode tears down what its run deployed, and an abandoned one is recorded `cancelled`. ## Tests 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 fc9fe3914..e241e44a2 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -413,7 +413,7 @@ async def _turns_stop_with_the_server(app: Starlette) -> AsyncIterator[None]: required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], - # Where an install on the host works, and as whom, for a stand-in installed in this agent's place. + # Where an install on the host works, and as whom, for a proxy installed in this agent's place. workdir_host="/app", user_host="root", ), diff --git a/src/agent_env/agents/agents/codex/a2a_server.py b/src/agent_env/agents/agents/codex/a2a_server.py index dc3701f73..a91a89711 100644 --- a/src/agent_env/agents/agents/codex/a2a_server.py +++ b/src/agent_env/agents/agents/codex/a2a_server.py @@ -609,6 +609,9 @@ async def _stage_granted_workspace(source: ReadObject | None) -> Path | None: required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], + # Where an install on the host works, and as whom, for a proxy installed in this agent's place. + workdir_host="/app", + user_host="root", ), enable( PEER_AGENTS_V1, diff --git a/src/agent_env/agents/agents/stand-in/Dockerfile b/src/agent_env/agents/agents/proxy/Dockerfile similarity index 86% rename from src/agent_env/agents/agents/stand-in/Dockerfile rename to src/agent_env/agents/agents/proxy/Dockerfile index 4e088a623..a10648162 100644 --- a/src/agent_env/agents/agents/stand-in/Dockerfile +++ b/src/agent_env/agents/agents/proxy/Dockerfile @@ -9,9 +9,9 @@ RUN pip install --no-cache-dir -r requirements.txt # install.sh is install/v1's: in the image, so it stays in its build context. COPY a2a_server.py install.sh ./ RUN useradd --create-home agent && mkdir -p /app && chown agent:agent /app -ENV A2A_PORT=8000 +ENV A2A_PORT=8000 PYTHONPATH=/opt/agent USER agent -# The commands the driver runs work where the stand-in does: in /app, as agent. +# The actions the driver has it run work where the proxy does: in /app, as agent. WORKDIR /app EXPOSE 8000 CMD ["python", "/opt/agent/a2a_server.py"] diff --git a/src/agent_env/agents/agents/stand-in/a2a_server.py b/src/agent_env/agents/agents/proxy/a2a_server.py similarity index 67% rename from src/agent_env/agents/agents/stand-in/a2a_server.py rename to src/agent_env/agents/agents/proxy/a2a_server.py index 3b7c29706..ba576c78f 100644 --- a/src/agent_env/agents/agents/stand-in/a2a_server.py +++ b/src/agent_env/agents/agents/proxy/a2a_server.py @@ -1,8 +1,8 @@ -"""Serves the stand-in agent. ``AGENT_CARD`` is what agent-env's install/v1 reads. ``install_agent`` formats each -command with ``str.format``.""" +"""Serves the proxy agent. ``AGENT_CARD`` is what agent-env's install/v1 reads. ``install_agent`` formats each command +with ``str.format``.""" from agentenv_protocol.a2a_agent import INSTALL_V1, a2a_agent, enable -from agentenv_protocol.a2a_agent.stand_in import EXTENSIONS, IDENTITY, StandInAgent, StandInConfig +from agentenv_protocol.a2a_agent.proxy import EXTENSIONS, IDENTITY, ProxyAgent, ProxyConfig, stopping_actions _UNPACK = "mkdir -p {work_dir}/agent && tar -xzf {agent_ctx_tar} -C {work_dir}/agent" # In the place of the agent it replaces: its {workdir}, as its {user}. @@ -11,19 +11,20 @@ @a2a_agent( identity=IDENTITY, - config=StandInConfig, + config=ProxyConfig, extensions=( *EXTENSIONS, enable(INSTALL_V1, install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], a2a_port=8000), ), + lifespan=stopping_actions, ) -class InstalledStandIn(StandInAgent): +class InstalledProxy(ProxyAgent): pass -agent = InstalledStandIn() +agent = InstalledProxy() app = agent.create_app() AGENT_CARD = app.state.agentenv_a2a.card diff --git a/src/agent_env/agents/agents/stand-in/agent.toml b/src/agent_env/agents/agents/proxy/agent.toml similarity index 100% rename from src/agent_env/agents/agents/stand-in/agent.toml rename to src/agent_env/agents/agents/proxy/agent.toml diff --git a/src/agent_env/agents/agents/stand-in/install.sh b/src/agent_env/agents/agents/proxy/install.sh similarity index 80% rename from src/agent_env/agents/agents/stand-in/install.sh rename to src/agent_env/agents/agents/proxy/install.sh index 7e152dde4..13eafe171 100644 --- a/src/agent_env/agents/agents/stand-in/install.sh +++ b/src/agent_env/agents/agents/proxy/install.sh @@ -1,5 +1,5 @@ #!/bin/sh -# Installs the stand-in onto this Linux host from the folder this script is in, and starts it on PORT, as USER in +# Installs the proxy onto this Linux host from the folder this script is in, and starts it on PORT, as USER in # WORKDIR, the place of the agent it replaces: # # install.sh NAME PORT WORKDIR USER @@ -8,7 +8,7 @@ set -eu name=$1 port=$2 workdir=$3 user=$4 here=$(cd "$(dirname "$0")" && pwd) -root=/opt/agentenv-stand-in-$name +root=/opt/agentenv-proxy-$name if ! python3 -c "import ensurepip, venv" 2>/dev/null; then export DEBIAN_FRONTEND=noninteractive @@ -24,4 +24,4 @@ id -u "$user" >/dev/null 2>&1 || useradd -m "$user" mkdir -p "$workdir" chown -R "$user": "$workdir" nohup runuser -u "$user" -- env -C "$workdir" A2A_PORT="$port" "$root/venv/bin/python" "$root/a2a_server.py" \ - >"/tmp/agentenv-stand-in-$name.log" 2>&1 "/tmp/agentenv-proxy-$name.log" 2>&1 =0.1.367 diff --git a/src/agent_env/agents/agents/stand-in/requirements.txt b/src/agent_env/agents/agents/stand-in/requirements.txt deleted file mode 100644 index cddf65a18..000000000 --- a/src/agent_env/agents/agents/stand-in/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -# A release pins agentenv-framework-protocol to the version it was released with. The floor is the first that ships -# StandInAgent and driver/v1. -agentenv-framework-protocol[agent]>=0.1.356 diff --git a/src/agent_env/task/episode/proxy.py b/src/agent_env/task/episode/proxy.py new file mode 100644 index 000000000..df3612eca --- /dev/null +++ b/src/agent_env/task/episode/proxy.py @@ -0,0 +1,125 @@ +"""The episode's end of a proxy agent: the driver calls of ``urn:agentenv:proxy/v1``, and a proxy served in this process +for an outside agent that works in no sandbox.""" + +from __future__ import annotations + +import asyncio +import secrets +import threading +from typing import Any + +import httpx +import uvicorn +from agentenv_protocol.a2a_agent import PROXY_V1 +from agentenv_protocol.a2a_agent.proxy import ProxyAgent + +from agent_env.a2a_agent.a2a_agent import A2AAgent + +_CALL_TIMEOUT = httpx.Timeout(60.0, connect=10.0) +_SERVE_TIMEOUT_SECONDS = 30.0 +# Short, so a long-poll in flight when the episode ends finishes within the server's shutdown. +_IN_PROCESS_LONG_POLL_SECONDS = 1.0 +_ATTEMPTS = 3 +_RETRY_SECONDS = 1.0 + + +class ProxyRefused(Exception): + """The proxy refused a call, a 400 or 409, saying why.""" + + +class ProxyClient: + """One proxy, claimed with a token of this episode's, at the endpoints its card names.""" + + def __init__(self, a2a_url: str, card: dict[str, Any]): + proxy = A2AAgent.find_extension(card, PROXY_V1.uri) + if proxy is None: + raise RuntimeError(f"the agent at {a2a_url} isn't a proxy: its card has no {PROXY_V1.uri}") + self._endpoints = {operation: a2a_url + A2AAgent.operation(proxy, operation)[1] + for operation in PROXY_V1.operations} + self._token = secrets.token_urlsafe(24) + self._client = httpx.AsyncClient(timeout=_CALL_TIMEOUT) + + async def claim(self) -> None: + await self._call("claim") + + async def next(self, after: int) -> dict[str, Any]: + return await self._call("next", after=after) + + async def tools(self, task_id: str) -> list[dict[str, Any]]: + return (await self._call("tools", task_id=task_id))["tools"] + + async def act(self, task_id: str, action: dict[str, Any]) -> dict[str, Any]: + """``action``'s result, once it has run in the turn. A call that didn't reach the proxy is tried again: the + action's id makes a retry return the action it already started.""" + action_id = secrets.token_hex(16) + state = await self._retried("act", task_id=task_id, action_id=action_id, action=action) + while not state["done"]: + state = await self._retried("result", action_id=action_id) + return state["result"] + + async def record(self, task_id: str, message: str | dict[str, Any], model: str | None, + usage: dict[str, Any] | None) -> None: + await self._call("record", task_id=task_id, message=message, model=model, usage=usage) + + async def reply(self, task_id: str, **reply: Any) -> bool: + """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone.""" + response = await self._client.post(self._endpoints["reply"], json={"token": self._token, "task_id": task_id, + **reply}) + if response.status_code == 404: + return False + _raise_for_status(response) + return True + + async def close(self) -> None: + await self._client.aclose() + + async def _retried(self, operation: str, **body: Any) -> dict[str, Any]: + for attempt in range(1, _ATTEMPTS + 1): + try: + return await self._call(operation, **body) + except httpx.TransportError: + if attempt == _ATTEMPTS: + raise + await asyncio.sleep(_RETRY_SECONDS) + raise AssertionError("unreachable") + + async def _call(self, operation: str, **body: Any) -> dict[str, Any]: + response = await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) + _raise_for_status(response) + return response.json() or {} + + +class InProcessProxy: + """A proxy served on loopback in a thread of its own, with no shell: one for an outside agent that works in no + sandbox, which only this process's steps message.""" + + def __init__(self) -> None: + agent = ProxyAgent(shell=False, long_poll_seconds=_IN_PROCESS_LONG_POLL_SECONDS) + self._server = uvicorn.Server(uvicorn.Config(agent.create_app(), host="127.0.0.1", port=0, log_level="warning", + timeout_graceful_shutdown=2 * _IN_PROCESS_LONG_POLL_SECONDS)) + self._thread = threading.Thread(target=self._server.run, name="agentenv-proxy", daemon=True) + + async def start(self) -> str: + """Serve it, and return its URL once it answers.""" + self._thread.start() + async with asyncio.timeout(_SERVE_TIMEOUT_SECONDS): + while not self._server.started: + if not self._thread.is_alive(): + raise RuntimeError("the in-process proxy stopped as it started") + await asyncio.sleep(0.01) + port = self._server.servers[0].sockets[0].getsockname()[1] + return f"http://127.0.0.1:{port}" + + def stop(self) -> None: + """Have it stop; its thread ends once the requests in flight have.""" + self._server.should_exit = True + + +def _raise_for_status(response: httpx.Response) -> None: + if response.status_code in (400, 409): + try: + detail = response.json().get("detail") + except ValueError: + detail = response.text + raise ProxyRefused(str(detail)) + response.raise_for_status() diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index e2a4819ca..cc6bfaf3f 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -1,7 +1,7 @@ """Which steps a task may have when some of its agents are played from outside AgentEnv, and why the others are refused. -An outside agent is played through a stand-in agent deployed in its place (the ``agents`` bundle's ``stand-in``), so -steps that act on the agent, or in its sandbox, act on the stand-in. Every registered step type is either checked +An outside agent is played through a proxy agent deployed in its place (the ``agents`` bundle's ``proxy``), so steps +that act on the agent, or in its sandbox, act on the proxy. Every registered step type is either checked here (``CHECKS``) or listed as untouched by an outside agent (``UNAFFECTED``). A test holds the registry to the two, so a new step type is classified before an episode runs it. Plugin step types aren't inspected. """ @@ -51,7 +51,7 @@ def problems(steps: list[TaskStep], outside: list[str]) -> list[str]: def works_in_sandbox(steps: list[TaskStep], agent: str) -> bool: - """Whether ``steps`` use ``agent``'s sandbox, so its stand-in is deployed there rather than served in this + """Whether ``steps`` use ``agent``'s sandbox, so its proxy is deployed there rather than served in this process.""" return any(step.agent_name == agent and ( step.type in ("install_agent", "load_artifact", "run_code", "add_skills") @@ -92,7 +92,7 @@ def _deploy_agent(step: DeployAgentTaskStep, outside: Outside) -> list[str]: def _install_agent(step: InstallAgentTaskStep, outside: Outside) -> list[str]: if step.agent_name not in outside.agents or step.container_name is None: return [] - return [f"an outside agent's stand-in replaces an agent installed on the VM host, not in container " + return [f"an outside agent's proxy replaces an agent installed on the VM host, not in container " f"{step.container_name!r}"] @@ -166,7 +166,7 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s "agent_prompt_response_verifier", "verify_mcp_env_assessment", # It skips an agent that isn't deployed. "teardown_sandboxes", - # They act on an outside agent's stand-in, in its sandbox, as on the agent. + # They act on an outside agent's proxy, in its sandbox, as on the agent. "add_skills", "load_artifact", "run_code", "verify_sandbox", }) diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index 050df47ba..f978e8d6e 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -1,23 +1,23 @@ from __future__ import annotations import asyncio -import json import logging from collections import deque from dataclasses import dataclass -from datetime import datetime, timezone from typing import Any import httpx -from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession +from agentenv_protocol import EnvironmentTool, ToolResult +from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url from agent_env.store.routing import namespace_routing -from agent_env.task.episode.stand_in import InProcessStandIn, StandInClient +from agent_env.task.episode.proxy import InProcessProxy, ProxyClient, ProxyRefused from agent_env.task.store import record_task_cancelled from agent_env.task.task import Task, recorded_scores, substitute_seed_in_steps from agent_env.task.teardown import teardown_run from agent_env.task_step.context import DeployedAgent, TaskStepContext +from agent_env.task_step.task_steps.deploy_agent import _mcp_add_body logger = logging.getLogger(__name__) @@ -25,13 +25,6 @@ _DRAIN_RETRY_SECONDS = 1.0 _CARD_TIMEOUT_SECONDS = 10.0 -BASH = EnvironmentTool( - name="bash", - description="Run a bash command where the agent works, as its user. Returns its stdout, then its stderr.", - inputSchema={"type": "object", "properties": {"command": {"type": "string", "description": "The command to run."}}, - "required": ["command"]}, -) - class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" @@ -43,130 +36,118 @@ class Skill: name: str description: str - path: str | None + path: str class Turn: - """An outside agent's turn. Its model reads ``system_prompt``, ``prompt`` and ``skills``, acts through ``tools()`` - (or an MCP client of its own on ``mcp_urls``, sending ``headers`` and ``role``), and hands its reply back with - ``end()``.""" + """An outside agent's turn. Its model reads ``system_prompt``, ``prompt`` and ``skills``, acts through ``tools()``, + each of its calls recorded with ``record_model_call()``, and hands its reply back with ``end()``. Every action runs + in the agent's proxy, which records the turn.""" def __init__( - self, *, agent: str, step_id: str, prompt: str, system_prompt: str | None, role: str | None, - mcp_urls: dict[str, str], skills: list[Skill], shell: bool, stand_in: StandInClient, task_id: str, + self, *, agent: str, step_id: str, prompt: str, system_prompt: str | None, skills: list[Skill], + proxy: ProxyClient, task_id: str, ): self.agent = agent self.step_id = step_id self.prompt = prompt self.system_prompt = system_prompt - self.role = role - self.mcp_urls = mcp_urls - self.headers = {env_id: sandbox_request_headers_for_url(url) for env_id, url in mcp_urls.items()} self.skills = skills - self._shell = shell - self._stand_in = stand_in + self._proxy = proxy self._task_id = task_id - self._spans: list[dict[str, Any]] = [] - self._started = _now() self._open = True def tools(self, *, timeout: float = 600.0) -> TurnTools: - """A client for the tools of the envs this agent acts in, calling them as its role with the headers each URL - needs, and ``bash`` when it works in a sandbox; use it with ``async with``. ``timeout`` bounds each env tool's + """A client for what this turn can call: the tools of the envs its agent acts in, and ``bash`` when it works + in a sandbox, each run where the agent works; use it with ``async with``. ``timeout`` bounds each env tool's call, in seconds.""" return TurnTools(self, timeout) + async def record_model_call(self, message: str | dict[str, Any], *, model: str | None = None, + usage: dict[str, Any] | None = None) -> None: + """Record one of the model's calls, once it has returned: an OpenAI chat message, an Anthropic message or a + string. From then on each call to the turn's tools must be one the last recorded call made, and all must run + before the next is recorded or the turn ends; else ``EpisodeError``. Calls to tools the turn doesn't offer + aren't checked.""" + await self._checked(self._proxy.record(self._task_id, message, model, usage)) + async def end(self, response: str | None = None) -> None: """Hand the turn back, and the run carries on. No reply records an empty one, as an agent that answers with - nothing leaves. Once the turn is over, ended or not, this does nothing; a stand-in that can't be reached - fails the run, which ``next_turn()`` and ``result()`` report.""" + nothing leaves. Once the turn is over, ended or not, this does nothing. Raises ``EpisodeError``, leaving the + turn open, while calls a recorded model call made haven't run; a proxy that can't be reached fails the run, + which ``next_turn()`` and ``result()`` report.""" if not self._open: return + try: + await self._checked(self._proxy.reply(self._task_id, parts=[{"kind": "text", "text": response or ""}])) + except httpx.HTTPError as e: + logger.warning("turn %r of agent %r: its reply didn't reach the proxy: %s", self.step_id, self.agent, e) self._open = False - reply = response or "" - spans = [*self._spans, _span("turn", "chain", self._started, completion={"content": [ - {"type": "text", "text": reply}]})] - await self._reply(text=reply, trajectory=spans, tool_call_count=len(self._spans)) async def _fail(self, error: Exception) -> None: if self._open: self._open = False - await self._reply(error=str(error)) + try: + await self._proxy.reply(self._task_id, error=str(error)) + except httpx.HTTPError as e: + logger.warning("turn %r of agent %r: its failure didn't reach the proxy: %s", self.step_id, self.agent, + e) - async def _reply(self, **reply: Any) -> None: + async def _checked(self, call: Any) -> Any: + """``call``'s answer, with the proxy's refusal, or a turn it no longer has, as an ``EpisodeError``.""" + if not self._open: + call.close() + raise EpisodeError(f"turn {self.step_id!r} of agent {self.agent!r} is over") try: - await self._stand_in.reply(self._task_id, **reply) - except httpx.HTTPError as e: - logger.warning("turn %r of agent %r: its reply didn't reach the stand-in: %s", self.step_id, self.agent, e) + return await call + except ProxyRefused as e: + raise EpisodeError(f"turn {self.step_id!r} of agent {self.agent!r}: {e}") from None + except httpx.HTTPStatusError as e: + if e.response.status_code == 404: + self._open = False + raise EpisodeError(f"turn {self.step_id!r} of agent {self.agent!r} is over") from None + raise class TurnTools: - """The tools of a turn, listed once when it opens: its envs', and ``bash`` when its agent works in a sandbox. - ``list()`` names every one, ``call()`` routes a call to what offers it, and each call goes in the turn's - trajectory.""" + """What a turn can call, listed once when it opens: its envs' tools, and ``bash`` when its agent works in a + sandbox. ``list()`` names every one, and ``call()`` has the agent's proxy run a call where the agent works.""" def __init__(self, turn: Turn, timeout: float): self._turn = turn - self._sessions = { - env_id: ToolSession(mcp_url=url, role=turn.role, headers=turn.headers[env_id], timeout=timeout) - for env_id, url in turn.mcp_urls.items() - } + self._timeout = timeout self._tools: list[EnvironmentTool] = [] - self._offered_by: dict[str, str | None] = {} + self._servers: dict[str, str | None] = {} async def __aenter__(self) -> TurnTools: - try: - for env_id, session in self._sessions.items(): - await session.open() - for tool in await session.list_tools(): - if (other := self._offered_by.get(tool.name)) is not None: - raise EpisodeError(f"tool {tool.name!r} is offered by both env {other!r} and env {env_id!r}") - self._offered_by[tool.name] = env_id - self._tools.append(tool) - if self._turn._shell: - if (env_id := self._offered_by.get(BASH.name)) is not None: - raise EpisodeError(f"tool {BASH.name!r} is offered by env {env_id!r}, and the agent's sandbox " - "offers its own") - self._offered_by[BASH.name] = None - self._tools.append(BASH) - except BaseException: - await self._close() - raise + for tool in await self._turn._checked(self._turn._proxy.tools(self._turn._task_id)): + self._tools.append(EnvironmentTool(name=tool["name"], description=tool.get("description"), + inputSchema=tool.get("input_schema"))) + self._servers[tool["name"]] = tool.get("server") return self async def __aexit__(self, *exc_info: object) -> None: - await self._close() + pass def list(self) -> list[EnvironmentTool]: return list(self._tools) async def call(self, name: str, arguments: dict[str, Any] | None = None) -> ToolResult: - if name not in self._offered_by: - raise EpisodeError( - f"no env of turn {self._turn.step_id!r} offers tool {name!r}; they offer {sorted(self._offered_by)}") - started, env_id = _now(), self._offered_by[name] - try: - if env_id is None: - result = await self._bash(arguments or {}) - else: - result = await self._sessions[env_id].call_tool(name, arguments) - except Exception as e: - self._record(name, arguments, started, {"error": f"{type(e).__name__}: {e}"}) - raise - self._record(name, arguments, started, {"output": result.content}) - return result - - def _record(self, name: str, arguments: dict[str, Any] | None, started: str, completion: dict[str, Any]) -> None: - self._turn._spans.append(_span(name, "execute_tool", started, prompt={"input": arguments or {}}, - completion=completion)) - - async def _bash(self, arguments: dict[str, Any]) -> ToolResult: - if not isinstance(command := arguments.get("command"), str): - raise EpisodeError(f"tool {BASH.name!r} takes a command string, not {command!r}") - return ToolResult(content=[{"type": "text", "text": await self._turn._stand_in.run(command)}]) - - async def _close(self) -> None: - await asyncio.gather(*(session.close() for session in self._sessions.values())) + if name not in self._servers: + raise EpisodeError(f"turn {self._turn.step_id!r} offers no tool {name!r}; it offers {sorted(self._servers)}") + if (server := self._servers[name]) is None: + if not isinstance(command := (arguments or {}).get("command"), str): + raise EpisodeError(f"tool {name!r} takes a command string, not {command!r}") + action = {"kind": "bash", "command": command} + else: + action = {"kind": "tool", "server": server, "name": name, "arguments": arguments or {}, + "timeout_seconds": self._timeout} + result = await self._turn._checked(self._turn._proxy.act(self._turn._task_id, action)) + if result["kind"] == "bash": + timed_out = result["timed_out"] + text = "error: the command timed out" if timed_out else result["stdout"] + result["stderr"] + return ToolResult(content=[{"type": "text", "text": text}], isError=timed_out) + return ToolResult(content=result["content"], isError=result["is_error"]) class Episode: @@ -185,10 +166,10 @@ def __init__(self, id: str, task: Task, context: TaskStepContext): self._changed: asyncio.Future[None] = asyncio.get_running_loop().create_future() self._tearing_down = False self._abandoned: str | None = None - self._stand_ins: list[StandInClient] = [] + self._proxies: list[ProxyClient] = [] self._drains: list[asyncio.Task[None]] = [] self._played: set[str] = set() - self._in_process: list[InProcessStandIn] = [] + self._in_process: list[InProcessProxy] = [] async def next_turn(self) -> Turn | None: """The next outside turn the run reaches, or None once the run has finished. Raises the run's error if it @@ -248,7 +229,7 @@ async def _launch(self) -> None: raise async def _drive(self, running: asyncio.Event) -> None: - """Run the task with namespace routing on, as ``agent-env run`` does, so the stand-in, an ``@local`` agent, + """Run the task with namespace routing on, as ``agent-env run`` does, so the proxy, an ``@local`` agent, resolves.""" running.set() with namespace_routing(): @@ -257,7 +238,7 @@ async def _drive(self, running: asyncio.Event) -> None: finally: self._tearing_down = True try: - await self._release_stand_ins() + await self._release_proxies() report = await teardown_run(self._context) if report.still_up: logger.warning("episode %s: %s may still be up", self.id, @@ -272,7 +253,7 @@ async def _wait(self, changed: asyncio.Future[None]) -> None: await asyncio.wait((self._run, changed), return_when=asyncio.FIRST_COMPLETED) def _close_turns_of_failed_steps(self) -> None: - """A step that failed has given up on its turns, which its stand-in drops only a long-poll later.""" + """A step that failed has given up on its turns, which its proxy drops only a long-poll later.""" failed = {failure.get("step_id") for failure in self._context.metadata.get("failed_steps") or []} for turn in self._turns: if turn.step_id in failed: @@ -285,34 +266,36 @@ def _raise_unless_completed(self) -> None: raise error async def _serve_in_process(self, agent: str, role: str | None, mcp_urls: dict[str, str]) -> DeployedAgent: - """A stand-in served in this process for ``agent``, which works in no sandbox, played as the others are.""" - server = InProcessStandIn() + """A proxy served in this process for ``agent``, which works in no sandbox, with each env server registered as + ``deploy_agent`` registers them with any agent, at the URL this process reaches: played as the others are.""" + server = InProcessProxy() self._in_process.append(server) url = await server.start() async with httpx.AsyncClient(timeout=_CARD_TIMEOUT_SECONDS) as client: response = await client.get(f"{url}/.well-known/agent.json") response.raise_for_status() - deployed = DeployedAgent(agent_name=agent, api_url=url, a2a_url=url, a2a_card=response.json(), role=role) - await self._attach(agent, deployed, role, mcp_urls, shell=False) + card = response.json() + mcp_config = A2AAgent.find_extension(card, A2AAgent.EXT_MCP_CONFIG) + for env_id, mcp_url in mcp_urls.items(): + body = _mcp_add_body(mcp_config, mcp_url, sandbox_request_headers_for_url(mcp_url), env_id, role) + (await client.post(url + mcp_config["params"]["endpoint"], json=body)).raise_for_status() + deployed = DeployedAgent(agent_name=agent, api_url=url, a2a_url=url, a2a_card=card, role=role) + await self._attach(agent, deployed) return deployed - async def _attach( - self, agent: str, deployed: DeployedAgent, role: str | None, mcp_urls: dict[str, str], *, shell: bool, - ) -> None: - """Claim ``agent``'s stand-in and hand out the turns it opens.""" - stand_in = StandInClient(deployed.a2a_url, deployed.a2a_card or {}) - self._stand_ins.append(stand_in) - await stand_in.claim() - self._drains.append(asyncio.create_task(self._drain(agent, stand_in, role, mcp_urls, shell))) - - async def _drain( - self, agent: str, stand_in: StandInClient, role: str | None, mcp_urls: dict[str, str], shell: bool, - ) -> None: - """Hand out each turn the stand-in opens, and close each one it drops, until the run ends.""" + async def _attach(self, agent: str, deployed: DeployedAgent) -> None: + """Claim ``agent``'s proxy and hand out the turns it opens.""" + proxy = ProxyClient(deployed.a2a_url, deployed.a2a_card or {}) + self._proxies.append(proxy) + await proxy.claim() + self._drains.append(asyncio.create_task(self._drain(agent, proxy))) + + async def _drain(self, agent: str, proxy: ProxyClient) -> None: + """Hand out each turn the proxy opens, and close each one it drops, until the run ends.""" after, open_turns = 0, {} while True: try: - answer = await stand_in.next(after) + answer = await proxy.next(after) except httpx.HTTPError as e: logger.warning("episode %s: reading agent %r's turns failed, retrying: %s", self.id, agent, e) await asyncio.sleep(_DRAIN_RETRY_SECONDS) @@ -326,10 +309,10 @@ async def _drain( for task_id, item in listed.items(): if task_id not in open_turns: open_turns[task_id] = turn = Turn( - agent=agent, step_id=self._step_of(agent, item["context_id"]), prompt=item["text"], - system_prompt=item.get("system_prompt"), role=role, mcp_urls=dict(mcp_urls), - skills=[Skill(**skill) for skill in item.get("skills", [])], shell=shell, stand_in=stand_in, - task_id=task_id, + agent=agent, step_id=self._step_of(agent, item["context_id"]), + prompt="\n".join(part["text"] for part in item["parts"]), + system_prompt=item.get("system_prompt"), + skills=[Skill(**skill) for skill in item.get("skills", [])], proxy=proxy, task_id=task_id, ) self._waiting.append(turn) self._turns.add(turn) @@ -350,7 +333,7 @@ def _step_of(self, agent: str, context_id: str) -> str: def _seeded_steps(self) -> list: return substitute_seed_in_steps(self._task.steps, self._context.metadata.get("seed") or {}) - async def _release_stand_ins(self) -> None: + async def _release_proxies(self) -> None: for drain in self._drains: drain.cancel() await asyncio.gather(*self._drains, return_exceptions=True) @@ -358,21 +341,7 @@ async def _release_stand_ins(self) -> None: turn._open = False for server in self._in_process: server.stop() - await asyncio.gather(*(stand_in.close() for stand_in in self._stand_ins), return_exceptions=True) - - -def _span(name: str, operation: str, started: str, *, prompt: Any = None, completion: Any = None) -> dict[str, Any]: - """An OTel GenAI span, the trajectory format rubrics read.""" - attributes = {"gen_ai.operation.name": operation} - if prompt is not None: - attributes["gen_ai.prompt"] = json.dumps(prompt, default=str) - if completion is not None: - attributes["gen_ai.completion"] = json.dumps(completion, default=str) - return {"name": name, "start_time": started, "end_time": _now(), "attributes": attributes} - - -def _now() -> str: - return datetime.now(timezone.utc).isoformat() + await asyncio.gather(*(proxy.close() for proxy in self._proxies), return_exceptions=True) def _episode_of(context: TaskStepContext) -> Episode: diff --git a/src/agent_env/task/episode/spec.py b/src/agent_env/task/episode/spec.py index c8b0c0263..1731edb42 100644 --- a/src/agent_env/task/episode/spec.py +++ b/src/agent_env/task/episode/spec.py @@ -37,7 +37,7 @@ def from_task( if found := problems(seeded, outside): raise EpisodeError(f"task {task.id!r} can't run as an episode: " + "; ".join(found)) sandboxed = {name: works_in_sandbox(seeded, name) for name in outside} - steps = [_stand_in(step, as_run, sandboxed, turn_timeout) for step, as_run in zip(task.steps, seeded)] + steps = [_in_episode(step, as_run, sandboxed, turn_timeout) for step, as_run in zip(task.steps, seeded)] return cls(copy=Task(id=task.id, version=task.version, steps=steps), trainer_agents=tuple(outside), seed=seed) async def start(self) -> Episode: @@ -69,9 +69,9 @@ def _trainer_agents(task: Task, steps: list[TaskStep], named: Sequence[str] | No return outside -def _stand_in(step: TaskStep, as_run: TaskStep, sandboxed: dict[str, bool], turn_timeout: float | None) -> TaskStep: +def _in_episode(step: TaskStep, as_run: TaskStep, sandboxed: dict[str, bool], turn_timeout: float | None) -> TaskStep: """The step the episode's copy of the task runs in place of ``step``, judged by ``as_run``, the step as the - seed makes it; ``sandboxed`` says, for each outside agent by name, whether its stand-in is deployed in its + seed makes it; ``sandboxed`` says, for each outside agent by name, whether its proxy is deployed in its sandbox. The copy keeps every step's own fields, which the run seeds.""" if as_run.type in DEPLOYS and as_run.agent_name in sandboxed: return _EpisodeAgentStep( diff --git a/src/agent_env/task/episode/stand_in.py b/src/agent_env/task/episode/stand_in.py deleted file mode 100644 index dbff4ecbe..000000000 --- a/src/agent_env/task/episode/stand_in.py +++ /dev/null @@ -1,89 +0,0 @@ -"""The episode's end of a stand-in agent: the driver calls of ``urn:agentenv:driver/v1``, and a stand-in served in this -process for an outside agent that works in no sandbox.""" - -from __future__ import annotations - -import asyncio -import secrets -import threading -from typing import Any - -import httpx -import uvicorn -from agentenv_protocol.a2a_agent import DRIVER_V1 -from agentenv_protocol.a2a_agent.stand_in import StandInAgent - -from agent_env.a2a_agent.a2a_agent import A2AAgent - -_CALL_TIMEOUT = httpx.Timeout(60.0, connect=10.0) -_SERVE_TIMEOUT_SECONDS = 30.0 -# Short, so a long-poll in flight when the episode ends finishes within the server's shutdown. -_IN_PROCESS_LONG_POLL_SECONDS = 1.0 - - -class StandInClient: - """One stand-in, claimed with a token of this episode's, at the endpoints its card names.""" - - def __init__(self, a2a_url: str, card: dict[str, Any]): - driver = A2AAgent.find_extension(card, DRIVER_V1.uri) - if driver is None: - raise RuntimeError(f"the agent at {a2a_url} isn't a stand-in: its card has no {DRIVER_V1.uri}") - self._endpoints = {operation: a2a_url + A2AAgent.operation(driver, operation)[1] - for operation in DRIVER_V1.operations} - self._token = secrets.token_urlsafe(24) - self._client = httpx.AsyncClient(timeout=_CALL_TIMEOUT) - - async def claim(self) -> None: - await self._call("claim") - - async def next(self, after: int) -> dict[str, Any]: - return await self._call("next", after=after) - - async def reply(self, task_id: str, **reply: Any) -> bool: - """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone.""" - response = await self._client.post(self._endpoints["reply"], json={"token": self._token, "task_id": task_id, - **reply}) - if response.status_code == 404: - return False - response.raise_for_status() - return True - - async def run(self, command: str) -> str: - exec_id = (await self._call("exec", command=command))["exec_id"] - while not (output := await self._call("output", exec_id=exec_id))["done"]: - pass - return output["output"] - - async def close(self) -> None: - await self._client.aclose() - - async def _call(self, operation: str, **body: Any) -> dict[str, Any]: - response = await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) - response.raise_for_status() - return response.json() - - -class InProcessStandIn: - """A stand-in served on loopback in a thread of its own, with no shell: one for an outside agent that works in no - sandbox, which only this process's steps message.""" - - def __init__(self) -> None: - agent = StandInAgent(shell=False, long_poll_seconds=_IN_PROCESS_LONG_POLL_SECONDS) - self._server = uvicorn.Server(uvicorn.Config(agent.create_app(), host="127.0.0.1", port=0, log_level="warning", - timeout_graceful_shutdown=2 * _IN_PROCESS_LONG_POLL_SECONDS)) - self._thread = threading.Thread(target=self._server.run, name="agentenv-stand-in", daemon=True) - - async def start(self) -> str: - """Serve it, and return its URL once it answers.""" - self._thread.start() - async with asyncio.timeout(_SERVE_TIMEOUT_SECONDS): - while not self._server.started: - if not self._thread.is_alive(): - raise RuntimeError("the in-process stand-in agent stopped as it started") - await asyncio.sleep(0.01) - port = self._server.servers[0].sockets[0].getsockname()[1] - return f"http://127.0.0.1:{port}" - - def stop(self) -> None: - """Have it stop; its thread ends once the requests in flight have.""" - self._server.should_exit = True diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py index 09cb9b1e6..e363175de 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -12,14 +12,14 @@ from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url -# The agents bundle's stand-in, deployed in the place of an agent an episode plays. -STAND_IN_AGENT_ID = "@local/agentenv-framework/agents/stand-in" +# The agents bundle's proxy, deployed in the place of an agent an episode plays. +PROXY_AGENT_ID = "@local/agentenv-framework/agents/proxy" class _EpisodeAgentStep(TaskStep): """An outside agent's ``deploy_agent`` or ``install_agent``. In a task that uses the agent's sandbox it runs that - step with the stand-in agent in the agent's place; otherwise it serves a stand-in in this process. Either way the - episode then plays the stand-in, which the agent's ``prompt_agent`` steps message as they would the agent.""" + step with the proxy agent in the agent's place; otherwise it serves a proxy in this process. Either way the episode + then plays the proxy, which the agent's ``prompt_agent`` steps message as they would the agent.""" type: ClassVar[str] = "episode_agent" entity_refs = () @@ -43,24 +43,22 @@ def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: episode = _episode_of(context) agent = self.step.agent_name - role = getattr(self.step, "role", None) - mcp_urls = {} - for env_id in getattr(self.step, "env_ids", None) or []: - deployed = bound_deployment(context, env_id, env_step_id=self.step.env_step_id, step_id=self.id) - mcp_urls[env_id] = deployed.mcp_url if deployed is not None else live_mcp_url(env_id) if not self.sandboxed: - deployed = await episode._serve_in_process(agent, role, mcp_urls) + mcp_urls = {} + for env_id in getattr(self.step, "env_ids", None) or []: + deployed = bound_deployment(context, env_id, env_step_id=self.step.env_step_id, step_id=self.id) + mcp_urls[env_id] = deployed.mcp_url if deployed is not None else live_mcp_url(env_id) + deployed = await episode._serve_in_process(agent, getattr(self.step, "role", None), mcp_urls) context.deployed_agents.append(deployed) if system_prompt := getattr(self.step, "system_prompt", None): config = A2AAgent.negotiate_agent_config(deployed.a2a_card, {"system_prompt": system_prompt}) await protocol.post_agent_config(deployed.a2a_url + config.endpoint, config.fields) return context - await asyncio.to_thread(write_installed, STAND_IN_AGENT_ID) - override = {"a2a_agent_id": STAND_IN_AGENT_ID} if self.step.type == "deploy_agent" else { - "stand_in": STAND_IN_AGENT_ID} + await asyncio.to_thread(write_installed, PROXY_AGENT_ID) + override = {"a2a_agent_id": PROXY_AGENT_ID} if self.step.type == "deploy_agent" else {"proxy": PROXY_AGENT_ID} step_params = context.metadata.setdefault("user_overrides", {}).setdefault("step_params", {}) step_params[self.step.id] = {**step_params.get(self.step.id, {}), **override} context = await self.step.execute(context) deployed = next(deployed for deployed in context.deployed_agents if deployed.agent_name == agent) - await episode._attach(agent, deployed, role, mcp_urls, shell=True) + await episode._attach(agent, deployed) return context diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 07d2e5540..3ecb6ac67 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -93,7 +93,7 @@ def _placed(name: str) -> str: if place is None: raise KeyError( f"agent declares required_param {name!r}, which only an install in another agent's place " - "(the step's `stand_in` param) gives" + "(the step's `proxy` param) gives" ) return place[name] @@ -240,9 +240,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if problem := _install_problem(agent): raise RuntimeError(problem) place = None - if stand_in := self.step_param_overrides(context).get("stand_in"): + if proxy := self.step_param_overrides(context).get("proxy"): place = await self._place_of(sandbox, agent) - agent = A2AAgent.get(stand_in) + agent = A2AAgent.get(proxy) if problem := _install_problem(agent): raise RuntimeError(problem) @@ -378,15 +378,15 @@ async def _trust_lines(sandbox, container: Optional[str], work_dir: str) -> str: return "".join(f"{name}={shlex.quote(value)}\n" for name, value in env.items()) async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: - """Where ``agent``, which a stand-in replaces, works on the VM host and as which user: its install/v1 params + """Where ``agent``, which a proxy replaces, works on the VM host and as which user: its install/v1 params ``workdir_host`` and ``user_host``, templates of ``{agent_name}`` as its install commands are.""" if self.container_name is not None: raise RuntimeError( - f"install_agent '{self.id}': a stand-in replaces an agent installed on the VM host, not in a container" + f"install_agent '{self.id}': a proxy replaces an agent installed on the VM host, not in a container" ) if sandbox.ON_THIS_MACHINE: raise RuntimeError( - f"install_agent '{self.id}': a stand-in isn't installed on a sandbox that runs on this machine, where " + f"install_agent '{self.id}': a proxy isn't installed on a sandbox that runs on this machine, where " "its player's commands would run here; deploy the sandbox on a remote VM provider" ) logger.info(f"Loading agent image {agent.docker_image_artifact.image_name} onto VM (to read where it works)...") @@ -397,7 +397,7 @@ async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: if not all(isinstance(value, str) and value for value in place.values()): raise RuntimeError( f"Agent {agent.id} v{agent.version} install/v1 params don't say where it works on the VM host: a " - f"stand-in installed in its place needs '{WORKDIR_HOST_KEY}' and '{USER_HOST_KEY}'" + f"proxy installed in its place needs '{WORKDIR_HOST_KEY}' and '{USER_HOST_KEY}'" ) return {name: value.format(agent_name=self.agent_name) for name, value in place.items()} diff --git a/tst/integration/agent/stand_in_modal_vm_test.py b/tst/integration/agent/proxy_modal_vm_test.py similarity index 93% rename from tst/integration/agent/stand_in_modal_vm_test.py rename to tst/integration/agent/proxy_modal_vm_test.py index 425671b63..264504b48 100644 --- a/tst/integration/agent/stand_in_modal_vm_test.py +++ b/tst/integration/agent/proxy_modal_vm_test.py @@ -1,4 +1,4 @@ -"""An agent installed on a Modal VM's host, played from this machine through the stand-in installed in its place: the +"""An agent installed on a Modal VM's host, played from this machine through the proxy installed in its place: the shape of a task that loads its inputs onto the VM, installs a harness there and checks the outputs it leaves.""" import asyncio @@ -16,7 +16,7 @@ pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_remote_sandbox("modal_vm")] -# An agent the stand-in replaces needs only a card declaring where it works on a VM host. +# An agent the proxy replaces needs only a card declaring where it works on a VM host. REPLACED_SERVER = """ class _Card: def model_dump(self): @@ -34,9 +34,9 @@ def model_dump(self): @pytest.mark.asyncio -async def test_an_agent_installed_on_a_vm_host_is_played_through_a_stand_in_installed_in_its_place(): +async def test_an_agent_installed_on_a_vm_host_is_played_through_a_proxy_installed_in_its_place(): uid = uuid.uuid4().hex[:8] - greeting = FileArtifact.put_bytes(id=f"greeting-vm-{uid}", description="stand-in on a VM", filename="greeting.txt", + greeting = FileArtifact.put_bytes(id=f"greeting-vm-{uid}", description="proxy on a VM", filename="greeting.txt", content=b"hello", content_type="text/plain") task = Task.from_dict({"id": f"copy-vm-{uid}", "steps": [ {"id": "vm", "type": "deploy_sandbox", "sandbox_name": "vm", "sandbox_mode": "vm", "sandbox_type": "modal_vm", @@ -76,6 +76,6 @@ def _replaced(uid: str) -> A2AAgent: context = Path(build_dir) (context / "a2a_server.py").write_text(REPLACED_SERVER) (context / "Dockerfile").write_text(REPLACED_DOCKERFILE) - image = DockerImageArtifact.put_context(f"replaced-agent-image-{uid}", description="stand-in on a VM", + image = DockerImageArtifact.put_context(f"replaced-agent-image-{uid}", description="proxy on a VM", context_path=str(context), dockerfile_path="Dockerfile", platform=None) return A2AAgent.put(id=f"replaced-agent-{uid}", docker_image_artifact=image) diff --git a/tst/integration/env/envs/server_env_local_test.py b/tst/integration/env/envs/server_env_local_test.py index f5907ac9e..702ad5012 100644 --- a/tst/integration/env/envs/server_env_local_test.py +++ b/tst/integration/env/envs/server_env_local_test.py @@ -26,7 +26,7 @@ from agent_env.artifact import DockerImageArtifact, EnvironmentArtifact, FileArtifact from agent_env.artifact.store import reset_artifact_store -from agent_env.config import configure, reset_config, set_image_store +from agent_env.config import configure, get_config, reset_config, set_image_store from agent_env.env import legacy_protocol from agent_env.env.env import DeployedSandboxEnv, Env from agent_env.env.envs.mcp_server import MCPServerEnv @@ -208,13 +208,17 @@ async def verify(mcp_url): """ +ADD_A_CRATE = {"role": "assistant", "content": "Adding one.", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "items_add_item", "arguments": '{"item": "crate"}'}}, + {"id": "c2", "type": "function", "function": {"name": "items_whoami", "arguments": "{}"}}]} + + @pytest.mark.asyncio async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(local_stack, tmp_path): """A trainer's view of a seeded task: the env deployed and loaded and the echo agent prompted around the outside - agent, which plays its turn through the env's tools as its role through a stand-in served in this process; its - reply recorded and - scored under the verifier's step; and every episode, played, not played, abandoned or forgotten past its turn's - deadline, leaving nothing behind.""" + agent, which plays its turn through the env's tools as its role through a proxy served in this process, which + records the model's calls and the tools they ran; its reply scored under the verifier's step; and every episode, + played, not played, abandoned or forgotten past its turn's deadline, leaving nothing behind.""" uid = uuid.uuid4().hex[:8] env = MCPServerEnv.put(id=f"server-items-{uid}", docker_image_artifact=_put_items_image(f"server-items-{uid}"), environment_name="items", env_provider_type="server") @@ -237,15 +241,18 @@ async def test_a_tasks_agent_is_played_from_a_loop_of_ones_own_as_an_episode(loc async def play(episode, add): turn = await episode.next_turn() - assert turn.prompt == "Add a crate." and turn.role == "alice@example.com" + assert turn.prompt == "Add a crate." assert [(agent["agent_name"], bool(agent["sandbox_id"])) for agent in Task.get_instance(episode.id).context["deployed_agents"]] == [("policy", False), ("echo", True)] if add: async with turn.tools() as tools: - assert {"items_add_item", "items_whoami"} <= {tool.name for tool in tools.list()} + names = {tool.name for tool in tools.list()} + assert {"items_add_item", "items_whoami"} <= names and "bash" not in names + await turn.record_model_call(ADD_A_CRATE, model="policy-model") await tools.call("items_add_item", {"item": "crate"}) assert json.loads((await tools.call("items_whoami")).text)["role"] == "alice@example.com" + await turn.record_model_call("Added a crate.", model="policy-model") await turn.end("Added a crate." if add else None) return await episode.result() @@ -258,6 +265,9 @@ async def play(episode, add): assert [agent["agent_name"] for agent in record.context["deployed_agents"]] == ["policy", "echo"] replies = {response["step_id"]: response["response"] for response in record.context["prompt_responses"]} assert replies["play"] == "Added a crate." and replies["ask-echo"].startswith("Echo:") + [play] = [response for response in record.context["prompt_responses"] if response["step_id"] == "play"] + assert [span["name"] for span in _stored(play["agent_trajectory_object_url"])] == [ + "chat", "items_add_item", "items_whoami", "chat", "turn"] abandoned = await spec.start() await abandoned.next_turn() @@ -305,40 +315,66 @@ def _items_artifact(uid: str) -> EnvironmentArtifact: @pytest.mark.asyncio -async def test_an_agent_that_works_in_its_sandbox_is_played_through_a_stand_in_deployed_there(local_stack, tmp_path): - """A task that loads files into its agent and checks what the agent left there: the agent is never deployed, - a stand-in is, in its place, and the loop's ``bash`` runs in it, as its user and in its folder, with no model - credentials; the check scores what the commands left, and nothing is left running.""" +async def test_an_agent_that_works_in_its_sandbox_is_played_through_a_proxy_deployed_there(local_stack, tmp_path): + """A task that loads files into its agent, gives it an env, and checks what the agent left in both: the agent is + never deployed, a proxy is, in its place, and every call the loop makes runs in it, ``bash`` as the agent's user in + its folder with no model credentials, the env's tools as its role; the checks score what the calls left, the proxy + records the turn, and nothing is left running.""" uid = uuid.uuid4().hex[:8] - greeting = FileArtifact.put_bytes(id=f"greeting-{uid}", description="stand-in e2e", filename="greeting.txt", + env = MCPServerEnv.put(id=f"server-items-{uid}", docker_image_artifact=_put_items_image(f"server-items-{uid}"), + environment_name="items", env_provider_type="server") + verify = FileArtifact.put_bytes(id=f"verify-crate-{uid}", description="proxy e2e", filename="verify.py", + content=VERIFY_CRATE.encode(), content_type="text/x-python") + greeting = FileArtifact.put_bytes(id=f"greeting-{uid}", description="proxy e2e", filename="greeting.txt", content=b"hello", content_type="text/plain") task = Task.from_dict({"id": f"copy-{uid}", "steps": [ - {"id": "agent", "type": "deploy_agent", "agent_name": "solver", "a2a_agent_id": f"never-deployed-{uid}"}, + {"id": "env", "type": "deploy_env", "env_id": env.id}, + {"id": "agent", "type": "deploy_agent", "agent_name": "solver", "a2a_agent_id": f"never-deployed-{uid}", + "env_ids": [env.id], "role": "bob@example.com"}, {"id": "load", "type": "load_artifact", "agent_name": "solver", "artifact_id": greeting.id, "destination_path": "/app/inputs"}, {"id": "solve", "type": "prompt_agent", "agent_name": "solver", "prompt_id": "solve", - "prompt": "Copy inputs/greeting.txt to outputs/answer.txt.", "poll_interval_seconds": 0.5}, + "prompt": "Copy inputs/greeting.txt to outputs/answer.txt, and add a crate.", "poll_interval_seconds": 0.5}, {"id": "check", "type": "verify_sandbox", "agent_name": "solver", "criteria": [ {"type": "probe_file_contains", "criterion": "the greeting is copied", "paths": ["outputs/answer.txt"], "expected": "^hello$"}]}, + {"id": "crate", "type": "env_outcome_verifier", "env_id": env.id, "file_artifact_id": verify.id, + "verifier_id": "crate"}, ]}) + copy = {"role": "assistant", "content": None, "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "bash", "arguments": json.dumps( + {"command": "mkdir -p outputs && cp inputs/greeting.txt outputs/answer.txt"})}}, + {"id": "c2", "type": "function", "function": {"name": "items_add_item", "arguments": '{"item": "crate"}'}}]} episode = await EpisodeSpec.from_task(task).start() turn = await episode.next_turn() async with turn.tools() as tools: - assert [tool.name for tool in tools.list()] == ["bash"] + names = [tool.name for tool in tools.list()] + assert names[0] == "bash" and {"items_add_item", "items_whoami"} <= set(names) assert (await tools.call("bash", {"command": "pwd; whoami; cat inputs/greeting.txt"})).text == ( "/app\nagent\nhello") assert (await tools.call("bash", {"command": "env | grep -c ^LITELLM_ || true"})).text == "0\n" + assert json.loads((await tools.call("items_whoami")).text)["role"] == "bob@example.com" + await turn.record_model_call(copy, model="solver-model", usage={"prompt_tokens": 12, "completion_tokens": 5}) await tools.call("bash", {"command": "mkdir -p outputs && cp inputs/greeting.txt outputs/answer.txt"}) - await turn.end("Copied.") + await tools.call("items_add_item", {"item": "crate"}) + await turn.record_model_call("Copied, and added a crate.", model="solver-model") + await turn.end("Copied, and added a crate.") - assert await episode.result() == {"check": 1.0} + assert await episode.result() == {"check": 1.0, "crate": 1.0} record = Task.get_instance(episode.id) [solver] = record.context["deployed_agents"] assert solver["agent_name"] == "solver" and solver["sandbox_id"] [response] = record.context["prompt_responses"] - assert (response["response"], response["tool_call_count"]) == ("Copied.", 3) - assert response["agent_trajectory_object_url"] + assert (response["response"], response["tool_call_count"]) == ("Copied, and added a crate.", 5) + stored = _stored(response["agent_trajectory_object_url"]) + assert [span["name"] for span in stored] == [ + "bash", "bash", "items_whoami", "chat", "bash", "items_add_item", "chat", "turn"] + assert stored[-1]["attributes"]["gen_ai.request.model"] == "solver-model" assert list((tmp_path / "sandboxes").glob("agent-env-*")) == [] - assert _docker("ps", "-aq", "--filter", f"name=^/agent-{solver['sandbox_id']}$").stdout.strip() == "" + for entry in (*record.context["deployed_agents"], *record.context["deployed_envs"]): + assert _docker("ps", "-aq", "--filter", f"name=^/agent-{entry['sandbox_id']}$").stdout.strip() == "" + + +def _stored(url: str) -> list[dict]: + return json.loads(get_config().get_object_store().get(url)) diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index 2eb2c0023..5c063da8a 100644 --- a/tst/integration/env/envs/server_env_modal_test.py +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -98,7 +98,6 @@ async def start(spec): try: played = await start(spec) turn = await played.next_turn() - assert turn.mcp_urls[env.id].startswith("https://") async with turn.tools() as tools: await tools.call("items_add_item", {"item": "crate"}) assert json.loads((await tools.call("items_whoami")).text)["role"] == "alice@example.com" diff --git a/tst/unit/agents/codex/test_card_contract.py b/tst/unit/agents/codex/test_card_contract.py index 8d0d95e35..0f3432057 100644 --- a/tst/unit/agents/codex/test_card_contract.py +++ b/tst/unit/agents/codex/test_card_contract.py @@ -169,6 +169,7 @@ def test_install_params(codex): params = _extensions(codex)["urn:agentenv:install/v1"]["params"] assert list(params) == [ "install_commands", "a2a_port", "required_params", "install_commands_host", "required_params_host", + "workdir_host", "user_host", ] assert params["a2a_port"] == 8000 endpoint = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] @@ -176,6 +177,7 @@ def test_install_params(codex): assert params["required_params_host"] == ["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *endpoint] assert params["install_commands"] == a2a_server._INSTALL_COMMANDS assert params["install_commands_host"] == a2a_server._INSTALL_COMMANDS_HOST + assert (params["workdir_host"], params["user_host"]) == ("/app", "root") def test_the_card_and_health_are_served(codex): diff --git a/tst/unit/agents/stand_in/__init__.py b/tst/unit/agents/proxy/__init__.py similarity index 100% rename from tst/unit/agents/stand_in/__init__.py rename to tst/unit/agents/proxy/__init__.py diff --git a/tst/unit/agents/stand_in/agent.py b/tst/unit/agents/proxy/agent.py similarity index 50% rename from tst/unit/agents/stand_in/agent.py rename to tst/unit/agents/proxy/agent.py index 7bf5d791a..522ab08dc 100644 --- a/tst/unit/agents/stand_in/agent.py +++ b/tst/unit/agents/proxy/agent.py @@ -1,4 +1,4 @@ -"""The stand-in's a2a_server, imported from its folder by its path: Claude Code's, imported by the same name, may be +"""The proxy's a2a_server, imported from its folder by its path: Claude Code's, imported by the same name, may be loaded already.""" import importlib.util @@ -6,7 +6,7 @@ import agent_env -AGENT_DIR = Path(agent_env.__file__).parent / "agents" / "agents" / "stand-in" -_spec = importlib.util.spec_from_file_location("stand_in_a2a_server", AGENT_DIR / "a2a_server.py") +AGENT_DIR = Path(agent_env.__file__).parent / "agents" / "agents" / "proxy" +_spec = importlib.util.spec_from_file_location("proxy_a2a_server", AGENT_DIR / "a2a_server.py") a2a_server = importlib.util.module_from_spec(_spec) _spec.loader.exec_module(a2a_server) diff --git a/tst/unit/agents/stand_in/test_stand_in_agent.py b/tst/unit/agents/proxy/test_proxy_agent.py similarity index 80% rename from tst/unit/agents/stand_in/test_stand_in_agent.py rename to tst/unit/agents/proxy/test_proxy_agent.py index d01cd7c31..8d090cac7 100644 --- a/tst/unit/agents/stand_in/test_stand_in_agent.py +++ b/tst/unit/agents/proxy/test_proxy_agent.py @@ -1,4 +1,4 @@ -"""The agents bundle's stand-in: its card, how it installs on a VM's host, and how the episode writes it.""" +"""The agents bundle's proxy: its card, how it installs on a VM's host, and how the episode writes it.""" from __future__ import annotations @@ -7,10 +7,10 @@ from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.bundle.installed import write_installed -from agent_env.task.episode.steps import STAND_IN_AGENT_ID +from agent_env.task.episode.steps import PROXY_AGENT_ID from agent_env.task_step.task_steps.install_agent import _build_param_resolvers -from tst.unit.agents.stand_in.agent import AGENT_DIR, a2a_server +from tst.unit.agents.proxy.agent import AGENT_DIR, a2a_server def _extensions() -> dict: @@ -18,9 +18,9 @@ def _extensions() -> dict: return {ext["uri"]: ext for ext in card["capabilities"]["extensions"]} -def test_the_card_offers_the_driver_and_a_host_install_only(): +def test_the_card_offers_the_proxy_and_a_host_install_only(): extensions = _extensions() - assert {"urn:agentenv:driver/v1", "urn:agentenv:mcp-config/v1", "urn:agentenv:trajectory/v1", + assert {"urn:agentenv:proxy/v1", "urn:agentenv:mcp-config/v1", "urn:agentenv:trajectory/v1", "urn:agentenv:install/v1"} <= set(extensions) params = extensions["urn:agentenv:install/v1"]["params"] assert params["install_commands_host"] and not params.get("install_commands") @@ -48,9 +48,9 @@ def test_the_image_works_in_app_as_agent(): def test_the_episode_writes_it_from_the_bundle_taking_no_credentials(local_stores, cli_routing): - version = write_installed(STAND_IN_AGENT_ID) - agent = A2AAgent.get(STAND_IN_AGENT_ID) + version = write_installed(PROXY_AGENT_ID) + agent = A2AAgent.get(PROXY_AGENT_ID) assert (agent.version, agent.metadata) == (version, {"takes_credentials": False}) assert agent.docker_image_artifact.context_only - assert write_installed(STAND_IN_AGENT_ID) == version + assert write_installed(PROXY_AGENT_ID) == version diff --git a/tst/unit/task/episode/conftest.py b/tst/unit/task/episode/conftest.py index 15854fe44..6afb74908 100644 --- a/tst/unit/task/episode/conftest.py +++ b/tst/unit/task/episode/conftest.py @@ -1,4 +1,4 @@ -"""An episode's turns reach a stand-in agent this process serves on loopback, so these tests allow loopback +"""An episode's turns reach a proxy agent this process serves on loopback, so these tests allow loopback connections and nothing else.""" import pytest diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index ded6a205d..f3137bb1e 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -6,7 +6,8 @@ import httpx import pytest import pytest_asyncio -from agentenv_protocol import EnvironmentTool, ToolResult +from agentenv_protocol import EnvironmentTool, ToolResult, ToolSessionError +from agentenv_protocol.a2a_agent import proxy as proxy_module import agent_env.task.episode.run as episode_run from agent_env.config import get_config @@ -107,8 +108,8 @@ async def __call__(self, context, **_): class _ToolSessions: - """Stands in for ``ToolSession``: the session on each URL serves the tools ``offered`` names for it, and every - session opened, call made and session closed is recorded.""" + """Stands in for the proxy's ``ToolSession``: the session on each URL serves the tools ``offered`` names for it, + and every session opened, call made and session closed is recorded.""" def __init__(self, offered, raises=None): self.offered = offered @@ -117,12 +118,12 @@ def __init__(self, offered, raises=None): self.calls = [] self.closed = [] - def __call__(self, *, mcp_url, role, headers, timeout): + def __call__(self, *, mcp_url, headers=None, timeout): sessions = self class _Session: async def open(self): - sessions.opened.append((mcp_url, role, headers, timeout)) + sessions.opened.append((mcp_url, headers)) return self async def list_tools(self): @@ -187,9 +188,8 @@ async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(tear assert Task.get_instance(episode.id).status == "running" turn = await episode.next_turn() - assert (turn.agent, turn.step_id, turn.prompt, turn.system_prompt, turn.role) == ( - "alice", "turn", "Add a crate.", "You keep the items.", "alice@example.com") - assert (turn.mcp_urls, turn.headers) == ({"items": "http://127.0.0.1:20031/items/mcp"}, {"items": {}}) + assert (turn.agent, turn.step_id, turn.prompt, turn.system_prompt) == ( + "alice", "turn", "Add a crate.", "You keep the items.") await turn.end("Added a crate.") assert await episode.result() == {"replies": 1.0} @@ -204,30 +204,21 @@ async def test_the_run_waits_at_an_outside_turn_and_carries_on_once_it_ends(tear ("alice", None)] -@pytest.mark.asyncio -async def test_a_turn_carries_the_headers_each_url_needs(teardowns, monkeypatch): - monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) - episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() - - turn = await episode.next_turn() - - assert turn.headers == {"items": {"X-Sandbox": "http://127.0.0.1:20031/items/mcp"}} - - ITEMS_MCP, PANTRY_MCP = "http://127.0.0.1:20031/items/mcp", "http://127.0.0.1:20031/pantry/mcp" PANTRY = {"id": "pantry", "type": "episode_test_env", "env_id": "pantry"} STOCKED = {**ALICE, "env_ids": ["items", "pantry"]} +ROLE = "AgentEnv-Role" def _tool_sessions(monkeypatch, offered, raises=None): sessions = _ToolSessions(offered, raises) - monkeypatch.setattr(episode_run, "ToolSession", sessions) + monkeypatch.setattr(proxy_module, "ToolSession", sessions) monkeypatch.setattr(episode_run, "sandbox_request_headers_for_url", lambda url: {"X-Sandbox": url}) return sessions @pytest.mark.asyncio -async def test_a_turns_tools_call_each_env_as_its_role_with_its_headers(teardowns, monkeypatch): +async def test_a_turns_env_tools_run_in_its_proxy_with_the_headers_each_url_needs_and_the_role(teardowns, monkeypatch): sessions = _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"], PANTRY_MCP: ["pantry_count"]}) episode = await _spec(ITEMS, PANTRY, STOCKED, TURN, REPLIES).start() turn = await episode.next_turn() @@ -235,26 +226,27 @@ async def test_a_turns_tools_call_each_env_as_its_role_with_its_headers(teardown async with turn.tools(timeout=30) as tools: assert [tool.name for tool in tools.list()] == ["items_add_item", "pantry_count"] assert (await tools.call("pantry_count")).text == f"pantry_count on {PANTRY_MCP}" + await turn.end("Counted.") + await episode.result() - assert sessions.opened == [(url, "alice@example.com", {"X-Sandbox": url}, 30) for url in (ITEMS_MCP, PANTRY_MCP)] - assert sessions.calls == [(PANTRY_MCP, "pantry_count", None)] + assert sessions.opened == [(url, {"X-Sandbox": url, ROLE: "alice@example.com"}) for url in (ITEMS_MCP, PANTRY_MCP)] + assert sessions.calls == [(PANTRY_MCP, "pantry_count", {})] assert sorted(sessions.closed) == [ITEMS_MCP, PANTRY_MCP] @pytest.mark.asyncio -async def test_a_tool_two_envs_offer_is_refused_as_the_tools_open_and_nothing_is_left_open(teardowns, monkeypatch): - sessions = _tool_sessions(monkeypatch, {ITEMS_MCP: ["count"], PANTRY_MCP: ["count"]}) +async def test_a_tool_two_envs_offer_is_refused_as_the_tools_open(teardowns, monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["count"], PANTRY_MCP: ["count"]}) episode = await _spec(ITEMS, PANTRY, STOCKED, TURN, REPLIES).start() turn = await episode.next_turn() - with pytest.raises(EpisodeError, match="tool 'count' is offered by both env 'items' and env 'pantry'"): + with pytest.raises(EpisodeError, match="tool 'count' is offered by both 'items' and 'pantry'"): async with turn.tools(): pass - assert sorted(sessions.closed) == [ITEMS_MCP, PANTRY_MCP] @pytest.mark.asyncio -async def test_a_tool_no_env_offers_is_refused(teardowns, monkeypatch): +async def test_a_tool_the_turn_doesnt_offer_is_refused(teardowns, monkeypatch): _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() turn = await episode.next_turn() @@ -262,7 +254,26 @@ async def test_a_tool_no_env_offers_is_refused(teardowns, monkeypatch): async with turn.tools() as tools: with pytest.raises(EpisodeError) as raised: await tools.call("nope") - assert str(raised.value) == "no env of turn 'turn' offers tool 'nope'; they offer ['items_add_item']" + assert str(raised.value) == "turn 'turn' offers no tool 'nope'; it offers ['items_add_item']" + + +@pytest.mark.asyncio +async def test_an_env_tool_that_fails_answers_an_error_and_is_recorded(teardowns, monkeypatch): + _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}, + raises={"items_add_item": ToolSessionError("the env didn't answer")}) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() + + async with turn.tools() as tools: + result = await tools.call("items_add_item", {"name": "crate"}) + await turn.end("Gave up.") + await episode.result() + + assert (result.isError, result.text) == (True, "error: ToolSessionError: the env didn't answer") + [response] = Task.get_instance(episode.id).context["prompt_responses"] + stored = json.loads(get_config().get_object_store().get(response["agent_trajectory_object_url"])) + assert json.loads(stored[0]["attributes"]["gen_ai.completion"]) == { + "error": "error: ToolSessionError: the env didn't answer"} @pytest.mark.asyncio @@ -301,7 +312,8 @@ async def play(n): @pytest.mark.asyncio -async def test_a_seed_reaches_the_turns_prompt_env_binding_and_role(teardowns): +async def test_a_seed_reaches_the_turns_prompt_env_binding_and_role(teardowns, monkeypatch): + sessions = _tool_sessions(monkeypatch, {PANTRY_MCP: ["pantry_count"]}) episode = await _spec( {**ITEMS, "env_id": ""}, {**ALICE, "env_ids": [""], "role": "", "system_prompt": "You keep the ."}, @@ -310,9 +322,11 @@ async def test_a_seed_reaches_the_turns_prompt_env_binding_and_role(teardowns): ).start() turn = await episode.next_turn() + async with turn.tools(): + pass assert (turn.prompt, turn.system_prompt) == ("Add a crate to the pantry.", "You keep the pantry.") - assert (turn.mcp_urls, turn.role) == ({"pantry": "http://127.0.0.1:20031/pantry/mcp"}, "bob@example.com") + assert sessions.opened == [(PANTRY_MCP, {"X-Sandbox": PANTRY_MCP, ROLE: "bob@example.com"})] @pytest.mark.asyncio @@ -596,90 +610,118 @@ async def test_a_cancelled_start_takes_its_run_down_with_it(teardowns): assert teardowns.finished == teardowns.started and len(teardowns.finished) == 1 -class _SandboxedStandIn: - """Stands in for the client of a stand-in deployed in a sandbox: runs each command as ``ran ``, and - records each reply.""" +class _SandboxedProxy: + """Stands in for the client of a proxy deployed in a sandbox: offers ``items_add_item`` and ``bash``, answers each + action from ``results`` or ran ````, and records each action and reply.""" - def __init__(self): - self.commands = [] + def __init__(self, results=None): + self.results = results or {} + self.actions = [] self.replies = [] - async def run(self, command): - self.commands.append(command) - return f"ran {command}" + async def tools(self, task_id): + return [{"name": "items_add_item", "server": "items"}, {"name": "bash", "description": "Run a command."}] + + async def act(self, task_id, action): + self.actions.append(action) + if action["kind"] == "bash": + return self.results.get(action["command"]) or { + "kind": "bash", "exit_code": 0, "stdout": f"ran {action['command']}\n", "stderr": "", "timed_out": False} + return {"kind": "tool", "content": [{"type": "text", "text": "added"}], "is_error": False} async def reply(self, task_id, **reply): self.replies.append((task_id, reply)) return True -def _sandboxed_turn(stand_in): - return episode_run.Turn(agent="alice", step_id="turn", prompt="Add a crate.", system_prompt=None, role=None, - mcp_urls={"items": ITEMS_MCP}, skills=[], shell=True, stand_in=stand_in, task_id="t-1") +def _sandboxed_turn(proxy): + return episode_run.Turn(agent="alice", step_id="turn", prompt="Add a crate.", system_prompt=None, skills=[], + proxy=proxy, task_id="t-1") @pytest.mark.asyncio -async def test_a_turn_in_the_agents_sandbox_offers_bash_beside_its_envs_tools(monkeypatch): - _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) - stand_in = _SandboxedStandIn() +async def test_a_turn_in_the_agents_sandbox_offers_bash_and_shows_its_output_as_verifiers_does(): + proxy = _SandboxedProxy({"slow": {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True}, + "both": {"kind": "bash", "exit_code": 2, "stdout": "out\n", "stderr": "err\n", + "timed_out": False}}) - async with _sandboxed_turn(stand_in).tools() as tools: + async with _sandboxed_turn(proxy).tools() as tools: assert [tool.name for tool in tools.list()] == ["items_add_item", "bash"] - assert (await tools.call("bash", {"command": "ls inputs"})).text == "ran ls inputs" + assert (await tools.call("bash", {"command": "ls inputs"})).text == "ran ls inputs\n" + assert (await tools.call("bash", {"command": "both"})).text == "out\nerr\n" + timed_out = await tools.call("bash", {"command": "slow"}) with pytest.raises(EpisodeError, match="tool 'bash' takes a command string, not 3"): await tools.call("bash", {"command": 3}) + assert (await tools.call("items_add_item", {"name": "crate"})).text == "added" - assert stand_in.commands == ["ls inputs"] + assert (timed_out.isError, timed_out.text) == (True, "error: the command timed out") + assert proxy.actions == [{"kind": "bash", "command": "ls inputs"}, {"kind": "bash", "command": "both"}, + {"kind": "bash", "command": "slow"}, + {"kind": "tool", "server": "items", "name": "items_add_item", + "arguments": {"name": "crate"}, "timeout_seconds": 600.0}] @pytest.mark.asyncio -async def test_an_env_tool_named_bash_clashes_with_the_sandboxes_own(monkeypatch): - _tool_sessions(monkeypatch, {ITEMS_MCP: ["bash"]}) +async def test_a_turn_ends_with_its_reply_as_parts(): + proxy = _SandboxedProxy() + turn = _sandboxed_turn(proxy) - with pytest.raises(EpisodeError, match="tool 'bash' is offered by env 'items', and the agent's sandbox offers"): - async with _sandboxed_turn(_SandboxedStandIn()).tools(): - pass + await turn.end("Added a crate.") + await turn.end("Again.") + + assert proxy.replies == [("t-1", {"parts": [{"kind": "text", "text": "Added a crate."}]})] + + +@pytest.mark.asyncio +async def test_a_reply_that_cant_reach_the_proxy_leaves_the_runs_own_failure_to_report(caplog): + class _Unreachable(_SandboxedProxy): + async def reply(self, task_id, **reply): + raise httpx.ConnectError("the sandbox is gone") + + turn = _sandboxed_turn(_Unreachable()) + + await turn.end("done") + + assert "its reply didn't reach the proxy: the sandbox is gone" in caplog.text + + +ADD_A_CRATE = {"role": "assistant", "content": "Adding one.", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "items_add_item", "arguments": '{"name": "crate"}'}}]} @pytest.mark.asyncio -async def test_a_turns_calls_and_its_reply_are_handed_back_as_its_trajectory(monkeypatch): +async def test_recorded_model_calls_are_held_to_the_calls_that_run(teardowns, monkeypatch): _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}) - stand_in = _SandboxedStandIn() - turn = _sandboxed_turn(stand_in) + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() async with turn.tools() as tools: + await turn.record_model_call(ADD_A_CRATE, model="qwen", usage={"prompt_tokens": 9, "completion_tokens": 4}) + with pytest.raises(EpisodeError, match="didn't make this call to 'items_add_item'"): + await tools.call("items_add_item", {"name": "box"}) + with pytest.raises(EpisodeError, match="haven't run"): + await turn.end("Added.") await tools.call("items_add_item", {"name": "crate"}) - await tools.call("bash", {"command": "ls"}) + await turn.record_model_call("Added a crate.", model="qwen") await turn.end("Added a crate.") + assert await episode.result() == {"replies": 1.0} - [(task_id, reply)] = stand_in.replies - assert (task_id, reply["text"], reply["tool_call_count"]) == ("t-1", "Added a crate.", 2) - spans = reply["trajectory"] - assert [(span["name"], span["attributes"]["gen_ai.operation.name"]) for span in spans] == [ - ("items_add_item", "execute_tool"), ("bash", "execute_tool"), ("turn", "chain")] - assert json.loads(spans[0]["attributes"]["gen_ai.prompt"]) == {"input": {"name": "crate"}} - assert json.loads(spans[1]["attributes"]["gen_ai.completion"]) == {"output": [{"type": "text", "text": "ran ls"}]} - assert json.loads(spans[2]["attributes"]["gen_ai.completion"]) == { - "content": [{"type": "text", "text": "Added a crate."}]} - assert all(span["start_time"] <= span["end_time"] for span in spans) + [response] = Task.get_instance(episode.id).context["prompt_responses"] + stored = json.loads(get_config().get_object_store().get(response["agent_trajectory_object_url"])) + assert [(span["name"], span["attributes"]["gen_ai.operation.name"]) for span in stored] == [ + ("chat", "chat"), ("items_add_item", "execute_tool"), ("chat", "chat"), ("turn", "chain")] + assert json.loads(stored[-1]["attributes"]["gen_ai.prompt"]) == { + "prompt": "Add a crate.", "system_prompt": "You keep the items."} + assert stored[-1]["attributes"]["gen_ai.request.model"] == "qwen" @pytest.mark.asyncio -async def test_a_call_that_raises_is_handed_back_in_the_trajectory_too(monkeypatch): - _tool_sessions(monkeypatch, {ITEMS_MCP: ["items_add_item"]}, - raises={"items_add_item": TimeoutError("the env didn't answer")}) - stand_in = _SandboxedStandIn() - turn = _sandboxed_turn(stand_in) - - async with turn.tools() as tools: - with pytest.raises(TimeoutError, match="the env didn't answer"): - await tools.call("items_add_item", {"name": "crate"}) - await turn.end("Gave up.") +async def test_a_model_call_in_a_shape_the_proxy_cant_read_is_refused(teardowns): + episode = await _spec(ITEMS, ALICE, TURN, REPLIES).start() + turn = await episode.next_turn() - [(_, reply)] = stand_in.replies - assert reply["tool_call_count"] == 1 - assert json.loads(reply["trajectory"][0]["attributes"]["gen_ai.completion"]) == { - "error": "TimeoutError: the env didn't answer"} + with pytest.raises(EpisodeError, match="OpenAI chat message, an Anthropic message or a string"): + await turn.record_model_call({"role": "assistant", "content": {"not": "readable"}}) @pytest.mark.asyncio @@ -720,16 +762,3 @@ async def test_a_prompts_own_system_prompt_wins_over_the_agents(teardowns): episode = await _spec(ITEMS, ALICE, {**TURN, "system_prompt": "Count first."}, REPLIES).start() assert (await episode.next_turn()).system_prompt == "Count first." - - -@pytest.mark.asyncio -async def test_a_reply_that_cant_reach_the_stand_in_leaves_the_runs_own_failure_to_report(caplog): - class _Unreachable(_SandboxedStandIn): - async def reply(self, task_id, **reply): - raise httpx.ConnectError("the sandbox is gone") - - turn = _sandboxed_turn(_Unreachable()) - - await turn.end("done") - - assert "its reply didn't reach the stand-in: the sandbox is gone" in caplog.text diff --git a/tst/unit/task/episode/spec_test.py b/tst/unit/task/episode/spec_test.py index 1246b673c..e4e560d71 100644 --- a/tst/unit/task/episode/spec_test.py +++ b/tst/unit/task/episode/spec_test.py @@ -89,12 +89,12 @@ def test_the_copy_stands_in_for_the_outside_agents_steps_and_keeps_everything_el ] assert [copy.steps[i] is task.steps[i] for i in range(len(task.steps))] == [ True, True, False, True, True, True, True] - for original, stand_in in zip(task.steps, copy.steps): - assert (stand_in.id, stand_in.version, stand_in.depends_on, stand_in.fail_task_on_error) == ( + for original, played in zip(task.steps, copy.steps): + assert (played.id, played.version, played.depends_on, played.fail_task_on_error) == ( original.id, original.version, original.depends_on, original.fail_task_on_error) -def test_the_outside_agents_deploy_runs_in_a_stand_in_unseeded_and_round_trips(): +def test_the_outside_agents_deploy_runs_with_a_proxy_unseeded_and_round_trips(): task = _portsim() agent = EpisodeSpec.from_task(task, seed=SEED)._copy.steps[2] @@ -153,12 +153,12 @@ def test_an_inside_agent_keeps_its_own_steps(): @pytest.mark.parametrize(("extra", "replace"), IN_THE_SANDBOX.values(), ids=IN_THE_SANDBOX.keys()) -def test_a_task_using_the_outside_agents_sandbox_deploys_its_stand_in_there(extra, replace): +def test_a_task_using_the_outside_agents_sandbox_deploys_its_proxy_there(extra, replace): assert EpisodeSpec.from_task(_portsim(*extra, **replace))._copy.steps[2].sandboxed @pytest.mark.parametrize(("extra", "replace"), OUTSIDE_THE_SANDBOX.values(), ids=OUTSIDE_THE_SANDBOX.keys()) -def test_a_task_using_only_the_outside_agents_envs_serves_its_stand_in_here(extra, replace): +def test_a_task_using_only_the_outside_agents_envs_serves_its_proxy_here(extra, replace): assert not EpisodeSpec.from_task(_portsim(*extra, **replace))._copy.steps[2].sandboxed @@ -173,7 +173,7 @@ def test_a_task_using_only_the_outside_agents_envs_serves_its_stand_in_here(extr "deploy_human_agent": (({"id": "human", "type": "deploy_human_agent", **AT_THE_AGENT},), {"agent": None}, "step 'human': deploy_human_agent deploys 'default-agent', which is played from outside"), "an install into a container": ({}, {"agent": {**INSTALL, "container_name": "task"}}, - "step 'agent': an outside agent's stand-in replaces an agent installed on the VM " + "step 'agent': an outside agent's proxy replaces an agent installed on the VM " "host, not in container 'task'"), "parts": ({}, {"play": {**PLAY, "prompt": None, "parts": [{"kind": "text", "text": "Plan."}]}}, "step 'play': an outside agent's prompt can't have parts, only a text prompt"), diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index 6cb064023..c43deeffc 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -312,25 +312,25 @@ def test_the_trust_lines_resolve_as_the_trust_env_param(): resolvers = _build_param_resolvers(container_name="c", work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, config=None, agent_name="a", trust_lines="SSL_CERT_FILE=/x\n") assert resolvers["trust_env"]() == "SSL_CERT_FILE=/x\n" -STAND_IN = "@local/agentenv-framework/agents/stand-in" +PROXY = "@local/agentenv-framework/agents/proxy" CODEX_INSTALL = {"params": {"workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", "install_commands_host": ["install codex"], "a2a_port": 8000}} -STAND_IN_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], +PROXY_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], "required_params_host": ["user", "workdir", "a2a_port"], "a2a_port": 8000}} def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, container_name=None): - """An install of ``agent-x`` that the step's ``stand_in`` param replaces with the stand-in.""" + """An install of ``agent-x`` that the step's ``proxy`` param replaces with the proxy.""" sandbox = sandbox or _InstallSandbox() monkeypatch.setattr(sandbox_provider, "build_sandbox_provider", lambda name: SimpleNamespace(get_sandbox=_resolved(sandbox))) agents = {agent_id: SimpleNamespace(id=agent_id, version=1, metadata={}, docker_image_artifact=SimpleNamespace( image_name=f"{agent_id}-image", build_context_object_url=f"s3://bucket/{agent_id}.tar.gz")) - for agent_id in ("agent-x", STAND_IN)} + for agent_id in ("agent-x", PROXY)} monkeypatch.setattr(A2AAgent, "get", staticmethod(lambda agent_id, version=None: agents[agent_id])) step = InstallAgentTaskStep(id="i", version=None, sandbox_name="vm", a2a_agent_id="agent-x", agent_name="solver", container_name=container_name) - cards = {"agent-x-image": codex, f"{STAND_IN}-image": STAND_IN_INSTALL} + cards = {"agent-x-image": codex, f"{PROXY}-image": PROXY_INSTALL} monkeypatch.setattr(step, "_extract_install_extension", lambda sandbox, image: _resolved(cards[image])()) monkeypatch.setattr(step, "_wait_for_agent_card", _resolved({})) monkeypatch.setattr(step, "_register_instance", lambda *args: SimpleNamespace(instance_id="solver-1")) @@ -338,14 +338,14 @@ def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, cont deployed_sandboxes=[DeployedSandbox( sandbox_name="vm", sandbox_id="sb-1", sandbox_mode="vm", sandbox_type="modal_vm", tunnel_urls={"8000": "https://vm-8000.example"})], - metadata={"user_overrides": {"step_params": {"i": {"stand_in": STAND_IN}}}, + metadata={"user_overrides": {"step_params": {"i": {"proxy": PROXY}}}, "deployed_docker_containers": [{"container_name": "task", "sandbox_name": "vm"}]}, ) return step, context, sandbox @pytest.mark.asyncio -async def test_a_stand_in_is_installed_where_the_agent_it_replaces_says_it_works(monkeypatch, local_stores): +async def test_a_proxy_is_installed_where_the_agent_it_replaces_says_it_works(monkeypatch, local_stores): step, context, sandbox = _installing_in_place(monkeypatch) await step.execute(context) @@ -356,7 +356,7 @@ async def test_a_stand_in_is_installed_where_the_agent_it_replaces_says_it_works @pytest.mark.asyncio -async def test_a_stand_in_needs_the_replaced_agent_to_say_where_it_works(monkeypatch, local_stores): +async def test_a_proxy_needs_the_replaced_agent_to_say_where_it_works(monkeypatch, local_stores): step, context, _ = _installing_in_place(monkeypatch, codex={"params": {"install_commands_host": ["install codex"], "workdir_host": "/opt/agent"}}) @@ -365,7 +365,7 @@ async def test_a_stand_in_needs_the_replaced_agent_to_say_where_it_works(monkeyp @pytest.mark.asyncio -async def test_a_stand_in_isnt_installed_into_a_container(monkeypatch, local_stores): +async def test_a_proxy_isnt_installed_into_a_container(monkeypatch, local_stores): step, context, _ = _installing_in_place(monkeypatch, container_name="task") with pytest.raises(RuntimeError, match="installed on the VM host, not in a container"): @@ -373,7 +373,7 @@ async def test_a_stand_in_isnt_installed_into_a_container(monkeypatch, local_sto @pytest.mark.asyncio -async def test_a_stand_in_isnt_installed_on_a_sandbox_that_runs_on_this_machine(monkeypatch, local_stores): +async def test_a_proxy_isnt_installed_on_a_sandbox_that_runs_on_this_machine(monkeypatch, local_stores): here = _InstallSandbox() here.ON_THIS_MACHINE = True step, context, _ = _installing_in_place(monkeypatch, sandbox=here) @@ -383,7 +383,7 @@ async def test_a_stand_in_isnt_installed_on_a_sandbox_that_runs_on_this_machine( def test_only_an_install_in_another_agents_place_gives_its_workdir_and_user(): - with pytest.raises(KeyError, match="stand_in"): + with pytest.raises(KeyError, match="proxy"): _resolvers(None)["workdir"]() placed = _build_param_resolvers( container_name=None, work_dir="/tmp/w", agent_ctx_tar="/tmp/w/c.tgz", a2a_port=8000, config=object(), diff --git a/tst/unit/test_infra_images_script.py b/tst/unit/test_infra_images_script.py index 2b4cb82b1..0998ddf32 100644 --- a/tst/unit/test_infra_images_script.py +++ b/tst/unit/test_infra_images_script.py @@ -60,7 +60,8 @@ def test_plan_only_lists_the_infra_images_or_the_agents_images(script, monkeypat agents = [image.name for image in AGENT_IMAGES] assert planned == {"infra": [image.name for image in bootstrap.PUBLISHED_IMAGES if image.name not in agents], "agents": agents} - assert "agentenv-gateway" in planned["infra"] and agents == ["agentenv-claude-code", "agentenv-codex", "agentenv-openhands", "agentenv-pi"] + assert "agentenv-gateway" in planned["infra"] and agents == [ + "agentenv-claude-code", "agentenv-codex", "agentenv-openhands", "agentenv-pi", "agentenv-proxy"] def test_plan_missing_lists_the_images_the_registry_lacks_or_refuses_to_show(script, monkeypatch, capsys): From 27bb3220e7f040858b1225ef05e98630c352c801 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 10:24:37 -0700 Subject: [PATCH 14/39] fix(protocol): a tool action's deadline covers connecting, and turns hold one connection per server - A tool action's timeout_seconds bounds connecting to its env server and the call together; the proxy's sessions set no deadline of their own, which cut a longer request off at 600 s. Listing a turn's tools is bounded the same way. - A turn opens one connection per env server, under a lock, shared by its actions; none is opened or kept once the turn has closed. - record checks and sets the turn's pending calls with no wait in between, so of two model calls recorded at once the second is refused. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/proxy.py | 52 +++++++++---- .../agentenv-protocol/tests/test_a2a_proxy.py | 73 ++++++++++++++++++- 2 files changed, 109 insertions(+), 16 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 5798a3a09..5d33b6955 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -13,9 +13,11 @@ import contextlib import hmac import json +import math import os import shutil import signal +from collections import defaultdict from collections.abc import AsyncIterator from dataclasses import dataclass, field from datetime import datetime, timezone @@ -56,8 +58,9 @@ ) from .framework import AgentEnvAgent, AgentIdentity, a2a_agent from .tasks.v1 import AgentConfig, TaskRequest, TaskResult, TextPart, TrajectoryLog, Usage -from ..tool_client import ToolSession, ToolSessionError +from ..tool_client import ToolResult, ToolSession, ToolSessionError from ..transfers import download +from ..types import EnvironmentTool LONG_POLL_SECONDS = 25.0 BASH_TIMEOUT_SECONDS = 3600 @@ -93,6 +96,8 @@ class _Turn: ) actions: dict[str, asyncio.Task[ProxyBashResult | ProxyToolResult]] = field(default_factory=dict) sessions: dict[str, ToolSession] = field(default_factory=dict) + connecting: defaultdict[str, asyncio.Lock] = field(default_factory=lambda: defaultdict(asyncio.Lock)) + closed: bool = False tools: dict[str, ProxyTool] | None = None listing: asyncio.Lock = field(default_factory=asyncio.Lock) recorded: bool = False @@ -223,11 +228,11 @@ async def result(self, request: ProxyResultRequest) -> ProxyActionState: @extension(PROXY_V1.record) async def record(self, request: ProxyRecordRequest) -> None: - turn = self._open_turn(request.token, request.task_id) content, calls = _read_message(request.message) + offered = await self._tools(self._open_turn(request.token, request.task_id)) + turn = self._open_turn(request.token, request.task_id) if turn.pending: raise HTTPException(status_code=409, detail=_unrun(turn)) - offered = await self._tools(turn) turn.pending = [(name, _canonical(arguments)) for name, arguments in calls if name in offered] turn.recorded = True turn.model = request.model or turn.model @@ -279,7 +284,10 @@ async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: offered = {BASH: BASH_TOOL} if self._shell else {} for server in turn.servers: try: - listed = await (await self._session(turn, server)).list_tools() + listed = await asyncio.wait_for(self._listed(turn, server), TOOL_TIMEOUT_SECONDS) + except TimeoutError: + raise HTTPException(status_code=502, detail=( + f"env server {server!r} didn't list its tools within {TOOL_TIMEOUT_SECONDS}s")) from None except (ToolSessionError, httpx.HTTPError) as e: raise HTTPException(status_code=502, detail=f"env server {server!r}: {e}") from None for tool in listed: @@ -292,15 +300,29 @@ async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: turn.tools = offered return turn.tools + async def _listed(self, turn: _Turn, server: str) -> list[EnvironmentTool]: + return await (await self._session(turn, server)).list_tools() + + async def _called(self, turn: _Turn, action: ProxyToolAction) -> ToolResult: + return await (await self._session(turn, action.server)).call_tool(action.name, action.arguments) + async def _session(self, turn: _Turn, server: str) -> ToolSession: - if server not in turn.sessions: - registration = turn.servers.get(server) - if registration is None: - raise ToolSessionError(f"no env server {server!r} is registered") - turn.sessions[server] = await ToolSession(mcp_url=registration["url"], - headers=registration.get("headers") or None, - timeout=TOOL_TIMEOUT_SECONDS).open() - return turn.sessions[server] + """The turn's one connection to ``server``, which its first caller opens. It sets no deadline of its own: each + caller bounds what it waits for. None is opened, or kept, once the turn has closed.""" + async with turn.connecting[server]: + if turn.closed: + raise ToolSessionError("the turn is over") + if server not in turn.sessions: + registration = turn.servers.get(server) + if registration is None: + raise ToolSessionError(f"no env server {server!r} is registered") + session = await ToolSession(mcp_url=registration["url"], headers=registration.get("headers") or None, + timeout=math.inf).open() + if turn.closed: + await session.close() + raise ToolSessionError("the turn is over") + turn.sessions[server] = session + return turn.sessions[server] async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) -> ProxyBashResult | ProxyToolResult: started = _now() @@ -315,10 +337,9 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) timeout = action.timeout_seconds or TOOL_TIMEOUT_SECONDS name, arguments = action.name, action.arguments try: - session = await self._session(turn, action.server) - called = await asyncio.wait_for(session.call_tool(action.name, action.arguments), timeout) + called = await asyncio.wait_for(self._called(turn, action), timeout) result = ProxyToolResult(content=called.content, is_error=called.isError) - except asyncio.TimeoutError: + except TimeoutError: result = _tool_error(f"the tool call timed out after {timeout:g}s") except (ToolSessionError, httpx.HTTPError) as e: result = _tool_error(f"{type(e).__name__}: {e}") @@ -337,6 +358,7 @@ async def _state(self, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> async def _close(self, turn: _Turn) -> None: """Stop what the turn still runs and forget its actions, so what the driver never collected goes with it.""" + turn.closed = True for action_id, task in turn.actions.items(): task.cancel() self._actions.pop(action_id, None) diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index dcc5ce9a7..478bab14b 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json import socket import subprocess @@ -11,7 +12,7 @@ import uvicorn from starlette.testclient import TestClient -from agentenv_protocol import EnvironmentCard, ToolContext, create_fastmcp_app, tool +from agentenv_protocol import EnvironmentCard, ToolContext, ToolSession, create_fastmcp_app, tool from agentenv_protocol.a2a_agent import proxy from agentenv_protocol.a2a_agent.proxy import ProxyAgent @@ -267,6 +268,53 @@ def test_a_tool_action_calls_the_env_with_the_headers_registered_for_the_proxy(c assert failed["is_error"] is True +def test_a_tool_actions_deadline_covers_connecting_to_its_server(client): + with socket.socket() as silent: + silent.bind(("127.0.0.1", 0)) + silent.listen() + _register(client, f"http://127.0.0.1:{silent.getsockname()[1]}/mcp") + task_id = _send(client, "go") + + started = time.monotonic() + result = _done(client, task_id, "t", {"kind": "tool", "server": "items", "name": "add_item", + "timeout_seconds": 0.5}) + + assert result["is_error"] and "timed out after 0.5s" in result["content"][0]["text"] + assert time.monotonic() - started < 5 + + +def test_a_turns_actions_share_one_connection_to_each_env_server_which_closes_with_the_turn(client, env, monkeypatch): + opened, closed, results = [], [], [] + open_session, close_session = ToolSession.open, ToolSession.close + + async def opening(self): + opened.append(self) + await asyncio.sleep(0.2) + return await open_session(self) + + async def closing(self): + if self._client is not None: + closed.append(self) + await close_session(self) + + monkeypatch.setattr(ToolSession, "open", opening) + monkeypatch.setattr(ToolSession, "close", closing) + _register(client, env[1]) + task_id = _send(client, "go") + + calls = [threading.Thread(target=lambda n=n: results.append(_done( + client, task_id, f"w{n}", {"kind": "tool", "server": "items", "name": "whoami"}))) for n in range(3)] + for call in calls: + call.start() + for call in calls: + call.join() + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) + _settled(client, task_id) + + assert [result["is_error"] for result in results] == [False] * 3 + assert len(opened) == 1 and closed == opened + + def test_a_bash_action_runs_in_the_proxys_folder_and_returns_a_structured_result(client, tmp_path): (tmp_path / "inputs").mkdir() task_id = _send(client, "go") @@ -368,6 +416,29 @@ def test_recorded_model_calls_are_checked_against_the_actions_that_run(client, e assert task["status"]["message"]["parts"][1]["data"]["usage"]["tool_call_count"] == 2 +def test_two_model_calls_recorded_at_once_are_held_one_after_the_other(client, monkeypatch): + listing, statuses = ProxyAgent._tools, [] + + async def slow(self, turn): + await asyncio.sleep(0.3) + return await listing(self, turn) + + def record(command: str) -> None: + message = {"role": "assistant", "content": None, "tool_calls": [ + {"id": command, "type": "function", "function": {"name": "bash", "arguments": json.dumps({"command": command})}}]} + statuses.append(_proxy(client, "record", task_id=task_id, message=message).status_code) + + monkeypatch.setattr(ProxyAgent, "_tools", slow) + task_id = _send(client, "go") + records = [threading.Thread(target=record, args=(command,)) for command in ("true", "false")] + for thread in records: + thread.start() + for thread in records: + thread.join() + + assert sorted(statuses) == [200, 409] + + @pytest.mark.parametrize(("message", "content"), [ ("Plain words.", [{"type": "text", "text": "Plain words."}]), ({"role": "assistant", "content": [{"type": "thinking", "thinking": "Hmm."}, {"type": "text", "text": "Hi."}, From 845d7ed70d012c9caefab0924ab431e950aef45a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 10:26:42 -0700 Subject: [PATCH 15/39] fix(task): a turn's bash result shows a failed command's exit code and is an error Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/run.py | 9 ++++++--- tst/unit/task/episode/run_test.py | 16 ++++++++++------ 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index f978e8d6e..e49876d29 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -144,9 +144,12 @@ async def call(self, name: str, arguments: dict[str, Any] | None = None) -> Tool "timeout_seconds": self._timeout} result = await self._turn._checked(self._turn._proxy.act(self._turn._task_id, action)) if result["kind"] == "bash": - timed_out = result["timed_out"] - text = "error: the command timed out" if timed_out else result["stdout"] + result["stderr"] - return ToolResult(content=[{"type": "text", "text": text}], isError=timed_out) + if result["timed_out"]: + return ToolResult(content=[{"type": "text", "text": "error: the command timed out"}], isError=True) + text = result["stdout"] + result["stderr"] + if failed := result["exit_code"] != 0: + text = "\n".join(filter(None, (text.removesuffix("\n"), f"exit code {result['exit_code']}"))) + return ToolResult(content=[{"type": "text", "text": text}], isError=failed) return ToolResult(content=result["content"], isError=result["is_error"]) diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index f3137bb1e..a1b87e33c 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -640,23 +640,27 @@ def _sandboxed_turn(proxy): @pytest.mark.asyncio -async def test_a_turn_in_the_agents_sandbox_offers_bash_and_shows_its_output_as_verifiers_does(): +async def test_a_turn_in_the_agents_sandbox_offers_bash_and_shows_its_output_and_a_failed_exit(): proxy = _SandboxedProxy({"slow": {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True}, "both": {"kind": "bash", "exit_code": 2, "stdout": "out\n", "stderr": "err\n", - "timed_out": False}}) + "timed_out": False}, + "quiet": {"kind": "bash", "exit_code": 1, "stdout": "", "stderr": "", "timed_out": False}}) async with _sandboxed_turn(proxy).tools() as tools: assert [tool.name for tool in tools.list()] == ["items_add_item", "bash"] - assert (await tools.call("bash", {"command": "ls inputs"})).text == "ran ls inputs\n" - assert (await tools.call("bash", {"command": "both"})).text == "out\nerr\n" + ran = await tools.call("bash", {"command": "ls inputs"}) + both = await tools.call("bash", {"command": "both"}) + quiet = await tools.call("bash", {"command": "quiet"}) timed_out = await tools.call("bash", {"command": "slow"}) with pytest.raises(EpisodeError, match="tool 'bash' takes a command string, not 3"): await tools.call("bash", {"command": 3}) assert (await tools.call("items_add_item", {"name": "crate"})).text == "added" - assert (timed_out.isError, timed_out.text) == (True, "error: the command timed out") + assert [(result.isError, result.text) for result in (ran, both, quiet, timed_out)] == [ + (False, "ran ls inputs\n"), (True, "out\nerr\nexit code 2"), (True, "exit code 1"), + (True, "error: the command timed out")] assert proxy.actions == [{"kind": "bash", "command": "ls inputs"}, {"kind": "bash", "command": "both"}, - {"kind": "bash", "command": "slow"}, + {"kind": "bash", "command": "quiet"}, {"kind": "bash", "command": "slow"}, {"kind": "tool", "server": "items", "name": "items_add_item", "arguments": {"name": "crate"}, "timeout_seconds": 600.0}] From 88f94827958af0a4a4326ea8aee013e88b9bc101 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 10:36:49 -0700 Subject: [PATCH 16/39] fix(bundle): plan_entities takes no agent override; plan passes its own through Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/plan.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index cbc41393f..48a4559e5 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -482,15 +482,16 @@ def plan(self, tasks: Sequence[str], evals: Sequence[str], agent: str | None = N self._check_writers() self._check_loops() selected_tasks, selected_evals = self._selection(tasks, evals) - return self._planned(selected_tasks + selected_evals, selected_tasks, selected_evals, aliases) + return self._planned(selected_tasks + selected_evals, selected_tasks, selected_evals, aliases, override) def plan_entities(self, ids: Sequence[str]) -> Plan: self._check_writers() self._check_loops() - return self._planned(self._select_entities(ids), [], [], {}) + return self._planned(self._select_entities(ids), [], [], {}, None) def _planned(self, roots: list[ResolvedEntry], selected_tasks: list[ResolvedEntry], - selected_evals: list[ResolvedEntry], aliases: dict[str, AliasUse]) -> Plan: + selected_evals: list[ResolvedEntry], aliases: dict[str, AliasUse], + override: AgentOverride | None) -> Plan: closure = self._closure(roots) store_refs, store_latest = self._check_store_refs(closure) if override is not None: From e1aef46d5b6ea8b5483dfabac0bc5d7b5d511716 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 13:42:28 -0700 Subject: [PATCH 17/39] test(env): the Modal load test builds its server from a context Modal can build Its image went to localhost:5000, which the Modal provider refuses, so the test failed before it ran whenever Modal keys were present. Co-Authored-By: Claude Opus 5.5 --- .../env/envs/server_env_modal_test.py | 21 ++++--------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/tst/integration/env/envs/server_env_modal_test.py b/tst/integration/env/envs/server_env_modal_test.py index 5c063da8a..2fef6904a 100644 --- a/tst/integration/env/envs/server_env_modal_test.py +++ b/tst/integration/env/envs/server_env_modal_test.py @@ -2,7 +2,8 @@ A load is timed by the size of the payload it staged, measured in the server's container through the deploy's own handle and a restored one; and the server's state reads as JSON whether ``data/get`` answers with data or with a file -bundle. Needs the resolved config to reach Modal and an image store Modal can pull from; skipped otherwise. +bundle. Needs the resolved config to reach Modal; skipped otherwise. The server's image is a build context, which +Modal builds where it runs. """ import asyncio @@ -25,7 +26,6 @@ from agent_env.task_step.context import TaskStepContext from agent_env.task_step.task_steps.teardown_sandboxes import TeardownSandboxesTaskStep from tst.util.capabilities import skip_without_remote_sandbox -from tst.util.image_cache import build_or_reuse pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_remote_sandbox("modal")] @@ -37,7 +37,7 @@ async def test_a_modal_load_is_timed_by_its_payload_and_the_servers_state_reads_ staged_sizes, timeout_for = [], constants.data_plane_load_timeout_s monkeypatch.setattr(constants, "data_plane_load_timeout_s", lambda size: staged_sizes.append(size) or timeout_for(size)) uid = uuid.uuid4().hex[:8] - env = MCPServerEnv.put(id=f"server-items-modal-{uid}", docker_image_artifact=_items_image(), + env = MCPServerEnv.put(id=f"server-items-modal-{uid}", docker_image_artifact=_items_context(uid), environment_name="items", env_provider_type="server") deployed = await env.deploy(sandbox_type="modal", ttl_seconds=900) try: @@ -138,23 +138,10 @@ def _items_context(uid: str) -> DockerImageArtifact: ignore=shutil.ignore_patterns("__pycache__")) for name in ("server.py", "Dockerfile", "seed.json"): shutil.copy(data / name, context / name) - return DockerImageArtifact.put_context(f"server-items-episode-image-{uid}", description="episode on Modal", + return DockerImageArtifact.put_context(f"server-items-modal-image-{uid}", description="server env on Modal", context_path=str(context), dockerfile_path="Dockerfile") -def _items_image() -> DockerImageArtifact: - """The in-memory items server (card 'items'), built for linux/amd64 with the vendored agentenv_protocol.""" - data = _REPO / "tst" / "data" / "agentenv_mcp" - with tempfile.TemporaryDirectory() as build_dir: - context = Path(build_dir) - shutil.copytree(_REPO / "packages" / "agentenv-protocol" / "src" / "agentenv_protocol", context / "agentenv_protocol", - ignore=shutil.ignore_patterns("__pycache__")) - for name in ("server.py", "Dockerfile", "seed.json"): - shutil.copy(data / name, context / name) - return build_or_reuse(artifact_id="server-items-modal", description="server env on Modal", - dockerfile=context / "Dockerfile", context=context, tag="server-items-modal") - - def _items_artifact(uid: str) -> EnvironmentArtifact: return EnvironmentArtifact.put( id=f"server-items-modal-data-{uid}", From c9ae1fba9b94f2df150e173a05c0c4f0a5990615 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 15:06:01 -0700 Subject: [PATCH 18/39] fix(agents): the proxy needs protocol 0.1.370, and Gemini and OpenCode declare where they work on a VM host #280 shipped in protocol 0.1.370; 0.1.367 to 0.1.369 came out before it and have no proxy/v1. Gemini and OpenCode, new on main, install on a VM host as root, in /workspace and /app, so a proxy installed in their place works there. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/agents/agents/gemini/a2a_server.py | 3 +++ src/agent_env/agents/agents/opencode/a2a_server.py | 3 +++ src/agent_env/agents/agents/proxy/requirements.txt | 2 +- tst/unit/agents/gemini/test_card_contract.py | 2 ++ tst/unit/agents/opencode/test_card_contract.py | 1 + 5 files changed, 10 insertions(+), 1 deletion(-) diff --git a/src/agent_env/agents/agents/gemini/a2a_server.py b/src/agent_env/agents/agents/gemini/a2a_server.py index 702dab112..e70fccd9c 100644 --- a/src/agent_env/agents/agents/gemini/a2a_server.py +++ b/src/agent_env/agents/agents/gemini/a2a_server.py @@ -403,6 +403,9 @@ async def _vertex_openai_bridge( install_commands_host=_INSTALL_COMMANDS_HOST, a2a_port=8000, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], + # Where an install on the host works, and as whom, for a proxy installed in this agent's place. + workdir_host="/workspace", + user_host="root", ), enable( ATTRIBUTION_PROBE_V1, diff --git a/src/agent_env/agents/agents/opencode/a2a_server.py b/src/agent_env/agents/agents/opencode/a2a_server.py index 82ae12383..f023b833e 100644 --- a/src/agent_env/agents/agents/opencode/a2a_server.py +++ b/src/agent_env/agents/agents/opencode/a2a_server.py @@ -291,6 +291,9 @@ async def lifespan(app: Starlette): required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], + # Where an install on the host works, and as whom, for a proxy installed in this agent's place. + workdir_host="/app", + user_host="root", ), # SDK default: the CLI reaches the task's other agents through the SDK's peer_agents MCP tools. PEER_AGENTS_V1, diff --git a/src/agent_env/agents/agents/proxy/requirements.txt b/src/agent_env/agents/agents/proxy/requirements.txt index c35eee2fe..dfa7130a1 100644 --- a/src/agent_env/agents/agents/proxy/requirements.txt +++ b/src/agent_env/agents/agents/proxy/requirements.txt @@ -1,3 +1,3 @@ # The image installs the newest agentenv-framework-protocol at or above this floor; a release doesn't pin it, so the # image is built again only when this folder changes. The floor is the first that ships ProxyAgent and proxy/v1. -agentenv-framework-protocol[agent]>=0.1.367 +agentenv-framework-protocol[agent]>=0.1.370 diff --git a/tst/unit/agents/gemini/test_card_contract.py b/tst/unit/agents/gemini/test_card_contract.py index 0f2c073e7..6d28f35aa 100644 --- a/tst/unit/agents/gemini/test_card_contract.py +++ b/tst/unit/agents/gemini/test_card_contract.py @@ -201,6 +201,7 @@ def test_install_params(app): params = _extensions(app)["urn:agentenv:install/v1"]["params"] assert list(params) == [ "install_commands", "required_params", "install_commands_host", "a2a_port", "required_params_host", + "workdir_host", "user_host", ] assert params["a2a_port"] == 8000 endpoint = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] @@ -208,6 +209,7 @@ def test_install_params(app): assert params["required_params_host"] == ["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *endpoint] assert params["install_commands"] == a2a_server._INSTALL_COMMANDS assert params["install_commands_host"] == a2a_server._INSTALL_COMMANDS_HOST + assert (params["workdir_host"], params["user_host"]) == ("/workspace", "root") def test_the_card_and_health_are_served(client): diff --git a/tst/unit/agents/opencode/test_card_contract.py b/tst/unit/agents/opencode/test_card_contract.py index b289a83d8..94582d36d 100644 --- a/tst/unit/agents/opencode/test_card_contract.py +++ b/tst/unit/agents/opencode/test_card_contract.py @@ -100,6 +100,7 @@ def test_install_mode_reads_the_card_by_importing_agent_card(): install = extensions["urn:agentenv:install/v1"]["params"] assert install["install_commands"] == a2a_server._INSTALL_COMMANDS assert install["install_commands_host"] == a2a_server._INSTALL_COMMANDS_HOST + assert (install["workdir_host"], install["user_host"]) == ("/app", "root") assert extensions["urn:agentenv:staging/v1"]["params"]["endpoint"] == "/ext/staging" From d4826013e874b02c30ec41492d09aab76365ec87 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 15:38:49 -0700 Subject: [PATCH 19/39] fix(task_steps): verify_sandbox and collect_artifacts reach an agent installed on a VM's host Aimed at an agent by name on a VM, both looked for the agent's container, which an agent install_agent put on the host doesn't have, so they failed. They now work on the host when the agent is on_host, as run_code does. Co-Authored-By: Claude Opus 5.5 --- .../task_step/task_steps/collect_artifacts.py | 5 ++++- .../task_steps/verifiers/verify_sandbox.py | 11 ++++++++--- tst/integration/agent/proxy_modal_vm_test.py | 5 ++++- .../test_container_sandbox_routing.py | 19 ++++++++++--------- 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index 153f4e733..efa3b3d36 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -778,7 +778,10 @@ async def _collect_via_agent_container(self, context, store, artifact_id, versio sandbox = await self._resolve_live_sandbox(provider, agent.sandbox_id) logger.info(f"Connected to sandbox {agent.sandbox_id} (mode={sandbox.mode})") - container = await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM else None + # An agent installed on a VM's host has no container: its files are on the host. + container = ( + await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host else None + ) if container: logger.info(f"Using agent container: {container}") diff --git a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py index 7290dbfa7..916a7887f 100644 --- a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py +++ b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py @@ -120,12 +120,17 @@ def from_dict(cls, data: dict) -> "VerifySandboxTaskStep": async def execute(self, context: TaskStepContext) -> TaskStepContext: on_agent = self.sandbox_name is None if on_agent: - sandbox = await self._agent_sandbox(context) + agent, sandbox = await self._agent_sandbox(context) else: deployed, sandbox = await self._deployed_sandbox(context) logger.info(f"Connected to sandbox {sandbox.sandbox_id} (mode={sandbox.mode})") if on_agent: - container = await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM else None + # An agent installed on a VM's host has no container: its files are on the host. + container = ( + await find_agent_container(sandbox) + if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host + else None + ) else: # A container-mode sandbox on a VM-backed provider runs its image beside the host, and reattaches as a VM: # its files are reached in the container, decided by the mode deploy_sandbox recorded. @@ -196,7 +201,7 @@ async def _agent_sandbox(self, context: TaskStepContext): if agent.sandbox_type else get_agent_sandbox_provider() ) - return await provider.get_sandbox(agent.sandbox_id) + return agent, await provider.get_sandbox(agent.sandbox_id) async def _deployed_sandbox(self, context: TaskStepContext): deployed = next( diff --git a/tst/integration/agent/proxy_modal_vm_test.py b/tst/integration/agent/proxy_modal_vm_test.py index 264504b48..d1f8d9fa8 100644 --- a/tst/integration/agent/proxy_modal_vm_test.py +++ b/tst/integration/agent/proxy_modal_vm_test.py @@ -50,6 +50,9 @@ async def test_an_agent_installed_on_a_vm_host_is_played_through_a_proxy_install {"id": "check", "type": "verify_sandbox", "sandbox_name": "vm", "base_dir": "/opt/agent-solver/workspace", "criteria": [{"type": "probe_file_contains", "criterion": "the greeting is copied", "paths": ["outputs/answer.txt"], "expected": "^hello$"}]}, + {"id": "check-agent", "type": "verify_sandbox", "agent_name": "solver", "base_dir": "/opt/agent-solver/workspace", + "criteria": [{"type": "probe_file_exists", "criterion": "the answer is on the agent's host", + "paths": ["outputs/answer.txt"]}]}, ]}) episode = await EpisodeSpec.from_task(task).start() @@ -61,7 +64,7 @@ async def test_an_agent_installed_on_a_vm_host_is_played_through_a_proxy_install "/opt/agent-solver/workspace\nagent-solver\nhello") await tools.call("bash", {"command": "mkdir -p outputs && cp greeting.txt outputs/answer.txt"}) await turn.end("Copied.") - assert await episode.result() == {"check": 1.0} + assert await episode.result() == {"check": 1.0, "check-agent": 1.0} finally: await episode.abandon("the test ended") [vm] = Task.get_instance(episode.id).context["deployed_sandboxes"] diff --git a/tst/unit/task_step/test_container_sandbox_routing.py b/tst/unit/task_step/test_container_sandbox_routing.py index 0c6bb92d2..799ccd783 100644 --- a/tst/unit/task_step/test_container_sandbox_routing.py +++ b/tst/unit/task_step/test_container_sandbox_routing.py @@ -2,7 +2,8 @@ A container sandbox on a container provider (Modal) is the runtime itself; one on a VM-backed provider (local, modal_vm, e2b) runs beside its host, so its files are only reachable through `docker exec `. A VM-backed provider -reattaches every sandbox as a VM, so a deploy_sandbox sandbox's mode is the one deploy_sandbox recorded. +reattaches every sandbox as a VM, so a deploy_sandbox sandbox's mode is the one deploy_sandbox recorded. An agent +install_agent put on a VM's host has no container, so both work on the host. """ from __future__ import annotations @@ -39,9 +40,9 @@ def _provide(monkeypatch, sandbox): monkeypatch.setattr(module, "get_sandbox_provider", lambda: provider) -def _agent_context(): +def _agent_context(on_host: bool = False): ctx = TaskStepContext() - ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1")] + ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1", on_host=on_host)] return ctx @@ -81,11 +82,11 @@ async def _eval(self, sandbox, container, criterion): @pytest.mark.asyncio -@pytest.mark.parametrize("kind,expected", [("vm", CONTAINER), ("container", None)]) -async def test_collect_reaches_an_agents_container(monkeypatch, kind, expected): +@pytest.mark.parametrize("kind,on_host,expected", [("vm", False, CONTAINER), ("vm", True, None), ("container", False, None)]) +async def test_collect_reaches_an_agents_container(monkeypatch, kind, on_host, expected): _provide(monkeypatch, _sandbox(kind)) step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver", artifact_paths=["out.txt"]) - assert await _collect_container(monkeypatch, step, _agent_context()) == expected + assert await _collect_container(monkeypatch, step, _agent_context(on_host)) == expected @pytest.mark.asyncio @@ -102,11 +103,11 @@ async def test_collect_reaches_a_named_sandboxs_container(monkeypatch, kind, rec @pytest.mark.asyncio -@pytest.mark.parametrize("kind,expected", [("vm", CONTAINER), ("container", None)]) -async def test_verify_sandbox_probes_an_agents_container(monkeypatch, kind, expected): +@pytest.mark.parametrize("kind,on_host,expected", [("vm", False, CONTAINER), ("vm", True, None), ("container", False, None)]) +async def test_verify_sandbox_probes_an_agents_container(monkeypatch, kind, on_host, expected): _provide(monkeypatch, _sandbox(kind)) step = VerifySandboxTaskStep(id="v", version=None, agent_name="solver", criteria=_CRITERIA) - assert await _verify_container(monkeypatch, step, _agent_context()) == expected + assert await _verify_container(monkeypatch, step, _agent_context(on_host)) == expected @pytest.mark.asyncio From 11523f212e37c26450d2030553d597dbb3e5e152 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 15:54:01 -0700 Subject: [PATCH 20/39] fix(protocol): the proxy's tool deadlines hold on Python 3.10 On 3.10, asyncio.wait_for raises asyncio.TimeoutError, which isn't the built-in TimeoutError until 3.11, so a tool call or listing that timed out escaped as a server error. A proxy installed on a VM host runs on the host's own Python, which can be 3.10. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/proxy.py | 4 ++-- packages/agentenv-protocol/tests/test_a2a_proxy.py | 14 ++++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 5d33b6955..401cad18e 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -285,7 +285,7 @@ async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: for server in turn.servers: try: listed = await asyncio.wait_for(self._listed(turn, server), TOOL_TIMEOUT_SECONDS) - except TimeoutError: + except asyncio.TimeoutError: raise HTTPException(status_code=502, detail=( f"env server {server!r} didn't list its tools within {TOOL_TIMEOUT_SECONDS}s")) from None except (ToolSessionError, httpx.HTTPError) as e: @@ -339,7 +339,7 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) try: called = await asyncio.wait_for(self._called(turn, action), timeout) result = ProxyToolResult(content=called.content, is_error=called.isError) - except TimeoutError: + except asyncio.TimeoutError: result = _tool_error(f"the tool call timed out after {timeout:g}s") except (ToolSessionError, httpx.HTTPError) as e: result = _tool_error(f"{type(e).__name__}: {e}") diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 478bab14b..5bfc52dee 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -283,6 +283,20 @@ def test_a_tool_actions_deadline_covers_connecting_to_its_server(client): assert time.monotonic() - started < 5 +def test_a_server_that_never_lists_its_tools_is_a_bad_gateway(client, monkeypatch): + monkeypatch.setattr(proxy, "TOOL_TIMEOUT_SECONDS", 0.5) + with socket.socket() as silent: + silent.bind(("127.0.0.1", 0)) + silent.listen() + _register(client, f"http://127.0.0.1:{silent.getsockname()[1]}/mcp") + task_id = _send(client, "go") + + response = _proxy(client, "tools", task_id=task_id) + + assert response.status_code == 502 + assert "didn't list its tools within 0.5s" in response.text + + def test_a_turns_actions_share_one_connection_to_each_env_server_which_closes_with_the_turn(client, env, monkeypatch): opened, closed, results = [], [], [] open_session, close_session = ToolSession.open, ToolSession.close From b63403642b07b4a9f04f7c7986a3ae3f4095b541 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 22:55:49 -0700 Subject: [PATCH 21/39] fix(protocol): the proxy answers a stopped action, parses its own results, renders, and matches calls by what runs - An action polled while its turn closes is answered 404, not a bare 500 from the CancelledError; a command that can't start (an argument list too long) is a failed bash result, exit code 126, with its span, not a 500 on every poll. - ProxyBashResult.exit_code defaults to None: a stopped command's result leaves it out, which the published model refused. - The proxy's spans carry an OTel status, ERROR for a failed action or turn, and the explorer reads a span with none; it threw on every proxy trajectory with an action. - A recorded bash call is matched by its command alone, since an action carries nothing else, and a whole float equals the int a driver writes for it; either mismatch left the turn stuck. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 17 +++-- .../agentenv_protocol/a2a_agent/extensions.py | 2 +- .../src/agentenv_protocol/a2a_agent/proxy.py | 53 +++++++++---- .../agentenv-protocol/tests/test_a2a_proxy.py | 62 +++++++++++++-- .../src/lib/parse-trajectory-proxy.smoke.ts | 75 +++++++++++++++++++ .../explorer/ui/src/lib/parse-trajectory.ts | 6 +- 6 files changed, 183 insertions(+), 32 deletions(-) create mode 100644 src/agent_env/explorer/ui/src/lib/parse-trajectory-proxy.smoke.ts diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 76ad0b97d..79ce0d712 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -719,14 +719,17 @@ operations: and headers registered for it, and `{"kind": "bash", "command"}` runs `bash -c` in the proxy's folder, as its user, in a process group of its own; either takes `"timeout_seconds"` (3600 for `bash`, 600 for a tool by default). The driver picks - `action_id`, so a retried `act` returns the action it already started. It answers - `{"done", "result"}` at once when the action ends within the long-poll; `result` - (`{"token", "action_id"}`) long-polls one that doesn't. A `bash` result is - `{"exit_code", "stdout", "stderr", "timed_out"}`, a tool's `{"content", "is_error"}`. + `action_id`, unique for the proxy, so a retried `act` returns the action it already + started. It answers `{"done", "result"}` at once when the action ends within the long-poll; + `result` (`{"token", "action_id"}`) long-polls one that doesn't, and gets 404 once the + action's turn has closed. A `bash` result is `{"exit_code", "stdout", "stderr", + "timed_out"}`, with no `exit_code` when the command was stopped and 126 when it couldn't + start; a tool's is `{"content", "is_error"}`. - `record` (`{"token", "task_id", "message", "model", "usage"}`) records one model call: an OpenAI chat message, an Anthropic message or a string; anything else gets 400. Once a turn records a call, each action must be one the last recorded call made, by name and - arguments, and all of them must run before the next `record` or the `reply`, else 409. + arguments (a `bash` call by its command alone, a whole float equal to its int), and all of + them must run before the next `record` or the `reply`, else 409. Calls to tools the turn doesn't offer aren't checked. - `reply` (`{"token", "task_id", "parts"}`, or `"error"` in place of `"parts"`) ends the turn: parts complete the task, and an error fails it as `proxy.driver_error`. @@ -735,7 +738,9 @@ A turn's actions belong to it: when it closes, by reply or withdrawal, those sti are stopped, their whole process groups killed, and its results forgotten; the app's lifespan stops what runs when it shuts down. The proxy writes the turn's trajectory as it goes, readable live, as OTel GenAI spans in format `agentenv-proxy/v1`: `chat` for each -recorded call, `execute_tool` for each action, and `chain` for the turn. Skills it is given +recorded call, `execute_tool` for each action, and `chain` for the turn, each with an +OTel status: `ERROR` for an action that failed, a command that exited nonzero or timed out, +and a turn that ended in an error. Skills it is given are written to `~/skills//`, once each, and a turn names each one's folder. `ProxyAgent(shell=False)` is a proxy for an agent that works in no sandbox: it offers no `bash` and keeps no skills. It holds the env servers' headers as any agent does, so run it diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py index 8e272050f..5a044b655 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/extensions.py @@ -367,7 +367,7 @@ class ProxyResultRequest(ExtensionRequest): class ProxyBashResult(ExtensionResponse): kind: Literal["bash"] = "bash" - exit_code: int | None # none when the command was stopped + exit_code: int | None = None # none when the command was stopped stdout: str stderr: str timed_out: bool diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 401cad18e..47c053319 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -164,7 +164,7 @@ async def run(self, request: TaskRequest[ProxyConfig]) -> TaskResult: prompt = {"prompt": "\n".join(part.text for part in turn.parts), "system_prompt": turn.system_prompt} if reply.error is not None: turn.trajectory.append(_span("turn", "chain", turn.started, turn.model, prompt=prompt, - completion={"error": reply.error})) + completion={"error": reply.error}, failed=True)) return TaskResult.failure("proxy.driver_error", reply.error) turn.trajectory.append(_span("turn", "chain", turn.started, turn.model, prompt=prompt, completion={ "content": [{"type": "text", "text": part.text} for part in reply.parts]})) @@ -202,14 +202,14 @@ async def act(self, request: ProxyActRequest) -> ProxyActionState: if (known := self._actions.get(request.action_id)) is not None: if known[0].task_id != request.task_id: raise HTTPException(status_code=409, detail=f"action {request.action_id!r} is another turn's") - return await self._state(known[1]) + return await self._state(request.action_id, known[1]) turn = self._open_turn(request.token, request.task_id) action = request.action if isinstance(action, ProxyBashAction) and not self._shell: raise HTTPException(status_code=409, detail="this proxy works in no sandbox, so it runs no commands") if turn.recorded: - call = (action.name, _canonical(action.arguments)) if isinstance(action, ProxyToolAction) else ( - BASH, _canonical({"command": action.command})) + call = (_call(action.name, action.arguments) if isinstance(action, ProxyToolAction) + else _call(BASH, {"command": action.command})) if call not in turn.pending: raise HTTPException(status_code=409, detail=f"the last recorded model call didn't make this call to " f"{call[0]!r}") @@ -217,14 +217,14 @@ async def act(self, request: ProxyActRequest) -> ProxyActionState: task = asyncio.create_task(self._perform(turn, action)) turn.actions[request.action_id] = task self._actions[request.action_id] = (turn, task) - return await self._state(task) + return await self._state(request.action_id, task) @extension(PROXY_V1.result) async def result(self, request: ProxyResultRequest) -> ProxyActionState: self._check(request.token) if (known := self._actions.get(request.action_id)) is None: raise HTTPException(status_code=404, detail=f"no action {request.action_id!r}") - return await self._state(known[1]) + return await self._state(request.action_id, known[1]) @extension(PROXY_V1.record) async def record(self, request: ProxyRecordRequest) -> None: @@ -233,7 +233,7 @@ async def record(self, request: ProxyRecordRequest) -> None: turn = self._open_turn(request.token, request.task_id) if turn.pending: raise HTTPException(status_code=409, detail=_unrun(turn)) - turn.pending = [(name, _canonical(arguments)) for name, arguments in calls if name in offered] + turn.pending = [_call(name, arguments) for name, arguments in calls if name in offered] turn.recorded = True turn.model = request.model or turn.model span = _span("chat", "chat", turn.last_event, turn.model, completion={"content": content}) @@ -328,11 +328,17 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) started = _now() if isinstance(action, ProxyBashAction): timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS - result = await _run(action.command, timeout) + try: + result = await _run(action.command, timeout) + except OSError as e: + # As a shell reports a command it can't execute: an argument list too long, say. + result = ProxyBashResult(exit_code=126, stdout="", stderr=f"error: the command couldn't start: {e}\n", + timed_out=False) name, arguments = BASH, {"command": action.command} output = (f"error: the command timed out after {timeout:g}s" if result.timed_out else result.stdout + result.stderr) completion = {"output": output, "exit_code": result.exit_code} + failed = result.timed_out or result.exit_code != 0 else: timeout = action.timeout_seconds or TOOL_TIMEOUT_SECONDS name, arguments = action.name, action.arguments @@ -345,15 +351,19 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) result = _tool_error(f"{type(e).__name__}: {e}") text = "\n".join(block["text"] for block in result.content if block.get("type") == "text") completion = {"error": text} if result.is_error else {"output": text} - span = _span(name, "execute_tool", started, None, prompt={"input": arguments}, completion=completion) + failed = result.is_error + span = _span(name, "execute_tool", started, None, prompt={"input": arguments}, completion=completion, + failed=failed) turn.trajectory.append(span) turn.last_event = span["end_time"] return result - async def _state(self, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> ProxyActionState: + async def _state(self, action_id: str, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> ProxyActionState: await asyncio.wait((task,), timeout=self._long_poll_seconds) if not task.done(): return ProxyActionState(done=False) + if task.cancelled(): + raise HTTPException(status_code=404, detail=f"action {action_id!r} was stopped: its turn is over") return ProxyActionState(done=True, result=task.result()) async def _close(self, turn: _Turn) -> None: @@ -471,8 +481,22 @@ def _unrun(turn: _Turn) -> str: return f"the last recorded model call's calls to {sorted({name for name, _ in turn.pending})} haven't run" -def _canonical(arguments: dict[str, Any]) -> str: - return json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str) +def _call(name: str, arguments: dict[str, Any]) -> tuple[str, str]: + """A call as recorded calls and actions are compared: ``bash`` by its command alone, since an action carries + nothing else, and a whole float equal to the int a driver may have written for it.""" + if name == BASH: + arguments = {"command": arguments.get("command")} + return name, json.dumps(_whole(arguments), sort_keys=True, separators=(",", ":"), default=str) + + +def _whole(value: Any) -> Any: + if isinstance(value, float) and value.is_integer(): + return int(value) + if isinstance(value, dict): + return {key: _whole(item) for key, item in value.items()} + if isinstance(value, list): + return [_whole(item) for item in value] + return value def _tool_error(message: str) -> ProxyToolResult: @@ -480,7 +504,7 @@ def _tool_error(message: str) -> ProxyToolResult: def _span(name: str, operation: str, started: str, model: str | None, *, prompt: Any = None, - completion: Any = None) -> dict[str, Any]: + completion: Any = None, failed: bool = False) -> dict[str, Any]: """An OTel GenAI span, the trajectory format rubrics and the explorer read.""" attributes: dict[str, Any] = {"gen_ai.operation.name": operation} if model: @@ -489,7 +513,8 @@ def _span(name: str, operation: str, started: str, model: str | None, *, prompt: attributes["gen_ai.prompt"] = json.dumps(prompt, default=str) if completion is not None: attributes["gen_ai.completion"] = json.dumps(completion, default=str) - return {"name": name, "start_time": started, "end_time": _now(), "attributes": attributes} + return {"name": name, "start_time": started, "end_time": _now(), + "status": {"status_code": "ERROR" if failed else "OK"}, "attributes": attributes} def _now() -> str: diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 5bfc52dee..e1f9ff0f4 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -13,7 +13,7 @@ from starlette.testclient import TestClient from agentenv_protocol import EnvironmentCard, ToolContext, ToolSession, create_fastmcp_app, tool -from agentenv_protocol.a2a_agent import proxy +from agentenv_protocol.a2a_agent import ProxyActionState, proxy from agentenv_protocol.a2a_agent.proxy import ProxyAgent TOKEN = "t" * 32 @@ -34,6 +34,11 @@ def explode(self) -> str: """Always fails.""" raise RuntimeError("the tool blew up") + @tool() + def count(self, n: int) -> str: + """Echo a number.""" + return str(n) + @tool() def whoami(self, ctx: ToolContext) -> str: """Who is calling.""" @@ -198,6 +203,7 @@ def test_an_error_reply_fails_the_turn(client): task = _settled(client, task_id) assert task["status"]["state"] == "failed" assert "the driver gave up" in task["status"]["message"]["parts"][0]["text"] + assert _trajectory(client, task_id)[-1]["status"] == {"status_code": "ERROR"} def test_a_reply_names_exactly_one_open_turn_and_carries_parts_or_an_error(client): @@ -237,7 +243,7 @@ def test_a_turn_offers_its_env_servers_tools_and_bash(client, env): tools = {item["name"]: item for item in _proxy(client, "tools", task_id=task_id).json()["tools"]} - assert set(tools) == {"add_item", "explode", "whoami", "bash"} + assert set(tools) == {"add_item", "count", "explode", "whoami", "bash"} assert (tools["add_item"]["server"], "server" in tools["bash"]) == ("items", False) assert tools["add_item"]["input_schema"]["properties"]["item"]["type"] == "string" @@ -366,10 +372,44 @@ def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(cli left = _done(client, task_id, "l", _bash("sleep 60 & echo $! > left", timeout_seconds=0.5)) assert waited == left == {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True} + assert ProxyActionState.model_validate({"done": True, "result": waited}).result.exit_code is None _stopped(int((tmp_path / "waited").read_text())) _stopped(int((tmp_path / "left").read_text())) +def test_a_command_too_long_to_start_fails_as_a_shell_reports_it(client): + task_id = _send(client, "go") + + result = _done(client, task_id, "huge", _bash("true " + "x" * 3_000_000)) + + assert (result["exit_code"], result["timed_out"], result["stdout"]) == (126, False, "") + assert "couldn't start" in result["stderr"] + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) + _settled(client, task_id) + [action, _] = _trajectory(client, task_id) + assert (action["name"], action["status"]) == ("bash", {"status_code": "ERROR"}) + + +def test_an_action_polled_while_its_turn_closes_is_answered_404(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + agent = ProxyAgent(long_poll_seconds=10) + with TestClient(agent.create_app()) as client: + client.post("/ext/proxy/claim", json={"token": TOKEN}) + task_id = _send(client, "go") + polled = {} + poll = threading.Thread(target=lambda: polled.update(response=_proxy( + client, "act", task_id=task_id, action_id="long", action=_bash("sleep 30")))) + poll.start() + while "long" not in agent._actions: + time.sleep(0.01) + + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) + poll.join(10) + + assert polled["response"].status_code == 404 + assert "'long' was stopped: its turn is over" in polled["response"].text + + def test_a_turn_that_closes_stops_its_running_actions_and_forgets_them(client, tmp_path): task_id = _send(client, "go") assert _proxy(client, "act", task_id=task_id, action_id="long", @@ -400,8 +440,10 @@ def test_recorded_model_calls_are_checked_against_the_actions_that_run(client, e message = {"role": "assistant", "content": "Adding it.", "reasoning_content": "The user wants a crate.", "tool_calls": [ {"id": "c1", "type": "function", "function": {"name": "add_item", "arguments": '{"item": "crate"}'}}, - {"id": "c2", "type": "function", "function": {"name": "bash", "arguments": '{"command": "true"}'}}, - {"id": "c3", "type": "function", "function": {"name": "my_own_tool", "arguments": "{}"}}]} + {"id": "c2", "type": "function", "function": {"name": "bash", + "arguments": '{"command": "true", "timeout": 30}'}}, + {"id": "c3", "type": "function", "function": {"name": "my_own_tool", "arguments": "{}"}}, + {"id": "c4", "type": "function", "function": {"name": "count", "arguments": '{"n": 2.0}'}}]} assert _proxy(client, "record", task_id=task_id, message=message, model="qwen", usage={"prompt_tokens": 12, "completion_tokens": 7}).status_code == 200 @@ -412,22 +454,26 @@ def test_recorded_model_calls_are_checked_against_the_actions_that_run(client, e assert _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "x"}]).status_code == 409 _done(client, task_id, "b", _bash("true")) _done(client, task_id, "a", {"kind": "tool", "server": "items", "name": "add_item", "arguments": {"item": "crate"}}) + _done(client, task_id, "n", {"kind": "tool", "server": "items", "name": "count", "arguments": {"n": 2}}) assert _proxy(client, "record", task_id=task_id, message="Done.").status_code == 200 _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "Done."}]) task = _settled(client, task_id) spans = _trajectory(client, task_id) assert [(span["name"], span["attributes"]["gen_ai.operation.name"]) for span in spans] == [ - ("chat", "chat"), ("bash", "execute_tool"), ("add_item", "execute_tool"), ("chat", "chat"), ("turn", "chain")] + ("chat", "chat"), ("bash", "execute_tool"), ("add_item", "execute_tool"), ("count", "execute_tool"), + ("chat", "chat"), ("turn", "chain")] + assert {span["status"]["status_code"] for span in spans} == {"OK"} assert _attrs(spans[0], "gen_ai.completion")["content"] == [ {"type": "thinking", "thinking": "The user wants a crate."}, {"type": "text", "text": "Adding it."}, {"type": "tool_use", "id": "c1", "name": "add_item", "input": {"item": "crate"}}, - {"type": "tool_use", "id": "c2", "name": "bash", "input": {"command": "true"}}, - {"type": "tool_use", "id": "c3", "name": "my_own_tool", "input": {}}] + {"type": "tool_use", "id": "c2", "name": "bash", "input": {"command": "true", "timeout": 30}}, + {"type": "tool_use", "id": "c3", "name": "my_own_tool", "input": {}}, + {"type": "tool_use", "id": "c4", "name": "count", "input": {"n": 2.0}}] assert (spans[0]["attributes"]["gen_ai.request.model"], spans[0]["attributes"]["gen_ai.usage.input_tokens"], spans[0]["attributes"]["gen_ai.usage.output_tokens"]) == ("qwen", 12, 7) assert _attrs(spans[2], "gen_ai.prompt") == {"input": {"item": "crate"}} - assert task["status"]["message"]["parts"][1]["data"]["usage"]["tool_call_count"] == 2 + assert task["status"]["message"]["parts"][1]["data"]["usage"]["tool_call_count"] == 3 def test_two_model_calls_recorded_at_once_are_held_one_after_the_other(client, monkeypatch): diff --git a/src/agent_env/explorer/ui/src/lib/parse-trajectory-proxy.smoke.ts b/src/agent_env/explorer/ui/src/lib/parse-trajectory-proxy.smoke.ts new file mode 100644 index 000000000..539de82e1 --- /dev/null +++ b/src/agent_env/explorer/ui/src/lib/parse-trajectory-proxy.smoke.ts @@ -0,0 +1,75 @@ +/** + * Smoke test for a ProxyAgent trajectory in `parseOtelTrajectory`: the spans the proxy writes, `chat` for each + * recorded model call, `execute_tool` for each action and `chain` for the turn, with none of an exporter's span + * context. A span with no status still renders. + * + * Runner: plain TS, exits non-zero on assertion failure. From this package: + * npx tsx src/lib/parse-trajectory-proxy.smoke.ts + */ +import { type OtelSpan, type ToolCallEvent, parseOtelTrajectory } from './parse-trajectory'; + +let failures = 0; +function assert(cond: unknown, msg: string): void { + if (cond) { + console.log(`✓ ${msg}`); + } else { + failures += 1; + console.error(`✗ ${msg}`); + } +} + +function span(name: string, operation: string, at: string, attributes: Record, status?: string) { + const attrs: Record = { 'gen_ai.operation.name': operation }; + for (const [key, value] of Object.entries(attributes)) { + attrs[key] = typeof value === 'string' ? value : JSON.stringify(value); + } + return { + name, + start_time: at, + end_time: at, + ...(status ? { status: { status_code: status } } : {}), + attributes: attrs, + }; +} + +const spans = [ + span('chat', 'chat', '2026-10-23T10:00:01+00:00', { + 'gen_ai.request.model': 'qwen', + 'gen_ai.completion': { content: [ + { type: 'text', text: 'Copying it.' }, + { type: 'tool_use', id: 'c1', name: 'bash', input: { command: 'cp a b' } }, + { type: 'tool_use', id: 'c2', name: 'items_add_item', input: { item: 'crate' } }, + ] }, + }, 'OK'), + span('bash', 'execute_tool', '2026-10-23T10:00:02+00:00', { + 'gen_ai.prompt': { input: { command: 'cp a b' } }, + 'gen_ai.completion': { output: 'cp: a: No such file or directory\n', exit_code: 1 }, + }, 'ERROR'), + span('items_add_item', 'execute_tool', '2026-10-23T10:00:03+00:00', { + 'gen_ai.prompt': { input: { item: 'crate' } }, + 'gen_ai.completion': { output: '{"count": 1}' }, + }), + span('turn', 'chain', '2026-10-23T10:00:00+00:00', { + 'gen_ai.request.model': 'qwen', + 'gen_ai.prompt': { prompt: 'Copy a to b and add a crate.', system_prompt: null }, + 'gen_ai.completion': { content: [{ type: 'text', text: 'Added a crate; a was missing.' }] }, + }, 'OK'), +]; + +const parsed = parseOtelTrajectory(spans as unknown as OtelSpan[]); +const calls = parsed.events.filter((e): e is ToolCallEvent => e.type === 'tool_call'); + +assert(parsed.model === 'qwen', 'the model comes from the turn'); +assert(parsed.userPrompt === 'Copy a to b and add a crate.', 'the prompt comes from the turn'); +assert(parsed.finalResponse === 'Added a crate; a was missing.', 'the reply comes from the turn'); +assert(calls.length === 2 && parsed.toolCallCount === 2, 'each recorded call is a tool call'); +const [bash, add] = calls; +assert(bash?.result?.isError === true, 'an action the proxy marked ERROR is an error'); +assert(add?.result?.output.includes('"count": 1') && add.result.isError === false, + 'an action with no status renders'); + +if (failures > 0) { + console.error(`\n${failures} assertion(s) failed`); + process.exit(1); +} +console.log('\nall proxy trajectory assertions passed'); diff --git a/src/agent_env/explorer/ui/src/lib/parse-trajectory.ts b/src/agent_env/explorer/ui/src/lib/parse-trajectory.ts index 43f7777d0..3130e0b71 100644 --- a/src/agent_env/explorer/ui/src/lib/parse-trajectory.ts +++ b/src/agent_env/explorer/ui/src/lib/parse-trajectory.ts @@ -46,7 +46,7 @@ export interface OtelSpan { parent_id: string | null; start_time: string; end_time: string; - status: { status_code: string }; + status?: { status_code: string }; attributes: Record; events: unknown[]; links: unknown[]; @@ -410,7 +410,7 @@ export function parseOtelTrajectory( startTime: ts.start_time, endTime: ts.end_time, durationMs: durationMs(ts.start_time, ts.end_time), - isError: ts.status.status_code === 'ERROR' || detectOutputError(output), + isError: ts.status?.status_code === 'ERROR' || detectOutputError(output), }); } @@ -622,7 +622,7 @@ function parseScreenshotTrajectory(spans: OtelSpan[]): ParsedTrajectory { startTime: ts.start_time, endTime: ts.end_time, durationMs: durationMs(ts.start_time, ts.end_time), - isError: ts.status.status_code === 'ERROR' || detectOutputError(output), + isError: ts.status?.status_code === 'ERROR' || detectOutputError(output), screenshot, }; }); From 326422a16edc27926e069b60e3b8af78c090e157 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Fri, 23 Oct 2026 23:08:54 -0700 Subject: [PATCH 22/39] fix(protocol): only the proxy's own bash is matched by its command alone An env tool that is itself named bash, which a proxy with no shell can offer, is held to all its recorded arguments. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/proxy.py | 13 ++++--- .../agentenv-protocol/tests/test_a2a_proxy.py | 39 ++++++++++++++++++- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 47c053319..1263ba776 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -209,7 +209,7 @@ async def act(self, request: ProxyActRequest) -> ProxyActionState: raise HTTPException(status_code=409, detail="this proxy works in no sandbox, so it runs no commands") if turn.recorded: call = (_call(action.name, action.arguments) if isinstance(action, ProxyToolAction) - else _call(BASH, {"command": action.command})) + else _call(BASH, {"command": action.command}, own_shell=True)) if call not in turn.pending: raise HTTPException(status_code=409, detail=f"the last recorded model call didn't make this call to " f"{call[0]!r}") @@ -233,7 +233,8 @@ async def record(self, request: ProxyRecordRequest) -> None: turn = self._open_turn(request.token, request.task_id) if turn.pending: raise HTTPException(status_code=409, detail=_unrun(turn)) - turn.pending = [_call(name, arguments) for name, arguments in calls if name in offered] + turn.pending = [_call(name, arguments, own_shell=offered[name].server is None) + for name, arguments in calls if name in offered] turn.recorded = True turn.model = request.model or turn.model span = _span("chat", "chat", turn.last_event, turn.model, completion={"content": content}) @@ -481,10 +482,10 @@ def _unrun(turn: _Turn) -> str: return f"the last recorded model call's calls to {sorted({name for name, _ in turn.pending})} haven't run" -def _call(name: str, arguments: dict[str, Any]) -> tuple[str, str]: - """A call as recorded calls and actions are compared: ``bash`` by its command alone, since an action carries - nothing else, and a whole float equal to the int a driver may have written for it.""" - if name == BASH: +def _call(name: str, arguments: dict[str, Any], *, own_shell: bool = False) -> tuple[str, str]: + """A call as recorded calls and actions are compared: the proxy's own ``bash`` by its command alone, since its + action carries nothing else, and a whole float equal to the int a driver may have written for it.""" + if own_shell: arguments = {"command": arguments.get("command")} return name, json.dumps(_whole(arguments), sort_keys=True, separators=(",", ":"), default=str) diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index e1f9ff0f4..388ee463f 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -45,11 +45,27 @@ def whoami(self, ctx: ToolContext) -> str: return json.dumps({"role": ctx.caller.role}) +class _Terminal: + @tool() + def bash(self, command: str, cwd: str = "/") -> str: + """Run a command in the env's own terminal.""" + return f"{cwd}$ {command}" + + @pytest.fixture(scope="module") def env(): """An items env's MCP server on loopback, as a task's env serves one.""" - handler = _Items() - app = create_fastmcp_app(handler, card=EnvironmentCard(name="items")).streamable_http_app() + yield from _served(_Items(), "items") + + +@pytest.fixture(scope="module") +def terminal(): + """An env whose own tool is named ``bash``.""" + yield from _served(_Terminal(), "terminal") + + +def _served(handler: Any, name: str): + app = create_fastmcp_app(handler, card=EnvironmentCard(name=name)).streamable_http_app() sock = socket.socket() sock.bind(("127.0.0.1", 0)) server = uvicorn.Server(uvicorn.Config(app, log_level="warning")) @@ -476,6 +492,25 @@ def test_recorded_model_calls_are_checked_against_the_actions_that_run(client, e assert task["status"]["message"]["parts"][1]["data"]["usage"]["tool_call_count"] == 3 +def test_an_env_tool_named_bash_is_held_to_all_its_recorded_arguments(terminal, monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + with TestClient(ProxyAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: + client.post("/ext/proxy/claim", json={"token": TOKEN}) + _register(client, terminal[1], name="terminal") + task_id = _send(client, "go") + message = {"role": "assistant", "content": None, "tool_calls": [{"id": "c1", "type": "function", "function": { + "name": "bash", "arguments": '{"command": "ls", "cwd": "/a"}'}}]} + + assert _proxy(client, "record", task_id=task_id, message=message).status_code == 200 + moved = _proxy(client, "act", task_id=task_id, action_id="b", action={ + "kind": "tool", "server": "terminal", "name": "bash", "arguments": {"command": "ls", "cwd": "/b"}}) + ran = _done(client, task_id, "a", {"kind": "tool", "server": "terminal", "name": "bash", + "arguments": {"command": "ls", "cwd": "/a"}}) + + assert moved.status_code == 409 + assert ran["content"][0]["text"] == "/a$ ls" + + def test_two_model_calls_recorded_at_once_are_held_one_after_the_other(client, monkeypatch): listing, statuses = ProxyAgent._tools, [] From 2347de4857501251cbe3f3c2ad3c879c4ba11d0e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 24 Oct 2026 07:50:23 -0700 Subject: [PATCH 23/39] feat(task): an outside agent's prompt can be a conversation with a user simulator or counterpart Each message prompt_agent sends it is a turn of the same step, so an episode plays a task whose played agent converses over several turns, as Cooperation Bench tasks do. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/rules.py | 1 - tst/unit/task/episode/run_test.py | 39 +++++++++++++++++++++++++++++ tst/unit/task/episode/spec_test.py | 3 --- 3 files changed, 39 insertions(+), 4 deletions(-) diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index cc6bfaf3f..367931482 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -66,7 +66,6 @@ def _prompt_agent(step: PromptAgentTaskStep, outside: Outside) -> list[str]: return [ f"an outside agent's prompt can't have {what}" for what, present in ( ("parts, only a text prompt", step.parts is not None), - ("a user simulator (max_conversation_turns > 1)", step.max_conversation_turns > 1), ("output_format", step.output_format), ("agentenv_tools", step.agentenv_tools), ("snapshot_config", step.snapshot_config), diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index a1b87e33c..12cec2a5b 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -2,11 +2,15 @@ import asyncio import json +import threading +import time import httpx import pytest import pytest_asyncio +import uvicorn from agentenv_protocol import EnvironmentTool, ToolResult, ToolSessionError +from agentenv_protocol.a2a_agent import AgentEnvAgent, AgentIdentity, TaskRequest, TaskResult, a2a_agent from agentenv_protocol.a2a_agent import proxy as proxy_module import agent_env.task.episode.run as episode_run @@ -728,6 +732,41 @@ async def test_a_model_call_in_a_shape_the_proxy_cant_read_is_refused(teardowns) await turn.record_model_call({"role": "assistant", "content": {"not": "readable"}}) +@a2a_agent(identity=AgentIdentity(name="bob", description="Asks for one thing more.", version="1.0.0")) +class _Bob(AgentEnvAgent): + async def run(self, request: TaskRequest) -> TaskResult: + return TaskResult.text("And a box.") + + +@pytest.fixture +def bob(): + """A counterpart on loopback that the outside agent's prompt converses with.""" + server = uvicorn.Server(uvicorn.Config(_Bob().create_app(), host="127.0.0.1", port=0, log_level="warning")) + thread = threading.Thread(target=server.run, daemon=True) + thread.start() + while not server.started: + time.sleep(0.01) + yield f"http://127.0.0.1:{server.servers[0].sockets[0].getsockname()[1]}" + server.should_exit = True + thread.join(5) + + +@pytest.mark.asyncio +async def test_a_conversation_with_the_outside_agent_is_played_a_turn_per_message(teardowns, bob): + episode = await _spec(ITEMS, ALICE, {**TURN, "max_conversation_turns": 2, "user_a2a_url": bob}, REPLIES).start() + + first = await episode.next_turn() + await first.end("Added a crate.") + second = await episode.next_turn() + await second.end("Added a box.") + + assert [(turn.step_id, turn.prompt) for turn in (first, second)] == [ + ("turn", "Add a crate."), ("turn", "And a box.")] + assert await episode.result() == {"replies": 1.0} + [response] = Task.get_instance(episode.id).context["prompt_responses"] + assert response["response"] == "Added a box." + + @pytest.mark.asyncio async def test_a_prompt_continuing_a_conversation_gets_its_own_turn(teardowns): first = {**TURN, "context_id": "chat"} diff --git a/tst/unit/task/episode/spec_test.py b/tst/unit/task/episode/spec_test.py index e4e560d71..2d30e81ab 100644 --- a/tst/unit/task/episode/spec_test.py +++ b/tst/unit/task/episode/spec_test.py @@ -177,9 +177,6 @@ def test_a_task_using_only_the_outside_agents_envs_serves_its_proxy_here(extra, "host, not in container 'task'"), "parts": ({}, {"play": {**PLAY, "prompt": None, "parts": [{"kind": "text", "text": "Plan."}]}}, "step 'play': an outside agent's prompt can't have parts, only a text prompt"), - "user simulator": ({}, {"play": {**PLAY, "max_conversation_turns": 2}}, - "step 'play': an outside agent's prompt can't have a user simulator " - "(max_conversation_turns > 1)"), "output_format": ({}, {"play": {**PLAY, "output_format": {"type": "json_schema", "schema": {}}}}, "step 'play': an outside agent's prompt can't have output_format"), "agentenv_tools": ({}, {"play": {**PLAY, "agentenv_tools": ["complete_step"]}}, From ecb94a70f4f1f6952b03a511c0cc0d89c195b906 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 24 Oct 2026 10:47:45 -0700 Subject: [PATCH 24/39] feat(task): a proxy installs into the container an outside agent would be installed in An install_agent with a container_name now puts the proxy there, as root in the container's working directory, as an agent installed in a container works: install.sh's container mode brings uv, checked against its checksum, and a private Python, so the container's own are untouched. Artifact-generation tasks install their creators this way. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/agents/README.md | 3 +- .../agents/agents/proxy/a2a_server.py | 10 ++- src/agent_env/agents/agents/proxy/install.sh | 86 ++++++++++++++----- src/agent_env/task/episode/rules.py | 13 +-- .../task_step/task_steps/install_agent.py | 10 +-- tst/unit/agents/proxy/test_proxy_agent.py | 26 ++++-- tst/unit/task/episode/spec_test.py | 4 +- tst/unit/task_step/install_agent_test.py | 17 ++-- 8 files changed, 108 insertions(+), 61 deletions(-) diff --git a/src/agent_env/agents/README.md b/src/agent_env/agents/README.md index 024329940..d9a8cd422 100644 --- a/src/agent_env/agents/README.md +++ b/src/agent_env/agents/README.md @@ -81,7 +81,8 @@ which has no model: an episode deploys it in the place of the agent its caller p through it over `urn:agentenv:proxy/v1`. Env tool calls and `bash` run in it, where the agent would work, and it records the turn. Its image works in `/app` as user `agent`. On a `vm` sandbox's host, `install.sh` installs it where the agent it replaces declares it works (`workdir_host` and `user_host` in that agent's `install/v1`), with -the host's own Python in a venv. Its `agent.toml` sets `takes_credentials = false`, so it's given no model key or +the host's own Python in a venv; into a task's container, it installs as root in the container's working directory, +as an agent installed there works, with a private Python of uv's. Its `agent.toml` sets `takes_credentials = false`, so it's given no model key or store credentials. A release publishes its image at `ghcr.io/scaleapi/agentenv-proxy:src-`, as it does the others'. It has no alias and no smoke task: with nobody driving it, a prompt only waits out its deadline. diff --git a/src/agent_env/agents/agents/proxy/a2a_server.py b/src/agent_env/agents/agents/proxy/a2a_server.py index ba576c78f..0bf79e426 100644 --- a/src/agent_env/agents/agents/proxy/a2a_server.py +++ b/src/agent_env/agents/agents/proxy/a2a_server.py @@ -5,8 +5,10 @@ from agentenv_protocol.a2a_agent.proxy import EXTENSIONS, IDENTITY, ProxyAgent, ProxyConfig, stopping_actions _UNPACK = "mkdir -p {work_dir}/agent && tar -xzf {agent_ctx_tar} -C {work_dir}/agent" -# In the place of the agent it replaces: its {workdir}, as its {user}. -_INSTALL_COMMANDS_HOST = [_UNPACK, "sh {work_dir}/agent/install.sh {agent_name} {a2a_port} {workdir} {user}"] +# In the place of the agent it replaces: on a VM's host, its {workdir} as its {user}; in a container, as root in the +# container's working directory, as an agent installed there works. +_INSTALL_COMMANDS_HOST = [_UNPACK, "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {workdir} {user}"] +_INSTALL_COMMANDS = [_UNPACK, "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] @a2a_agent( @@ -14,7 +16,9 @@ config=ProxyConfig, extensions=( *EXTENSIONS, - enable(INSTALL_V1, install_commands_host=_INSTALL_COMMANDS_HOST, + enable(INSTALL_V1, install_commands=_INSTALL_COMMANDS, + required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port"], + install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], a2a_port=8000), ), diff --git a/src/agent_env/agents/agents/proxy/install.sh b/src/agent_env/agents/agents/proxy/install.sh index 13eafe171..17e701c3d 100644 --- a/src/agent_env/agents/agents/proxy/install.sh +++ b/src/agent_env/agents/agents/proxy/install.sh @@ -1,27 +1,67 @@ -#!/bin/sh -# Installs the proxy onto this Linux host from the folder this script is in, and starts it on PORT, as USER in -# WORKDIR, the place of the agent it replaces: +#!/bin/bash +# Installs the proxy from the folder this script is in, and starts it on PORT, in the place of the agent it replaces: # -# install.sh NAME PORT WORKDIR USER +# install.sh host NAME PORT WORKDIR USER onto this Linux host, as USER, in WORKDIR +# install.sh container NAME PORT into the running container NAME, as root, in its working directory # -# It runs in a venv of the host's own Python, which apt installs when the host has none. -set -eu -name=$1 port=$2 workdir=$3 user=$4 +# On a host it runs in a venv of the host's own Python, which apt installs when the host has none. In a container it +# brings uv, checked against its published checksum, and a private Python, so the container's own are never touched. +set -euo pipefail + +UV_VERSION=0.13.1 +UV_SHA256_X86_64=6c86040a8f045ab9adc9e3ec4cdafa03fd00e9a988ed2bf28971c7fc2d0debd0 +UV_SHA256_AARCH64=345a8308e4d1835cdfaadf692836fd3137a6d45864773f05ad7bf59442b1ba16 +PYTHON_VERSION=3.12 + +mode=$1 name=$2 port=$3 here=$(cd "$(dirname "$0")" && pwd) -root=/opt/agentenv-proxy-$name -if ! python3 -c "import ensurepip, venv" 2>/dev/null; then - export DEBIAN_FRONTEND=noninteractive - apt-get update -qq - apt-get install -y -qq --no-install-recommends python3 python3-venv -fi -rm -rf "$root" -mkdir -p "$root" -cp "$here/a2a_server.py" "$here/requirements.txt" "$root/" -python3 -m venv "$root/venv" -"$root/venv/bin/pip" install -q -r "$root/requirements.txt" -id -u "$user" >/dev/null 2>&1 || useradd -m "$user" -mkdir -p "$workdir" -chown -R "$user": "$workdir" -nohup runuser -u "$user" -- env -C "$workdir" A2A_PORT="$port" "$root/venv/bin/python" "$root/a2a_server.py" \ - >"/tmp/agentenv-proxy-$name.log" 2>&1 /dev/null; then + export DEBIAN_FRONTEND=noninteractive + apt-get update -qq + apt-get install -y -qq --no-install-recommends python3 python3-venv + fi + rm -rf "$root" + mkdir -p "$root" + cp "$here/a2a_server.py" "$here/requirements.txt" "$root/" + python3 -m venv "$root/venv" + "$root/venv/bin/pip" install -q -r "$root/requirements.txt" + id -u "$user" >/dev/null 2>&1 || useradd -m "$user" + mkdir -p "$workdir" + chown -R "$user": "$workdir" + nohup runuser -u "$user" -- env -C "$workdir" A2A_PORT="$port" "$root/venv/bin/python" "$root/a2a_server.py" \ + >"/tmp/agentenv-proxy-$name.log" 2>&1 &2; exit 1 ;; + esac + stage=$(mktemp -d "$here/../stage.XXXXXX") + trap 'rm -rf "$stage"' EXIT + uv_archive=uv-$uv_cpu-unknown-linux-musl.tar.gz + curl -fsSL --retry 3 "https://github.com/astral-sh/uv/releases/download/$UV_VERSION/$uv_archive" -o "$stage/$uv_archive" + sum=$(if command -v sha256sum >/dev/null; then sha256sum "$stage/$uv_archive"; else shasum -a 256 "$stage/$uv_archive"; fi) + [ "${sum%% *}" = "$uv_sha256" ] || { echo "$uv_archive doesn't match its published checksum" >&2; exit 1; } + tar -xzf "$stage/$uv_archive" -C "$stage" + # An earlier install's proxy, a retried step's say, stops first, so its port is free. + on "if [ -f $root/proxy.pid ]; then kill \$(cat $root/proxy.pid) 2>/dev/null || true; fi; rm -rf $root; mkdir -p $root/bin" + docker cp "$stage/uv-$uv_cpu-unknown-linux-musl/uv" "$target:$root/bin/uv" >/dev/null + docker cp "$here/a2a_server.py" "$target:$root/a2a_server.py" >/dev/null + docker cp "$here/requirements.txt" "$target:$root/requirements.txt" >/dev/null + on "cd $root && export UV_PYTHON_PREFERENCE=only-managed UV_PYTHON_INSTALL_DIR=$root/python UV_CACHE_DIR=$root/cache \ + && bin/uv venv --quiet --python $PYTHON_VERSION venv && bin/uv pip install --quiet --python venv/bin/python \ + -r requirements.txt && rm -rf cache" + docker exec -d -u 0 -e A2A_PORT="$port" "$target" \ + sh -c "echo \$\$ > $root/proxy.pid; exec $root/venv/bin/python $root/a2a_server.py > $root/proxy.log 2>&1" + ;; + *) echo "usage: install.sh host NAME PORT WORKDIR USER | install.sh container NAME PORT" >&2; exit 2 ;; +esac diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 367931482..5090f6efe 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -16,7 +16,6 @@ from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.collect_artifacts import CollectArtifactsTaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep - from agent_env.task_step.task_steps.install_agent import InstallAgentTaskStep from agent_env.task_step.task_steps.peer_agents import PeerAgentsTaskStep from agent_env.task_step.task_steps.prompt_agent import PromptAgentTaskStep from agent_env.task_step.task_steps.register_env_triggers import RegisterEnvTriggersStep @@ -88,13 +87,6 @@ def _deploy_agent(step: DeployAgentTaskStep, outside: Outside) -> list[str]: ] -def _install_agent(step: InstallAgentTaskStep, outside: Outside) -> list[str]: - if step.agent_name not in outside.agents or step.container_name is None: - return [] - return [f"an outside agent's proxy replaces an agent installed on the VM host, not in container " - f"{step.container_name!r}"] - - def _deploys_an_agent(step: TaskStep, outside: Outside) -> list[str]: if step.agent_name not in outside.agents: return [] @@ -139,7 +131,6 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s CHECKS: dict[str, Callable[..., list[str]]] = { "prompt_agent": _prompt_agent, "deploy_agent": _deploy_agent, - "install_agent": _install_agent, "deploy_human_agent": _deploys_an_agent, "register_agent_triggers": _needs_the_agent, "snapshot_agent_state": _needs_the_agent, @@ -165,8 +156,8 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s "agent_prompt_response_verifier", "verify_mcp_env_assessment", # It skips an agent that isn't deployed. "teardown_sandboxes", - # They act on an outside agent's proxy, in its sandbox, as on the agent. - "add_skills", "load_artifact", "run_code", "verify_sandbox", + # They act on an outside agent's proxy, in its sandbox, as on the agent; install_agent installs the proxy in its place. + "add_skills", "install_agent", "load_artifact", "run_code", "verify_sandbox", }) DEPLOYS = ("deploy_agent", "install_agent") diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 9f825faee..c8ddd12b2 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -250,7 +250,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError(problem) place = None if proxy := self.step_param_overrides(context).get("proxy"): - place = await self._place_of(sandbox, agent) + if self.container_name is None: + place = await self._place_of(sandbox, agent) agent = A2AAgent.get(proxy) if problem := _install_problem(agent): raise RuntimeError(problem) @@ -388,11 +389,8 @@ async def _trust_lines(sandbox, container: Optional[str], work_dir: str) -> str: async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: """Where ``agent``, which a proxy replaces, works on the VM host and as which user: its install/v1 params - ``workdir_host`` and ``user_host``, templates of ``{agent_name}`` as its install commands are.""" - if self.container_name is not None: - raise RuntimeError( - f"install_agent '{self.id}': a proxy replaces an agent installed on the VM host, not in a container" - ) + ``workdir_host`` and ``user_host``, templates of ``{agent_name}`` as its install commands are. In a container + no place is read: the proxy works there as root in the container's working directory, as an agent does.""" if sandbox.ON_THIS_MACHINE: raise RuntimeError( f"install_agent '{self.id}': a proxy isn't installed on a sandbox that runs on this machine, where " diff --git a/tst/unit/agents/proxy/test_proxy_agent.py b/tst/unit/agents/proxy/test_proxy_agent.py index 8d090cac7..9c4f73e3a 100644 --- a/tst/unit/agents/proxy/test_proxy_agent.py +++ b/tst/unit/agents/proxy/test_proxy_agent.py @@ -5,6 +5,8 @@ import re import subprocess +import pytest + from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.bundle.installed import write_installed from agent_env.task.episode.steps import PROXY_AGENT_ID @@ -18,25 +20,31 @@ def _extensions() -> dict: return {ext["uri"]: ext for ext in card["capabilities"]["extensions"]} -def test_the_card_offers_the_proxy_and_a_host_install_only(): +def test_the_card_offers_the_proxy_and_an_install_on_a_host_or_into_a_container(): extensions = _extensions() assert {"urn:agentenv:proxy/v1", "urn:agentenv:mcp-config/v1", "urn:agentenv:trajectory/v1", "urn:agentenv:install/v1"} <= set(extensions) params = extensions["urn:agentenv:install/v1"]["params"] - assert params["install_commands_host"] and not params.get("install_commands") + assert params["install_commands_host"] and params["install_commands"] -def test_the_host_install_asks_only_for_params_install_agent_supplies_and_runs_the_folders_script(): +@pytest.mark.parametrize(("commands", "required", "container", "script"), [ + ("install_commands_host", "required_params_host", None, + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {workdir} {user}"), + ("install_commands", "required_params", "task", "bash {work_dir}/agent/install.sh container {container} {a2a_port}"), +]) +def test_each_install_asks_only_for_params_install_agent_supplies_and_runs_the_folders_script(commands, required, + container, script): params = _extensions()["urn:agentenv:install/v1"]["params"] - supplied = set(_build_param_resolvers(container_name=None, work_dir="/w", agent_ctx_tar="/w/a.tgz", a2a_port=8000, - config=None, agent_name="a")) - placeholders = {name for command in params["install_commands_host"] for name in re.findall(r"\{(\w+)\}", command)} - assert placeholders == set(params["required_params_host"]) and placeholders <= supplied - assert params["install_commands_host"][-1] == "sh {work_dir}/agent/install.sh {agent_name} {a2a_port} {workdir} {user}" + supplied = set(_build_param_resolvers(container_name=container, work_dir="/w", agent_ctx_tar="/w/a.tgz", + a2a_port=8000, config=None, agent_name="a")) + placeholders = {name for command in params[commands] for name in re.findall(r"\{(\w+)\}", command)} + assert placeholders == set(params[required]) and placeholders <= supplied + assert params[commands][-1] == script def test_the_install_script_parses_and_the_image_keeps_it_in_its_build_context(): - subprocess.run(["sh", "-n", str(AGENT_DIR / "install.sh")], check=True) + subprocess.run(["bash", "-n", str(AGENT_DIR / "install.sh")], check=True) assert "COPY a2a_server.py install.sh ./" in (AGENT_DIR / "Dockerfile").read_text() diff --git a/tst/unit/task/episode/spec_test.py b/tst/unit/task/episode/spec_test.py index 2d30e81ab..b6f154e23 100644 --- a/tst/unit/task/episode/spec_test.py +++ b/tst/unit/task/episode/spec_test.py @@ -172,9 +172,6 @@ def test_a_task_using_only_the_outside_agents_envs_serves_its_proxy_here(extra, "needs exactly one deploy_agent or install_agent step, not 2"), "deploy_human_agent": (({"id": "human", "type": "deploy_human_agent", **AT_THE_AGENT},), {"agent": None}, "step 'human': deploy_human_agent deploys 'default-agent', which is played from outside"), - "an install into a container": ({}, {"agent": {**INSTALL, "container_name": "task"}}, - "step 'agent': an outside agent's proxy replaces an agent installed on the VM " - "host, not in container 'task'"), "parts": ({}, {"play": {**PLAY, "prompt": None, "parts": [{"kind": "text", "text": "Plan."}]}}, "step 'play': an outside agent's prompt can't have parts, only a text prompt"), "output_format": ({}, {"play": {**PLAY, "output_format": {"type": "json_schema", "schema": {}}}}, @@ -247,6 +244,7 @@ def test_every_refusal_is_reported_at_once(): {"verify": {"id": "rubric", "type": "rubrics_verifier", "prompt_id": "play", "criteria": [], "use_agent_judge": False}},), "installing the outside agent on a VM host": ({"agent": INSTALL},), + "installing the outside agent into a container": ({"agent": {**INSTALL, "container_name": "task"}},), **{name: ({"agent": replace["agent"]} if replace else {"extra": extra[0]},) for name, (extra, replace) in IN_THE_SANDBOX.items() if name != "an install"}, "tearing down an outside agent": ({"down": {"id": "down", "type": "teardown_sandboxes", diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index c43deeffc..d8dc73e54 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -130,6 +130,7 @@ def test_a2a_port_roundtrip(): class _InstallSandbox: ON_THIS_MACHINE = False + type = "modal_vm" def __init__(self): self.scripts: list[str] = [] @@ -316,7 +317,9 @@ def test_the_trust_lines_resolve_as_the_trust_env_param(): CODEX_INSTALL = {"params": {"workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", "install_commands_host": ["install codex"], "a2a_port": 8000}} PROXY_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], - "required_params_host": ["user", "workdir", "a2a_port"], "a2a_port": 8000}} + "required_params_host": ["user", "workdir", "a2a_port"], + "install_commands": ["serve in {container} on {a2a_port}"], + "required_params": ["container", "a2a_port"], "a2a_port": 8000}} def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, container_name=None): @@ -365,11 +368,15 @@ async def test_a_proxy_needs_the_replaced_agent_to_say_where_it_works(monkeypatc @pytest.mark.asyncio -async def test_a_proxy_isnt_installed_into_a_container(monkeypatch, local_stores): - step, context, _ = _installing_in_place(monkeypatch, container_name="task") +async def test_a_proxy_is_installed_into_the_container_the_agent_it_replaces_would_be(monkeypatch, local_stores): + step, context, sandbox = _installing_in_place(monkeypatch, container_name="task", + codex={"params": {"install_commands": ["install codex"]}}) - with pytest.raises(RuntimeError, match="installed on the VM host, not in a container"): - await step.execute(context) + await step.execute(context) + + assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ + f"serve in {sandbox.scoped_name('task')} on 8000"] + assert [agent.agent_name for agent in context.deployed_agents] == ["solver"] @pytest.mark.asyncio From 4a674fe276396667fa0152161d52ce47ab4a1e14 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sat, 24 Oct 2026 11:05:43 -0700 Subject: [PATCH 25/39] fix(task): the episode sends a proxy call again when it doesn't reach the proxy A reply posted over a connection a tunnel had dropped failed with 'Server disconnected', so its turn waited out its deadline. Every call but record, which isn't idempotent, is now tried again on a transport error: a reply that reached the proxy but whose answer was lost finds its turn ended. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/proxy.py | 34 ++++++++-------- tst/unit/task/episode/proxy_test.py | 60 +++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+), 16 deletions(-) create mode 100644 tst/unit/task/episode/proxy_test.py diff --git a/src/agent_env/task/episode/proxy.py b/src/agent_env/task/episode/proxy.py index df3612eca..4a4b688f1 100644 --- a/src/agent_env/task/episode/proxy.py +++ b/src/agent_env/task/episode/proxy.py @@ -28,7 +28,8 @@ class ProxyRefused(Exception): class ProxyClient: - """One proxy, claimed with a token of this episode's, at the endpoints its card names.""" + """One proxy, claimed with a token of this episode's, at the endpoints its card names. A call that doesn't reach + the proxy, over a connection a tunnel dropped say, is tried again, except ``record``'s, which isn't idempotent.""" def __init__(self, a2a_url: str, card: dict[str, Any]): proxy = A2AAgent.find_extension(card, PROXY_V1.uri) @@ -52,19 +53,20 @@ async def act(self, task_id: str, action: dict[str, Any]) -> dict[str, Any]: """``action``'s result, once it has run in the turn. A call that didn't reach the proxy is tried again: the action's id makes a retry return the action it already started.""" action_id = secrets.token_hex(16) - state = await self._retried("act", task_id=task_id, action_id=action_id, action=action) + state = await self._call("act", task_id=task_id, action_id=action_id, action=action) while not state["done"]: - state = await self._retried("result", action_id=action_id) + state = await self._call("result", action_id=action_id) return state["result"] async def record(self, task_id: str, message: str | dict[str, Any], model: str | None, usage: dict[str, Any] | None) -> None: - await self._call("record", task_id=task_id, message=message, model=model, usage=usage) + _raise_for_status(await self._post("record", {"task_id": task_id, "message": message, "model": model, + "usage": usage}, attempts=1)) async def reply(self, task_id: str, **reply: Any) -> bool: - """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone.""" - response = await self._client.post(self._endpoints["reply"], json={"token": self._token, "task_id": task_id, - **reply}) + """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone, as is one a reply that + reached it but whose answer was lost ended.""" + response = await self._post("reply", {"task_id": task_id, **reply}) if response.status_code == 404: return False _raise_for_status(response) @@ -73,21 +75,21 @@ async def reply(self, task_id: str, **reply: Any) -> bool: async def close(self) -> None: await self._client.aclose() - async def _retried(self, operation: str, **body: Any) -> dict[str, Any]: - for attempt in range(1, _ATTEMPTS + 1): + async def _call(self, operation: str, **body: Any) -> dict[str, Any]: + response = await self._post(operation, body) + _raise_for_status(response) + return response.json() or {} + + async def _post(self, operation: str, body: dict[str, Any], *, attempts: int = _ATTEMPTS) -> httpx.Response: + for attempt in range(1, attempts + 1): try: - return await self._call(operation, **body) + return await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) except httpx.TransportError: - if attempt == _ATTEMPTS: + if attempt == attempts: raise await asyncio.sleep(_RETRY_SECONDS) raise AssertionError("unreachable") - async def _call(self, operation: str, **body: Any) -> dict[str, Any]: - response = await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) - _raise_for_status(response) - return response.json() or {} - class InProcessProxy: """A proxy served on loopback in a thread of its own, with no shell: one for an outside agent that works in no diff --git a/tst/unit/task/episode/proxy_test.py b/tst/unit/task/episode/proxy_test.py new file mode 100644 index 000000000..28b97b6bd --- /dev/null +++ b/tst/unit/task/episode/proxy_test.py @@ -0,0 +1,60 @@ +"""The episode's client of a proxy: which calls it sends again when one doesn't reach the proxy.""" + +import httpx +import pytest +from agentenv_protocol.a2a_agent.proxy import ProxyAgent + +from agent_env.task.episode import proxy as client_module +from agent_env.task.episode.proxy import ProxyClient + +CARD = ProxyAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json", by_alias=True, exclude_none=True) +PARTS = [{"kind": "text", "text": "Added a crate."}] + + +def _client(monkeypatch, answer): + """A client whose transport drops the first request to each endpoint, as a tunnel drops an idle connection, and + answers the next with ``answer``.""" + monkeypatch.setattr(client_module, "_RETRY_SECONDS", 0) + sent: list[str] = [] + + def handle(request: httpx.Request) -> httpx.Response: + sent.append(request.url.path) + if sent.count(request.url.path) == 1: + raise httpx.RemoteProtocolError("Server disconnected without sending a response.") + return answer + + client = ProxyClient("http://proxy", CARD) + client._client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + return client, sent + + +@pytest.mark.asyncio +async def test_a_reply_that_doesnt_reach_the_proxy_is_sent_again(monkeypatch): + client, sent = _client(monkeypatch, httpx.Response(200, json={})) + + assert await client.reply("task-1", parts=PARTS) is True + assert sent == ["/ext/proxy/reply", "/ext/proxy/reply"] + + +@pytest.mark.asyncio +async def test_a_reply_whose_answer_was_lost_finds_its_turn_ended(monkeypatch): + client, _ = _client(monkeypatch, httpx.Response(404, json={"detail": "no open turn 'task-1'"})) + + assert await client.reply("task-1", parts=PARTS) is False + + +@pytest.mark.asyncio +async def test_a_listing_that_doesnt_reach_the_proxy_is_asked_again(monkeypatch): + client, sent = _client(monkeypatch, httpx.Response(200, json={"tools": [{"name": "bash"}]})) + + assert await client.tools("task-1") == [{"name": "bash"}] + assert sent == ["/ext/proxy/tools", "/ext/proxy/tools"] + + +@pytest.mark.asyncio +async def test_a_recorded_model_call_is_never_sent_twice(monkeypatch): + client, sent = _client(monkeypatch, httpx.Response(200, json={})) + + with pytest.raises(httpx.RemoteProtocolError): + await client.record("task-1", "Adding one.", None, None) + assert sent == ["/ext/proxy/record"] From 8de2c27c02c046fa046dddaa527de241cf8fa1f4 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 09:58:26 -0700 Subject: [PATCH 26/39] fix(protocol): the proxy runs sh where there is no bash The proxy's bash tool ran every command with `bash -c`, so in a task container without bash (an Alpine or BusyBox image) every command failed with exit code 126. Where there is no bash on the PATH, it now runs `sh -c` and the tool's description says so; the tool's name, the action kind, spans and call matching are unchanged, and so is everything where bash exists. An integration test installs the proxy into an Alpine and a Debian container with install.sh, from this checkout's protocol, and checks which shell its bash runs. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 3 +- .../src/agentenv_protocol/a2a_agent/proxy.py | 15 ++- .../agentenv-protocol/tests/test_a2a_proxy.py | 28 +++++- .../agent/proxy_container_install_test.py | 99 +++++++++++++++++++ 4 files changed, 138 insertions(+), 7 deletions(-) create mode 100644 tst/integration/agent/proxy_container_install_test.py diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 79ce0d712..92748c991 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -717,7 +717,8 @@ operations: - `act` (`{"token", "task_id", "action_id", "action"}`) runs one action: `{"kind": "tool", "server", "name", "arguments"}` calls an env server through `ToolSession`, with the URL and headers registered for it, and `{"kind": "bash", "command"}` runs `bash -c` in the - proxy's folder, as its user, in a process group of its own; either takes + proxy's folder, as its user, in a process group of its own. Where there is no `bash`, as + in a minimal image, it runs `sh -c`, and the `bash` tool's description says so. Either takes `"timeout_seconds"` (3600 for `bash`, 600 for a tool by default). The driver picks `action_id`, unique for the proxy, so a retried `act` returns the action it already started. It answers `{"done", "result"}` at once when the action ends within the long-poll; diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 1263ba776..0ea6600e6 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -75,6 +75,11 @@ input_schema={"type": "object", "properties": {"command": {"type": "string", "description": "The command."}}, "required": ["command"]}, ) +BASH_TOOL_ON_SH = BASH_TOOL.model_copy(update={ + "description": "Run a command where the agent works, as its user, in its folder, with POSIX sh: this place has " + "no bash, so bash-only syntax such as [[ ]], arrays or <(...) fails. Returns its stdout, then its " + "stderr.", +}) class ProxyConfig(AgentConfig): @@ -127,10 +132,12 @@ async def stopping_actions(app: Starlette) -> AsyncIterator[None]: @a2a_agent(identity=IDENTITY, config=ProxyConfig, extensions=EXTENSIONS, lifespan=stopping_actions) class ProxyAgent(AgentEnvAgent): """``shell=False`` is a proxy for an agent that works in no sandbox: it offers no ``bash`` and keeps no skills. + Where there is no ``bash``, its ``bash`` runs commands with ``sh``, as the tool's description says. ``long_poll_seconds`` bounds how long ``next``, ``act`` and ``result`` wait for something to report.""" def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): self._shell = shell + self._program = "bash" if shutil.which("bash") else "sh" self._long_poll_seconds = long_poll_seconds self._token: bytes | None = None self._turns: dict[str, _Turn] = {} @@ -282,7 +289,7 @@ async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: name are refused.""" async with turn.listing: if turn.tools is None: - offered = {BASH: BASH_TOOL} if self._shell else {} + offered = {BASH: BASH_TOOL if self._program == "bash" else BASH_TOOL_ON_SH} if self._shell else {} for server in turn.servers: try: listed = await asyncio.wait_for(self._listed(turn, server), TOOL_TIMEOUT_SECONDS) @@ -330,7 +337,7 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) if isinstance(action, ProxyBashAction): timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS try: - result = await _run(action.command, timeout) + result = await _run(self._program, action.command, timeout) except OSError as e: # As a shell reports a command it can't execute: an argument list too long, say. result = ProxyBashResult(exit_code=126, stdout="", stderr=f"error: the command couldn't start: {e}\n", @@ -522,11 +529,11 @@ def _now() -> str: return datetime.now(timezone.utc).isoformat() -async def _run(command: str, timeout: float) -> ProxyBashResult: +async def _run(program: str, command: str, timeout: float) -> ProxyBashResult: """Run ``command`` in a process group of its own, stopped with whatever it started if it isn't done within ``timeout``, or if the call is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" process = await asyncio.create_subprocess_exec( - "bash", "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, + program, "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True, ) try: diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 388ee463f..c93aa39a9 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -2,6 +2,7 @@ import asyncio import json +import shutil import socket import subprocess import threading @@ -261,6 +262,13 @@ def test_a_turn_offers_its_env_servers_tools_and_bash(client, env): assert set(tools) == {"add_item", "count", "explode", "whoami", "bash"} assert (tools["add_item"]["server"], "server" in tools["bash"]) == ("items", False) + assert tools["bash"] == { + "name": "bash", + "description": "Run a bash command where the agent works, as its user, in its folder; returns its stdout, " + "then its stderr.", + "input_schema": {"type": "object", "properties": {"command": {"type": "string", "description": "The command."}}, + "required": ["command"]}, + } assert tools["add_item"]["input_schema"]["properties"]["item"]["type"] == "string" @@ -355,9 +363,25 @@ def test_a_bash_action_runs_in_the_proxys_folder_and_returns_a_structured_result (tmp_path / "inputs").mkdir() task_id = _send(client, "go") - result = _done(client, task_id, "b1", _bash("ls; echo out; echo err >&2; exit 3")) + result = _done(client, task_id, "b1", _bash("ls; echo $0; echo out; echo err >&2; exit 3")) - assert result == {"kind": "bash", "exit_code": 3, "stdout": "inputs\nout\n", "stderr": "err\n", "timed_out": False} + assert result == { + "kind": "bash", "exit_code": 3, "stdout": "inputs\nbash\nout\n", "stderr": "err\n", "timed_out": False} + + +def test_where_there_is_no_bash_the_bash_tool_runs_sh_and_says_so(monkeypatch, tmp_path): + which = shutil.which + monkeypatch.setattr(proxy.shutil, "which", lambda name, *args, **kwargs: None if name == "bash" else which(name, *args, **kwargs)) + monkeypatch.chdir(tmp_path) + with TestClient(ProxyAgent(long_poll_seconds=0.2).create_app()) as client: + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 + task_id = _send(client, "go") + + tools = {item["name"]: item for item in _proxy(client, "tools", task_id=task_id).json()["tools"]} + result = _done(client, task_id, "s1", _bash("echo $0")) + + assert "with POSIX sh: this place has no bash" in tools["bash"]["description"] + assert (result["exit_code"], result["stdout"]) == (0, "sh\n") def test_an_action_done_within_the_long_poll_answers_at_once_and_a_slow_one_is_polled(client): diff --git a/tst/integration/agent/proxy_container_install_test.py b/tst/integration/agent/proxy_container_install_test.py new file mode 100644 index 000000000..a9323bee8 --- /dev/null +++ b/tst/integration/agent/proxy_container_install_test.py @@ -0,0 +1,99 @@ +"""The proxy's folder installed into a running container by its own install.sh, as install_agent installs it there. It +brings its own Python and serves; its bash runs commands with bash, or with sh in a container that has only sh (Alpine, +musl), and says which. The container installs this checkout's protocol, so the test covers the code beside it.""" + +import shutil +import subprocess +import time +import uuid +from pathlib import Path +from typing import Any + +import httpx +import pytest + +import agent_env + +pytestmark = [pytest.mark.integration] + +AGENT_DIR = Path(agent_env.__file__).parent / "agents" / "agents" / "proxy" +PROTOCOL_DIR = Path(agent_env.__file__).parents[2] / "packages" / "agentenv-protocol" +PORT = 8000 +TOKEN = "t" * 32 + + +@pytest.fixture(params=[("alpine:3.20", "sh"), ("debian:12-slim", "bash")], ids=["sh-only", "with-bash"]) +def container(request): + image, shell = request.param + name = f"proxy-{shell}-{uuid.uuid4().hex[:8]}" + subprocess.run(["docker", "run", "-d", "--name", name, "-p", f"127.0.0.1::{PORT}", + f"public.ecr.aws/docker/library/{image}", "sleep", "infinity"], check=True, capture_output=True) + try: + yield name, shell + finally: + subprocess.run(["docker", "rm", "-f", name], capture_output=True) + + +def test_installed_into_a_container_the_proxy_serves_and_its_bash_runs_the_shell_there(container, tmp_path): + name, shell = container + has_bash = subprocess.run(["docker", "exec", name, "sh", "-c", "command -v bash"], capture_output=True).returncode == 0 + assert has_bash == (shell == "bash") + folder = tmp_path / "agent" + folder.mkdir() + for file in ("install.sh", "a2a_server.py"): + shutil.copy(AGENT_DIR / file, folder / file) + (folder / "requirements.txt").write_text("agentenv-framework-protocol[agent] @ file:///tmp/protocol\n") + subprocess.run(["docker", "cp", f"{PROTOCOL_DIR}/.", f"{name}:/tmp/protocol"], check=True, capture_output=True) + + subprocess.run(["bash", str(folder / "install.sh"), "container", name, str(PORT)], check=True) + + with httpx.Client(base_url=f"http://{_published(name)}", timeout=30) as client: + _serving(client) + assert client.post("/ext/proxy/claim", json={"token": TOKEN}).status_code == 200 + task_id = _send(client, "go") + tools = _proxy(client, "tools", task_id=task_id).json()["tools"] + result = _done(client, task_id, {"kind": "bash", "command": "echo $0; echo out"}) + + assert [tool["name"] for tool in tools] == ["bash"] + assert tools[0]["description"].startswith( + "Run a bash command where" if shell == "bash" else "Run a command where the agent works, as its user, in its " + "folder, with POSIX sh: this place has no bash") + assert (result["exit_code"], result["stdout"]) == (0, f"{shell}\nout\n") + + +def _published(container: str) -> str: + return subprocess.run(["docker", "port", container, str(PORT)], check=True, capture_output=True, + text=True).stdout.splitlines()[0] + + +def _serving(client: httpx.Client) -> None: + deadline = time.monotonic() + 120 + while True: + try: + if client.get("/.well-known/agent.json").status_code == 200: + return + except httpx.TransportError: + pass + assert time.monotonic() < deadline, "the proxy never served its card" + time.sleep(1) + + +def _send(client: httpx.Client, text: str) -> str: + response = client.post("/a2a", json={ + "jsonrpc": "2.0", "id": "1", "method": "message/send", + "params": {"message": {"kind": "message", "messageId": f"message-{text}", "role": "user", + "contextId": "context-1", "parts": [{"kind": "text", "text": text}]}, + "configuration": {"blocking": False}}, + }) + return response.json()["result"]["id"] + + +def _proxy(client: httpx.Client, operation: str, **body: Any) -> httpx.Response: + return client.post(f"/ext/proxy/{operation}", json={"token": TOKEN, **body}) + + +def _done(client: httpx.Client, task_id: str, action: dict[str, Any]) -> dict[str, Any]: + state = _proxy(client, "act", task_id=task_id, action_id="a1", action=action).json() + while not state["done"]: + state = _proxy(client, "result", action_id="a1").json() + return state["result"] From 20b694fe6a87e3f7b9a4a5a84d8d754a44effb7e Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 09:58:26 -0700 Subject: [PATCH 27/39] fix(task): steps aimed at an agent installed into a container reach that container install_agent didn't record the container it put an agent in, so run_code, verify_sandbox and collect_artifacts aimed at that agent by name went to the sandbox's default agent container. DeployedAgent now carries `container_name`, which install_agent sets for a container install, and the three steps use it before their existing lookup. An agent installed into the default container, `agent-api`, reaches the same container as before. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task_step/context.py | 3 ++ .../task_step/task_steps/collect_artifacts.py | 7 +++-- .../task_step/task_steps/install_agent.py | 1 + .../task_step/task_steps/run_code.py | 2 +- .../task_steps/verifiers/verify_sandbox.py | 5 +-- tst/unit/task_step/install_agent_test.py | 6 ++-- .../test_container_sandbox_routing.py | 31 +++++++++++++------ tst/unit/task_step/test_run_code.py | 17 ++++++++++ 8 files changed, 56 insertions(+), 16 deletions(-) diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index 48a0629ef..4f97e8abc 100644 --- a/src/agent_env/task_step/context.py +++ b/src/agent_env/task_step/context.py @@ -79,6 +79,8 @@ class DeployedAgent: network_policy: dict | None = None # Installed straight onto the sandbox host (install_agent host mode), not into a container. on_host: bool = False + # The container install_agent put it in, as the sandbox's host names it; none for a host install or a deploy. + container_name: str | None = None @classmethod def from_dict(cls, data: dict) -> DeployedAgent: @@ -93,6 +95,7 @@ def from_dict(cls, data: dict) -> DeployedAgent: role=data.get("role"), network_policy=data.get("network_policy"), on_host=bool(data.get("on_host", False)), + container_name=data.get("container_name"), ) def url_for(self, sandbox_type: str | None, *, on_host: bool = False) -> str: diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index efa3b3d36..8ead51ea2 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -778,9 +778,12 @@ async def _collect_via_agent_container(self, context, store, artifact_id, versio sandbox = await self._resolve_live_sandbox(provider, agent.sandbox_id) logger.info(f"Connected to sandbox {agent.sandbox_id} (mode={sandbox.mode})") - # An agent installed on a VM's host has no container: its files are on the host. + # An agent installed on a VM's host has no container: its files are on the host. One installed into a + # container works in that one. container = ( - await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host else None + agent.container_name or await find_agent_container(sandbox) + if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host + else None ) if container: logger.info(f"Using agent container: {container}") diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index c8ddd12b2..31302740e 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -361,6 +361,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_card=agent_card, instance_id=deployed.instance_id, on_host=host_mode, + container_name=container, )) default_model = agent.metadata.get("default_model") or get_config().get_model_for_role("agent") diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 450db74b3..a590e9924 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -453,7 +453,7 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: # A host-installed agent has no container: stage and run on the VM host, as root. return _HostTarget(sandbox) prefix = ( - ("sudo", "docker", "exec", "-u", "0", sandbox.container_name) + ("sudo", "docker", "exec", "-u", "0", agent.container_name or sandbox.container_name) if sandbox.mode == SANDBOX_MODE_VM else () ) diff --git a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py index 916a7887f..35ba0343f 100644 --- a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py +++ b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py @@ -125,9 +125,10 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed, sandbox = await self._deployed_sandbox(context) logger.info(f"Connected to sandbox {sandbox.sandbox_id} (mode={sandbox.mode})") if on_agent: - # An agent installed on a VM's host has no container: its files are on the host. + # An agent installed on a VM's host has no container: its files are on the host. One installed into a + # container works in that one. container = ( - await find_agent_container(sandbox) + agent.container_name or await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host else None ) diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index d8dc73e54..f76f4a9de 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -355,7 +355,8 @@ async def test_a_proxy_is_installed_where_the_agent_it_replaces_says_it_works(mo assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ "serve as agent-solver in /opt/agent-solver/workspace on 8000"] - assert [agent.agent_name for agent in context.deployed_agents] == ["solver"] + assert [(agent.agent_name, agent.on_host, agent.container_name) for agent in context.deployed_agents] == [ + ("solver", True, None)] @pytest.mark.asyncio @@ -376,7 +377,8 @@ async def test_a_proxy_is_installed_into_the_container_the_agent_it_replaces_wou assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ f"serve in {sandbox.scoped_name('task')} on 8000"] - assert [agent.agent_name for agent in context.deployed_agents] == ["solver"] + assert [(agent.agent_name, agent.on_host, agent.container_name) for agent in context.deployed_agents] == [ + ("solver", False, sandbox.scoped_name("task"))] @pytest.mark.asyncio diff --git a/tst/unit/task_step/test_container_sandbox_routing.py b/tst/unit/task_step/test_container_sandbox_routing.py index 799ccd783..7a58a3848 100644 --- a/tst/unit/task_step/test_container_sandbox_routing.py +++ b/tst/unit/task_step/test_container_sandbox_routing.py @@ -3,7 +3,7 @@ A container sandbox on a container provider (Modal) is the runtime itself; one on a VM-backed provider (local, modal_vm, e2b) runs beside its host, so its files are only reachable through `docker exec `. A VM-backed provider reattaches every sandbox as a VM, so a deploy_sandbox sandbox's mode is the one deploy_sandbox recorded. An agent -install_agent put on a VM's host has no container, so both work on the host. +install_agent put on a VM's host has no container, so both work on the host; one it put in a container works there. """ from __future__ import annotations @@ -40,9 +40,10 @@ def _provide(monkeypatch, sandbox): monkeypatch.setattr(module, "get_sandbox_provider", lambda: provider) -def _agent_context(on_host: bool = False): +def _agent_context(on_host: bool = False, container_name: str | None = None): ctx = TaskStepContext() - ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1", on_host=on_host)] + ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1", on_host=on_host, + container_name=container_name)] return ctx @@ -82,11 +83,17 @@ async def _eval(self, sandbox, container, criterion): @pytest.mark.asyncio -@pytest.mark.parametrize("kind,on_host,expected", [("vm", False, CONTAINER), ("vm", True, None), ("container", False, None)]) -async def test_collect_reaches_an_agents_container(monkeypatch, kind, on_host, expected): +@pytest.mark.parametrize("kind,on_host,installed_in,expected", [ + ("vm", False, None, CONTAINER), + ("vm", True, None, None), + ("container", False, None, None), + ("vm", False, CONTAINER, CONTAINER), # installed into the sandbox's own agent container, as tasks do today + ("vm", False, "task", "task"), +]) +async def test_collect_reaches_an_agents_container(monkeypatch, kind, on_host, installed_in, expected): _provide(monkeypatch, _sandbox(kind)) step = CollectArtifactsTaskStep(id="c", version=None, agent_name="solver", artifact_paths=["out.txt"]) - assert await _collect_container(monkeypatch, step, _agent_context(on_host)) == expected + assert await _collect_container(monkeypatch, step, _agent_context(on_host, installed_in)) == expected @pytest.mark.asyncio @@ -103,11 +110,17 @@ async def test_collect_reaches_a_named_sandboxs_container(monkeypatch, kind, rec @pytest.mark.asyncio -@pytest.mark.parametrize("kind,on_host,expected", [("vm", False, CONTAINER), ("vm", True, None), ("container", False, None)]) -async def test_verify_sandbox_probes_an_agents_container(monkeypatch, kind, on_host, expected): +@pytest.mark.parametrize("kind,on_host,installed_in,expected", [ + ("vm", False, None, CONTAINER), + ("vm", True, None, None), + ("container", False, None, None), + ("vm", False, CONTAINER, CONTAINER), # installed into the sandbox's own agent container, as tasks do today + ("vm", False, "task", "task"), +]) +async def test_verify_sandbox_probes_an_agents_container(monkeypatch, kind, on_host, installed_in, expected): _provide(monkeypatch, _sandbox(kind)) step = VerifySandboxTaskStep(id="v", version=None, agent_name="solver", criteria=_CRITERIA) - assert await _verify_container(monkeypatch, step, _agent_context(on_host)) == expected + assert await _verify_container(monkeypatch, step, _agent_context(on_host, installed_in)) == expected @pytest.mark.asyncio diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index ad544628a..bcee43c7f 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -658,6 +658,23 @@ async def test_host_installed_agent_on_a_vm_runs_on_the_host(): assert ctx.metadata["script_results"]["filter"] == {"keep": True} +@pytest.mark.asyncio +@pytest.mark.parametrize("installed_in,expected", [(None, "agent-api"), ("agent-api", "agent-api"), ("task", "task")]) +async def test_an_agent_on_a_vm_runs_the_script_in_the_container_it_was_installed_in(installed_in, expected): + ctx = _make_context() + ctx.deployed_agents[0].sandbox_type = "modal_vm" + ctx.deployed_agents[0].container_name = installed_in + provider = MagicMock() + provider.get_sandbox = AsyncMock(return_value=_vm_sandbox_mock()) + with patch( + "agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", + return_value=provider, + ): + target = await _step()._agent_target(ctx) + + assert target.prefix == ("sudo", "docker", "exec", "-u", "0", expected) + + @pytest.mark.asyncio async def test_timeout_is_reported(): with _sandbox([(124, "", "")]): # the runner's status at the limit From 07627abd4b048a958161e5a0c288ca332af69ae7 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 10:21:26 -0700 Subject: [PATCH 28/39] fix(task): every write into an agent installed in a container reaches that container run_code wrote its script and input into the sandbox's default agent container while running them in the agent's install container, and load_artifact (files, universes, URLs, environment payloads, CLIs) and rubrics_verifier's trajectory staging for a judge named by agent wrote there too. VmSandbox's write methods take an optional container, which those paths pass only for an agent install_agent put in a container, so a deployed agent's calls are unchanged and an agent installed into the default container, agent-api, runs the same commands as before. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/a2a_agent/a2a_agent.py | 27 ++++--- .../providers/sandbox_providers/sandbox.py | 38 +++++---- .../sandbox_providers/smol_vm/sandbox.py | 4 +- .../task_step/task_steps/load_artifact.py | 15 ++-- .../task_step/task_steps/run_code.py | 9 ++- .../task_steps/verifiers/rubrics_verifier.py | 28 ++++--- .../sandbox_providers/vm_sandbox_test.py | 30 +++++++- ...st_load_artifact_agent_container_target.py | 13 ++-- tst/unit/task_step/test_load_artifact_file.py | 77 +++++++++++++++++++ tst/unit/task_step/test_run_code.py | 6 +- .../test_rubrics_verifier_multiturn.py | 14 +++- 11 files changed, 204 insertions(+), 57 deletions(-) diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 1c42a2865..89b6720cb 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -272,11 +272,13 @@ def _skill_extension(deployed: "DeployedA2AAgent") -> dict: return skill_ext @staticmethod - async def install_cli(deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", gateway_url: str) -> str: - """Install a CliArtifact bundle into the agent's filesystem at /opt/cli//. - Writes a sidecar bin/.env with the gateway URL and chmod +x's the entrypoint. - Idempotent. Returns the in-process path to the executable entrypoint.""" - from agent_env.providers.sandbox_providers.sandbox import VmSandbox + async def install_cli( + deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", gateway_url: str, container: str | None = None, + ) -> str: + """Install a CliArtifact bundle into the agent's filesystem at /opt/cli//, in ``container`` when + install_agent put the agent in one. Writes a sidecar bin/.env with the gateway URL and chmod +x's the + entrypoint. Idempotent. Returns the in-process path to the executable entrypoint.""" + from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider provider = build_sandbox_provider(deployed.sandbox_type) if deployed.sandbox_type else get_agent_sandbox_provider() sandbox = await provider.get_sandbox(deployed.sandbox_id) @@ -285,13 +287,16 @@ async def install_cli(deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", env_dir = posixpath.dirname(entrypoint_path) # write .env first — write_file_from_text creates the parent dir; write_file_from_object does not. - await sandbox.write_file_from_text(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n", f"{env_dir}/{GATEWAY_ENV_FILE}") + await sandbox.write_file_from_text(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n", f"{env_dir}/{GATEWAY_ENV_FILE}", + **into_container(container)) for rel_path, file_artifact in cli_artifact.get_cli_files().get_file_artifacts().items(): - await sandbox.write_file_from_object(file_artifact.object_url, f"{target_root}/{rel_path}") + await sandbox.write_file_from_object(file_artifact.object_url, f"{target_root}/{rel_path}", + **into_container(container)) if isinstance(sandbox, VmSandbox): - await sandbox.exec_script(f"docker exec -u 0 {shlex.quote(sandbox.container_name)} chmod +x {entrypoint_path}") + await sandbox.exec_script( + f"docker exec -u 0 {shlex.quote(container or sandbox.container_name)} chmod +x {entrypoint_path}") else: await sandbox.exec("chmod", "+x", entrypoint_path) @@ -303,8 +308,10 @@ async def load_file_artifact_universe( deployed: "DeployedA2AAgent", universe: "FileArtifactUniverse", destination_path: str, + container: str | None = None, ) -> dict[str, str]: - """Stage every FileArtifact in `universe` into the agent sandbox under destination_path.""" + """Stage every FileArtifact in `universe` into the agent sandbox under destination_path, in ``container`` when + install_agent put the agent in one.""" from agent_env.providers.sandbox_providers.sandbox import stage_files_into_container from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider @@ -323,7 +330,7 @@ async def load_file_artifact_universe( f"Loading FileArtifactUniverse '{universe.id}' v{universe.version} " f"({len(file_artifacts)} file(s)) into agent at {destination}" ) - return await stage_files_into_container(sandbox, file_artifacts, destination) + return await stage_files_into_container(sandbox, file_artifacts, destination, container) def __init__( self, diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 9ec670735..29bb4d8e7 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -544,11 +544,13 @@ async def copy_trust_dir_into_container(self, trust_dir: Path, container: str, d finally: await self.exec_script(f"rm -rf {shlex.quote(vm_dir)}") - async def _copy_into_container(self, vm_path: str, destination_path: str) -> None: + async def _copy_into_container(self, vm_path: str, destination_path: str, container: str | None = None) -> None: + """Into ``container``, an agent's that install_agent put in one of its own, else the sandbox's agent container.""" + container = container or self.container_name parent = os.path.dirname(destination_path) if parent: - await self.exec_script(f"docker exec -u 0 {shlex.quote(self.container_name)} mkdir -p {shlex.quote(parent)}") - await self.docker_cp(vm_path, f"{self.container_name}:{destination_path}") + await self.exec_script(f"docker exec -u 0 {shlex.quote(container)} mkdir -p {shlex.quote(parent)}") + await self.docker_cp(vm_path, f"{container}:{destination_path}") @staticmethod def _staging_path(kind: str, destination_path: str) -> str: @@ -556,19 +558,19 @@ def _staging_path(kind: str, destination_path: str) -> str: host's /tmp, as local ones do.""" return f"/tmp/_{kind}_{uuid.uuid4().hex[:12]}_{destination_path.replace('/', '_').lstrip('_')}" - async def write_file_from_object(self, object_url: str, destination_path: str) -> None: + async def write_file_from_object(self, object_url: str, destination_path: str, container: str | None = None) -> None: vm_path = self._staging_path("obj", destination_path) try: await self.load_object_file(object_url, vm_path) - await self._copy_into_container(vm_path, destination_path) + await self._copy_into_container(vm_path, destination_path, container) finally: await self._remove_vm_temp_file(vm_path) - async def write_file_from_url(self, url: str, destination_path: str) -> None: + async def write_file_from_url(self, url: str, destination_path: str, container: str | None = None) -> None: vm_path = self._staging_path("url", destination_path) try: await self.exec_script(f"curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(url)} -o {shlex.quote(vm_path)}") - await self._copy_into_container(vm_path, destination_path) + await self._copy_into_container(vm_path, destination_path, container) finally: await self._remove_vm_temp_file(vm_path) @@ -599,11 +601,11 @@ async def write_host_file(self, data: bytes, vm_path: str) -> None: await self.exec_script(f"mkdir -p {shlex.quote(parent)}") await self._write_bytes_to_vm_path(data, vm_path) - async def write_file_from_text(self, content: str, destination_path: str) -> None: + async def write_file_from_text(self, content: str, destination_path: str, container: str | None = None) -> None: vm_path = self._staging_path("wft", destination_path) try: await self._write_bytes_to_vm_path(content.encode(), vm_path) - await self._copy_into_container(vm_path, destination_path) + await self._copy_into_container(vm_path, destination_path, container) finally: await self._remove_vm_temp_file(vm_path, f"{vm_path}.b64") @@ -666,9 +668,18 @@ async def _pull(sandbox: VmSandbox, image_name: str) -> None: f"Pulling {image_name} for linux/amd64") -async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str) -> dict[str, str]: +def into_container(container: str | None) -> dict[str, str]: + """A write's ``container=``: given only for an agent install_agent put in a container of its own, which only a VM + sandbox hosts, so every other write reaches its sandbox as it always has.""" + return {"container": container} if container else {} + + +async def stage_files_into_container( + sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str, container: str | None = None +) -> dict[str, str]: """Write each object-store file into the sandbox's container under ``destination``; returns ``{name: path}``. - A VM-backed sandbox stages through its host into ``container_name``; a container sandbox is written directly.""" + A VM-backed sandbox stages through its host into ``container``, by default its ``container_name``; a container + sandbox is written directly.""" loaded: dict[str, str] = {} dirs_to_make: set[str] = {destination} for filename in file_artifacts: @@ -679,7 +690,8 @@ async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, for d in sorted(dirs_to_make): if isinstance(sandbox, VmSandbox): - await sandbox.exec_script(f"docker exec -u 0 {shlex.quote(sandbox.container_name)} mkdir -p {shlex.quote(d)}") + await sandbox.exec_script( + f"docker exec -u 0 {shlex.quote(container or sandbox.container_name)} mkdir -p {shlex.quote(d)}") else: await sandbox.exec("mkdir", "-p", d) @@ -689,7 +701,7 @@ async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, async def _stage(filename: str, file_artifact: Any) -> None: async with sem: logger.info(f" {file_artifact.object_url} -> {loaded[filename]}") - await sandbox.write_file_from_object(file_artifact.object_url, loaded[filename]) + await sandbox.write_file_from_object(file_artifact.object_url, loaded[filename], **into_container(container)) await asyncio.gather(*(_stage(fn, fa) for fn, fa in file_artifacts.items())) return loaded diff --git a/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py b/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py index bfc99bce9..31046017e 100644 --- a/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/smol_vm/sandbox.py @@ -84,7 +84,7 @@ async def _download_object_to_vm(self, object_url: str, vm_path: str) -> None: ca = "SSL_CERT_FILE=/etc/agentenv-ca-bundle.pem " if local_https else "" await self.exec_script(f"{ca}curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(rewritten)} -o {shlex.quote(vm_path)}") - async def write_file_from_url(self, url: str, destination_path: str) -> None: + async def write_file_from_url(self, url: str, destination_path: str, container: str | None = None) -> None: vm_path = self._staging_path("url", destination_path) rewritten = LocalSandboxProvider.get_external_url(url) local_https = urlsplit(rewritten).hostname == "host.docker.internal" and urlsplit(rewritten).scheme == "https" @@ -93,7 +93,7 @@ async def write_file_from_url(self, url: str, destination_path: str) -> None: await self.write_host_file(local_ca().bundle_path.read_bytes(), "/etc/agentenv-ca-bundle.pem") ca = "SSL_CERT_FILE=/etc/agentenv-ca-bundle.pem " if local_https else "" await self.exec_script(f"{ca}curl -fsSL {CURL_RETRY_FLAGS} {shlex.quote(rewritten)} -o {shlex.quote(vm_path)}") - await self._copy_into_container(vm_path, destination_path) + await self._copy_into_container(vm_path, destination_path, container) finally: await self._remove_vm_temp_file(vm_path) diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index 1555c9ae1..3c98a6327 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -404,7 +404,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: EnvironmentUniverseArtifact, ) from agent_env.env.env import Env, require_gateway_url - from agent_env.providers.sandbox_providers.sandbox import VmSandbox + from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_VM, build_sandbox_provider, @@ -517,7 +517,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" - files = await A2AAgent.load_file_artifact_universe(deployed_agent, artifact, destination) + files = await A2AAgent.load_file_artifact_universe( + deployed_agent, artifact, destination, **into_container(agent.container_name)) context.metadata.setdefault("loaded_file_artifact_universes", []).append({ "id": artifact.id, "version": artifact.version, @@ -559,7 +560,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) - container = sandbox.container_name if isinstance(sandbox, VmSandbox) else None + container = agent.container_name or ( + sandbox.container_name if isinstance(sandbox, VmSandbox) else None) destination = (destination_path or "/app/files").rstrip("/") or "/" files = await _stage_environment_payload_into_container( sandbox, container, artifact, destination, @@ -610,7 +612,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) gateway_url = reachable_url(require_gateway_url(deployed, "Installing a CliArtifact"), from_sandbox_type=deployed.sandbox_type, to_sandbox_type=agent.sandbox_type) - install_path = await A2AAgent.install_cli(deployed_agent, artifact, gateway_url) + install_path = await A2AAgent.install_cli( + deployed_agent, artifact, gateway_url, **into_container(agent.container_name)) agent_clis = context.metadata.setdefault("installed_clis", {}).setdefault(self.agent_name, {}) agent_clis[artifact.id] = { "install_path": install_path, @@ -620,6 +623,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise ValueError(f"Unsupported artifact type for loading: '{artifact.type}'") if downloads: + agent_container = None if self.sandbox_name is not None: sandbox = target_sandbox else: @@ -632,6 +636,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) + agent_container = agent.container_name destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" semaphore = asyncio.Semaphore(8) @@ -650,7 +655,7 @@ async def _load_one(url: str, filename: str) -> None: elif self.container_name is not None: await self._load_url_into_container(sandbox, sandbox.scoped_name(self.container_name), url, dest) else: - await sandbox.write_file_from_url(url, dest) + await sandbox.write_file_from_url(url, dest, **into_container(agent_container)) logger.info(f"Loaded URL into {target_desc}: {url} -> {dest}") await asyncio.gather(*(_load_one(u, f) for u, f in downloads)) diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index a590e9924..0b8e0f7b9 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -32,6 +32,7 @@ from agent_env.artifact import Artifact, FileArtifact from agent_env.artifact.artifacts.file_artifact_universe import FileArtifactUniverse from agent_env.env.env import DeployedSandboxEnv, require_sandbox +from agent_env.providers.sandbox_providers.sandbox import into_container from agent_env.task_step.context import TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -457,7 +458,7 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: if sandbox.mode == SANDBOX_MODE_VM else () ) - return _AgentTarget(sandbox, prefix) + return _AgentTarget(sandbox, prefix, agent.container_name) # --- errors --- @@ -537,10 +538,12 @@ async def run(self, *command: str) -> tuple[int, str, str]: @dataclass class _AgentTarget: - """Runs the script inside a container: an agent's, a container sandbox's, or an env's in one.""" + """Runs the script inside a container: an agent's, a container sandbox's, or an env's in one. ``container`` is + the one install_agent put the agent in, which its files are written into too.""" sandbox: Any prefix: tuple[str, ...] + container: str | None = None def run_dir(self, step_id: str) -> str: return f"{_RUN_DIR}/{step_id}" @@ -551,7 +554,7 @@ async def reset_dir(self, path: str) -> None: async def write_file(self, path: str, content: str) -> None: # write_file_from_text takes (content, path); we expose (path, content). - await self.sandbox.write_file_from_text(content, path) + await self.sandbox.write_file_from_text(content, path, **into_container(self.container)) async def run(self, *command: str) -> tuple[int, str, str]: return await self.sandbox.exec_with_output(*self.prefix, *command) 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 e908bb612..47fdc6d4f 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 @@ -19,6 +19,7 @@ from agent_env.a2a_agent.staging import transfer_store from agent_env.config import get_config from agent_env.config.model import ModelParam +from agent_env.providers.sandbox_providers.sandbox import into_container from agent_env.task_step.context import PromptResponse, TaskStepContext, dual_keyed from agent_env.task_step.task_steps.verifiers.judge_utils.judge_output_format import ( SCREENSHOT_GROUNDING, @@ -738,6 +739,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: judge_a2a_url: Optional[str] = None judge_sandbox_id: Optional[str] = None judge_agent_card: dict = {} + judge_container: Optional[str] = None if use_agent_judge: if self.agent_name is not None: agent = next((a for a in context.deployed_agents if a.agent_name == self.agent_name), None) @@ -749,6 +751,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: judge_sandbox_id = agent.sandbox_id judge_agent_card = agent.a2a_card or {} judge_agent = agent + judge_container = agent.container_name else: from agent_env.a2a_agent import A2AAgent from agent_env.config import get_config @@ -787,11 +790,13 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: 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: - await self._load_per_turn_trajectories(sandbox, per_turn_uris, multiturn_dir, filter_to_apply) + await self._load_per_turn_trajectories( + sandbox, per_turn_uris, multiturn_dir, filter_to_apply, judge_container) else: - await sandbox.write_file_from_object(compacted.container_object_url, self.TRAJECTORY_CONTAINER_PATH) + 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) + await self._load_tool_result_files(sandbox, compacted.tool_result_files, judge_container) compacted.tool_result_files = [] try: @@ -1077,10 +1082,13 @@ def _filter_trajectory(self, object_url: str, trajectory_filter: TrajectoryFilte compact_url = config.get_object_store().put(compact_key, json.dumps(filtered).encode(), content_type="application/json") return compact_url, tool_result_files - async def _load_tool_result_files(self, sandbox, tool_result_files: list[tuple[str, str]]) -> None: - """Write externalized tool result files into the agent container.""" + async def _load_tool_result_files( + self, sandbox, tool_result_files: list[tuple[str, str]], container: Optional[str] = None, + ) -> None: + """Write externalized tool result files into the agent container: ``container``, when install_agent put the + agent in one.""" for container_path, content in tool_result_files: - await sandbox.write_file_from_text(content, container_path) + await sandbox.write_file_from_text(content, container_path, **into_container(container)) logger.info(f"Loaded {len(tool_result_files)} tool result files into container") async def _load_trajectory_into_container(self, sandbox_id: str, object_url: str) -> None: @@ -1095,7 +1103,7 @@ def _read_trajectory_text(self, object_url: str) -> str: async def _load_per_turn_trajectories( self, sandbox, per_turn_uris: list[str], container_dir: str, - trajectory_filter: Optional[TrajectoryFilter] = None, + trajectory_filter: Optional[TrajectoryFilter] = None, container: Optional[str] = None, ) -> None: """Load each turn into the judge container as turn_NN.json under ``container_dir``. @@ -1111,12 +1119,12 @@ async def _load_per_turn_trajectories( content, turn_files = await asyncio.to_thread( self._filter_trajectory_content, uri, trajectory_filter, f"turn_{i:02d}_" ) - await sandbox.write_file_from_text(content, dest) + await sandbox.write_file_from_text(content, dest, **into_container(container)) tool_result_files.extend(turn_files) else: - await sandbox.write_file_from_object(uri, dest) + await sandbox.write_file_from_object(uri, dest, **into_container(container)) if tool_result_files: - await self._load_tool_result_files(sandbox, tool_result_files) + await self._load_tool_result_files(sandbox, tool_result_files, container) logger.info( f"Loaded {len(per_turn_uris)} per-turn trajectories into {container_dir}" f"{' (filtered)' if trajectory_filter else ''}, {len(tool_result_files)} tool-result files" diff --git a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py index 97e0414ab..721123012 100644 --- a/tst/unit/providers/sandbox_providers/vm_sandbox_test.py +++ b/tst/unit/providers/sandbox_providers/vm_sandbox_test.py @@ -594,12 +594,36 @@ async def test_the_store_is_called_off_the_event_loop(signs): _WRITES = { - "text": lambda sandbox, dest: sandbox.write_file_from_text("body", dest), - "object": lambda sandbox, dest: sandbox.write_file_from_object("s3://bucket/body.json", dest), - "url": lambda sandbox, dest: sandbox.write_file_from_url("https://example.com/body.json", dest), + "text": lambda sandbox, dest, **into: sandbox.write_file_from_text("body", dest, **into), + "object": lambda sandbox, dest, **into: sandbox.write_file_from_object("s3://bucket/body.json", dest, **into), + "url": lambda sandbox, dest, **into: sandbox.write_file_from_url("https://example.com/body.json", dest, **into), } +def _unstaged(scripts: list[str]) -> list[str]: + """``scripts`` with each write's staging path, unique per call, as one name.""" + return [re.sub(r"/tmp/_(\w+?)_[0-9a-f]{12}_", r"/tmp/_\1_", script) for script in scripts] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("write", _WRITES.values(), ids=_WRITES.keys()) +async def test_a_write_into_an_agent_installed_in_the_sandboxs_own_container_runs_what_it_always_has( + signing_store, write): + unnamed, named = _ScriptRecorder(), _ScriptRecorder() + await write(unnamed, "/work/out.json") + await write(named, "/work/out.json", container="agent-api") + assert _unstaged(named.scripts) == _unstaged(unnamed.scripts) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("write", _WRITES.values(), ids=_WRITES.keys()) +async def test_a_write_into_an_agent_installed_in_a_container_of_its_own_reaches_that_one(signing_store, write): + sandbox = _ScriptRecorder() + await write(sandbox, "/work/out.json", container="task") + assert "docker exec -u 0 task mkdir -p /work" in sandbox.scripts + assert [s.split(" ")[3] for s in sandbox.scripts if s.startswith("docker cp ")] == ["task:/work/out.json"] + + @pytest.mark.asyncio @pytest.mark.parametrize("write", _WRITES.values(), ids=_WRITES.keys()) async def test_writes_to_one_destination_stage_at_different_host_paths(signing_store, write): diff --git a/tst/unit/task_step/test_load_artifact_agent_container_target.py b/tst/unit/task_step/test_load_artifact_agent_container_target.py index c71dfe7ab..e66d9d600 100644 --- a/tst/unit/task_step/test_load_artifact_agent_container_target.py +++ b/tst/unit/task_step/test_load_artifact_agent_container_target.py @@ -77,7 +77,7 @@ async def _fake_stage(sandbox, container_name, environment_artifact, destination return seen -def _context_with_agent(monkeypatch, sandbox) -> TaskStepContext: +def _context_with_agent(monkeypatch, sandbox, container_name: str | None = None) -> TaskStepContext: async def _get_sandbox(sandbox_id): return sandbox @@ -92,7 +92,7 @@ async def _get_sandbox(sandbox_id): ctx.deployed_agents = [ DeployedAgent( agent_name="solver", api_url="http://agent", sandbox_id=sandbox.sandbox_id, - sandbox_type=sandbox.type, + sandbox_type=sandbox.type, container_name=container_name, ) ] return ctx @@ -137,8 +137,9 @@ async def test_vm_backed_agent_container_keeps_its_per_sandbox_name( @pytest.mark.asyncio -async def test_plain_vm_agent_gets_the_default_container_name( - monkeypatch, artifact, staged +@pytest.mark.parametrize("installed_in,expected", [(None, "agent-api"), ("agent-api", "agent-api"), ("task", "task")]) +async def test_plain_vm_agent_gets_the_default_container_name_or_the_one_it_was_installed_in( + monkeypatch, artifact, staged, installed_in, expected ): class _Vm(VmSandbox): type = "platform_vm" @@ -152,11 +153,11 @@ def __init__(self) -> None: async def terminate(self) -> None: # pragma: no cover - never called here pass - ctx = _context_with_agent(monkeypatch, _Vm()) + ctx = _context_with_agent(monkeypatch, _Vm(), installed_in) step = LoadArtifactTaskStep( id="stage", version=None, agent_name="solver", artifact_id=ARTIFACT_ID, ) await step.execute(ctx) - assert staged[0]["container"] == "agent-api" + assert staged[0]["container"] == expected assert staged[0]["destination"] == "/app/files" diff --git a/tst/unit/task_step/test_load_artifact_file.py b/tst/unit/task_step/test_load_artifact_file.py index b60787fe9..9d0d777b0 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -3,18 +3,22 @@ from __future__ import annotations +import re from types import SimpleNamespace import pytest from agent_env.a2a_agent import store as agent_store +from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.artifact.artifact import Artifact from agent_env.artifact.artifacts.file import FileArtifact from agent_env.env import env as env_module from agent_env.env.env import DeployedGatewayEnv, Env from agent_env.providers.sandbox_providers import sandbox_provider +from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep +from tst.util.exec_scripts import script_run CHECK = FileArtifact( id="check", version=2, description="checker", filename="check.py", @@ -155,3 +159,76 @@ async def test_into_an_agent(monkeypatch): assert (entry["agent_name"], entry["artifact_type"], entry["files"]) == ( "solver", "file", {"check.py": "/tmp/file_artifacts/check.py"}, ) + + +class _Vm(VmSandbox): + """A VM that records the scripts it runs on its host and the objects it pulls there.""" + + def __init__(self): + self.sandbox_id = "sb-agent" + self.scripts: list[str] = [] + + async def terminate(self) -> None: # pragma: no cover + pass + + async def exec(self, *command): # pragma: no cover + return None + + async def exec_with_output(self, *args): + self.scripts.append(script_run(args) if args[:3] == ("sudo", "bash", "-c") else " ".join(args)) + return 0, "", "" + + async def load_object_file(self, object_url, destination_path): + self.scripts.append(f"pull {object_url} {destination_path}") + + +async def _load_into_agent(monkeypatch, installed_in: str | None, **load) -> list[str]: + """The scripts a load aimed at an agent runs on its VM, its staging paths as one name.""" + sandbox = _Vm() + monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) + monkeypatch.setattr(agent_store, "get_a2a_agent_instance_store", + lambda: SimpleNamespace(get=lambda instance_id: SimpleNamespace(sandbox_type=None, + sandbox_id="sb-agent"))) + context = TaskStepContext() + context.deployed_agents = [DeployedAgent(agent_name="solver", api_url="https://agent", sandbox_id="sb-agent", + instance_id="inst-1", container_name=installed_in)] + await LoadArtifactTaskStep(id="stage", version=None, agent_name="solver", **load).execute(context) + return _unstaged(sandbox.scripts) + + +def _unstaged(scripts: list[str]) -> list[str]: + """``scripts`` with each write's staging path, unique per call, as one name.""" + return [re.sub(r"/tmp/_(\w+?)_[0-9a-f]{12}_", r"/tmp/_\1_", script) for script in scripts] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("load", [{"artifact_id": "check"}, {"urls": ["https://example.com/check.py"]}], + ids=["artifact", "url"]) +async def test_into_an_agent_installed_into_a_container(monkeypatch, load): + deployed = await _load_into_agent(monkeypatch, None, **load) + in_its_own = await _load_into_agent(monkeypatch, "agent-api", **load) + in_another = await _load_into_agent(monkeypatch, "task", **load) + + assert in_its_own == deployed + assert [s.split(" ")[3] for s in in_another if s.startswith("docker cp ")] == [ + "task:/tmp/file_artifacts/check.py"] + assert not any("agent-api" in s for s in in_another) + + +@pytest.mark.asyncio +async def test_a_cli_reaches_an_agent_installed_into_a_container(monkeypatch): + cli = SimpleNamespace(id="tool", version=1, command_name="tool", entrypoint="bin/tool", + get_cli_files=lambda: SimpleNamespace(get_file_artifacts=lambda: {"bin/tool": CHECK})) + + async def _install(installed_in: str | None) -> list[str]: + sandbox = _Vm() + monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) + await A2AAgent.install_cli(SimpleNamespace(sandbox_type=None, sandbox_id="sb-agent"), cli, "https://gateway", + installed_in) + return _unstaged(sandbox.scripts) + + deployed, in_its_own, in_another = await _install(None), await _install("agent-api"), await _install("task") + + assert in_its_own == deployed + assert "docker exec -u 0 task chmod +x /opt/cli/tool/bin/tool" in in_another + assert not any("agent-api" in s for s in in_another) diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index bcee43c7f..90f63496f 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -664,15 +664,19 @@ async def test_an_agent_on_a_vm_runs_the_script_in_the_container_it_was_installe ctx = _make_context() ctx.deployed_agents[0].sandbox_type = "modal_vm" ctx.deployed_agents[0].container_name = installed_in + sandbox = _vm_sandbox_mock() provider = MagicMock() - provider.get_sandbox = AsyncMock(return_value=_vm_sandbox_mock()) + provider.get_sandbox = AsyncMock(return_value=sandbox) with patch( "agent_env.providers.sandbox_providers.sandbox_provider.build_sandbox_provider", return_value=provider, ): target = await _step()._agent_target(ctx) + await target.write_file("/tmp/run/input.json", "{}") assert target.prefix == ("sudo", "docker", "exec", "-u", "0", expected) + sandbox.write_file_from_text.assert_awaited_once_with( + "{}", "/tmp/run/input.json", **({"container": installed_in} if installed_in else {})) @pytest.mark.asyncio diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py index 16c86db58..c1feaaba2 100644 --- a/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_multiturn.py @@ -185,22 +185,26 @@ def _no_merge(*a, **k): # --- _load_per_turn_trajectories: container delivery call-sequence ------------- @pytest.mark.asyncio -async def test_load_per_turn_trajectories_no_filter_writes_each_turn(): +@pytest.mark.parametrize("judge_container", [None, "task"]) +async def test_load_per_turn_trajectories_no_filter_writes_each_turn(judge_container): v = _verifier() sandbox = MagicMock() sandbox.write_file_from_object = AsyncMock() sandbox.write_file_from_text = AsyncMock() await v._load_per_turn_trajectories( - sandbox, ["s3://a/1.json", "s3://a/2.json"], "/tmp/d", None) + sandbox, ["s3://a/1.json", "s3://a/2.json"], "/tmp/d", None, judge_container) dests = [call.args[1] for call in sandbox.write_file_from_object.call_args_list] assert dests == ["/tmp/d/turn_01.json", "/tmp/d/turn_02.json"] + into = {"container": judge_container} if judge_container else {} + assert [call.kwargs for call in sandbox.write_file_from_object.call_args_list] == [into, into] sandbox.write_file_from_text.assert_not_called() @pytest.mark.asyncio -async def test_load_per_turn_trajectories_default_filter_compacts_and_namespaces(monkeypatch): +@pytest.mark.parametrize("judge_container", [None, "task"]) +async def test_load_per_turn_trajectories_default_filter_compacts_and_namespaces(monkeypatch, judge_container): spans = json.dumps([_tool_span("t", {"result": "x" * 2000})]) v = _verifier() monkeypatch.setattr(v, "_read_trajectory_text", lambda uri: spans) @@ -208,7 +212,7 @@ async def test_load_per_turn_trajectories_default_filter_compacts_and_namespaces sandbox.write_file_from_object = AsyncMock() sandbox.write_file_from_text = AsyncMock() - await v._load_per_turn_trajectories(sandbox, ["u1", "u2"], "/tmp/d", TrajectoryFilter()) + await v._load_per_turn_trajectories(sandbox, ["u1", "u2"], "/tmp/d", TrajectoryFilter(), judge_container) sandbox.write_file_from_object.assert_not_called() # compacted -> written as text dests = [call.args[1] for call in sandbox.write_file_from_text.call_args_list] @@ -217,6 +221,8 @@ async def test_load_per_turn_trajectories_default_filter_compacts_and_namespaces assert any("turn_01_tool_call_result" in d for d in tool_dests) assert any("turn_02_tool_call_result" in d for d in tool_dests) assert len(set(tool_dests)) == len(tool_dests) # no cross-turn collisions + into = {"container": judge_container} if judge_container else {} + assert all(call.kwargs == into for call in sandbox.write_file_from_text.call_args_list) # --- trajectory reads and writes run off the event loop ----------------------- From dc96da41050ec390be1ba4d195787f3b5b4d5731 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 15:45:17 -0700 Subject: [PATCH 29/39] fix(protocol): a proxy command is done when its shell exits A command that left a process in the background holding its output, `npm run dev &` say, kept its action waiting until the 3600 s deadline, which then killed the process too. The proxy now takes a command as done when its shell exits: what it started keeps running, output that arrives within a second is returned, and later output is read and dropped so a server never blocks on a full pipe. A command whose shell outlives its deadline is still stopped with everything it started. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 4 +- .../src/agentenv_protocol/a2a_agent/proxy.py | 59 ++++++++++++++++--- .../agentenv-protocol/tests/test_a2a_proxy.py | 33 ++++++++++- 3 files changed, 84 insertions(+), 12 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 92748c991..9311f1ed1 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -718,7 +718,9 @@ operations: "server", "name", "arguments"}` calls an env server through `ToolSession`, with the URL and headers registered for it, and `{"kind": "bash", "command"}` runs `bash -c` in the proxy's folder, as its user, in a process group of its own. Where there is no `bash`, as - in a minimal image, it runs `sh -c`, and the `bash` tool's description says so. Either takes + in a minimal image, it runs `sh -c`, and the `bash` tool's description says so. A command + is done when its shell exits: what it started in the background keeps running, and output + that arrives more than a second later is read and dropped. Either takes `"timeout_seconds"` (3600 for `bash`, 600 for a tool by default). The driver picks `action_id`, unique for the proxy, so a retried `act` returns the action it already started. It answers `{"done", "result"}` at once when the action ends within the long-poll; diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 0ea6600e6..9b860743c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -64,6 +64,8 @@ LONG_POLL_SECONDS = 25.0 BASH_TIMEOUT_SECONDS = 3600 +OUTPUT_GRACE_SECONDS = 1.0 # how long a done command's output may still arrive, from what it left running +EXIT_POLL_SECONDS = 0.05 TOOL_TIMEOUT_SECONDS = 600 SKILLS_DIR = Path.home() / "skills" TRAJECTORY_FORMAT = "agentenv-proxy/v1" @@ -144,6 +146,7 @@ def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_S self._version = 0 self._changed: asyncio.Future[None] | None = None self._actions: dict[str, tuple[_Turn, asyncio.Task[ProxyBashResult | ProxyToolResult]]] = {} + self._draining: set[asyncio.Task[None]] = set() async def run(self, request: TaskRequest[ProxyConfig]) -> TaskResult: if any(not isinstance(part, TextPart) for part in request.parts): @@ -337,7 +340,7 @@ async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) if isinstance(action, ProxyBashAction): timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS try: - result = await _run(self._program, action.command, timeout) + result = await _run(self._program, action.command, timeout, self._draining) except OSError as e: # As a shell reports a command it can't execute: an argument list too long, say. result = ProxyBashResult(exit_code=126, stdout="", stderr=f"error: the command couldn't start: {e}\n", @@ -385,6 +388,8 @@ async def _close(self, turn: _Turn) -> None: async def _stop(self) -> None: await asyncio.gather(*(self._close(turn) for turn in list(self._turns.values())), return_exceptions=True) + for drain in list(self._draining): + drain.cancel() @contextlib.asynccontextmanager async def _skill_folder(self, name: str) -> AsyncIterator[Path]: @@ -529,21 +534,59 @@ def _now() -> str: return datetime.now(timezone.utc).isoformat() -async def _run(program: str, command: str, timeout: float) -> ProxyBashResult: - """Run ``command`` in a process group of its own, stopped with whatever it started if it isn't done within - ``timeout``, or if the call is cancelled, even when ``bash`` has exited and a child it left holds the output open.""" +async def _run(program: str, command: str, timeout: float, draining: set[asyncio.Task[None]]) -> ProxyBashResult: + """Run ``command`` in a process group of its own. It's done when its shell exits: what it started in the background + keeps running, and output that arrives after ``OUTPUT_GRACE_SECONDS`` is read and dropped (in ``draining``), so a + server it started never blocks on a full pipe. A command not done within ``timeout``, or whose call is cancelled, + is stopped with whatever it started.""" process = await asyncio.create_subprocess_exec( program, "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, start_new_session=True, ) + stdout, stderr = _Output(), _Output() + readers = [asyncio.create_task(_read(process.stdout, stdout)), asyncio.create_task(_read(process.stderr, stderr))] try: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout) + await asyncio.wait_for(_exited(process), timeout) except BaseException as e: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) - await process.wait() + await _exited(process) + for reader in readers: + reader.cancel() if isinstance(e, asyncio.TimeoutError): return ProxyBashResult(exit_code=None, stdout="", stderr="", timed_out=True) raise - return ProxyBashResult(exit_code=process.returncode, stdout=stdout.decode(errors="replace"), - stderr=stderr.decode(errors="replace"), timed_out=False) + await asyncio.wait(readers, timeout=OUTPUT_GRACE_SECONDS) + result = ProxyBashResult(exit_code=process.returncode, stdout=stdout.taken(), stderr=stderr.taken(), + timed_out=False) + for reader in readers: + if not reader.done(): + draining.add(reader) + reader.add_done_callback(draining.discard) + return result + + +async def _exited(process: asyncio.subprocess.Process) -> None: + """Until the shell exits: ``Process.wait()`` also waits for everything still holding its output.""" + while process.returncode is None: + await asyncio.sleep(EXIT_POLL_SECONDS) + + +class _Output: + """What a command writes to one stream, kept until it's taken and dropped after.""" + + def __init__(self) -> None: + self._data: bytearray | None = bytearray() + + def add(self, chunk: bytes) -> None: + if self._data is not None: + self._data.extend(chunk) + + def taken(self) -> str: + data, self._data = self._data or bytearray(), None + return data.decode(errors="replace") + + +async def _read(stream: asyncio.StreamReader, output: _Output) -> None: + while chunk := await stream.read(65536): + output.add(chunk) diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index c93aa39a9..e48ae72b9 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -2,7 +2,9 @@ import asyncio import json +import os import shutil +import signal import socket import subprocess import threading @@ -409,12 +411,37 @@ def test_a_command_that_outlives_its_timeout_is_stopped_with_what_it_started(cli task_id = _send(client, "go") waited = _done(client, task_id, "w", _bash("sleep 60 & echo $! > waited; wait", timeout_seconds=0.5)) - left = _done(client, task_id, "l", _bash("sleep 60 & echo $! > left", timeout_seconds=0.5)) - assert waited == left == {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True} + assert waited == {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True} assert ProxyActionState.model_validate({"done": True, "result": waited}).result.exit_code is None _stopped(int((tmp_path / "waited").read_text())) - _stopped(int((tmp_path / "left").read_text())) + + +def test_a_command_is_done_when_its_shell_exits_and_what_it_left_in_the_background_runs_on(client, tmp_path): + task_id = _send(client, "go") + + started = time.monotonic() + left = _done(client, task_id, "l", _bash("sleep 60 & echo $! > left; echo started")) + + assert left == {"kind": "bash", "exit_code": 0, "stdout": "started\n", "stderr": "", "timed_out": False} + assert time.monotonic() - started < proxy.OUTPUT_GRACE_SECONDS + 2 + pid = int((tmp_path / "left").read_text()) + assert _state(pid) and not _state(pid).startswith("Z") + os.kill(pid, signal.SIGKILL) + + +def test_what_a_background_process_writes_after_its_command_is_done_never_blocks_it(client, monkeypatch, tmp_path): + monkeypatch.setattr(proxy, "OUTPUT_GRACE_SECONDS", 0.1) + task_id = _send(client, "go") + writes = "import sys, time; time.sleep(0.5); sys.stdout.write('x' * 200_000); sys.stdout.flush(); open('done', 'w')" + + result = _done(client, task_id, "b", _bash(f"python3 -c \"{writes}\" & echo started")) + + assert (result["exit_code"], result["stdout"]) == (0, "started\n") + deadline = time.monotonic() + 10 + while not (tmp_path / "done").exists(): + assert time.monotonic() < deadline, "the background process blocked on its output" + time.sleep(0.05) def test_a_command_too_long_to_start_fails_as_a_shell_reports_it(client): From e90c83b4eab13d4fa206500894ea1d8a2c50b741 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:47:42 -0700 Subject: [PATCH 30/39] refactor(bundle): import Task and the task store from their modules agent_env.task imports the episode, which imports agent_env.bundle, which imported Task and record_task_cancelled from the agent_env.task package; so the package only imported while .episode came after .store. Importing them from agent_env.task.task and agent_env.task.store makes the order free. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/bundle/importer.py | 2 +- src/agent_env/bundle/materialize.py | 2 +- src/agent_env/bundle/plan.py | 2 +- src/agent_env/bundle/run.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/agent_env/bundle/importer.py b/src/agent_env/bundle/importer.py index c06b5ec3e..c1c6d1f55 100644 --- a/src/agent_env/bundle/importer.py +++ b/src/agent_env/bundle/importer.py @@ -64,7 +64,7 @@ from agent_env.store.base import NotFoundError from agent_env.store.ids import LOCAL_PREFIX from agent_env.store.image_store.oci_registry_credentials import names_registry -from agent_env.task import Task +from agent_env.task.task import Task from agent_env.task_step import registry as task_step_registry from agent_env.task_step.registry import get_task_step_registry from agent_env.utils.paths import os_reason diff --git a/src/agent_env/bundle/materialize.py b/src/agent_env/bundle/materialize.py index 446e6e50a..feffc6310 100644 --- a/src/agent_env/bundle/materialize.py +++ b/src/agent_env/bundle/materialize.py @@ -32,7 +32,7 @@ from agent_env.providers.sandbox_providers.host_images import build, build_staged from agent_env.store.base import NotFoundError from agent_env.store.routing import namespace_routing_enabled -from agent_env.task import Task +from agent_env.task.task import Task from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep from agent_env.utils.docker.build import DockerBuildError diff --git a/src/agent_env/bundle/plan.py b/src/agent_env/bundle/plan.py index 48a4559e5..9a2dac8bd 100644 --- a/src/agent_env/bundle/plan.py +++ b/src/agent_env/bundle/plan.py @@ -51,7 +51,7 @@ from agent_env.store.base import NotFoundError from agent_env.store.ids import LOCAL_PREFIX from agent_env.store.routing import namespace_routing_enabled -from agent_env.task import Task +from agent_env.task.task import Task from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import DeployAgentTaskStep diff --git a/src/agent_env/bundle/run.py b/src/agent_env/bundle/run.py index b05bd02cb..4b524b246 100644 --- a/src/agent_env/bundle/run.py +++ b/src/agent_env/bundle/run.py @@ -28,9 +28,9 @@ from agent_env.providers import build_sandbox_provider from agent_env.store.base import NotFoundError from agent_env.store.routing import namespace_routing, run_scope -from agent_env.task import Task, record_task_cancelled +from agent_env.task.store import record_task_cancelled from agent_env.task.interrupts import Interrupts -from agent_env.task.task import recorded_scores +from agent_env.task.task import Task, recorded_scores from agent_env.task.teardown import TeardownReport, teardown_run from agent_env.task_step.context import TaskStepContext From ccd2ded5cae00dd9a3d20cd8736b4d03653091b7 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:47:52 -0700 Subject: [PATCH 31/39] refactor(task): plainer names for where an installed agent works - DeployedAgent.container_name is installed_container: it holds the scoped name, beside the steps' own, unscoped container_name fields. - install_agent's placeholders for a replaced agent's place are host_workdir and host_user, apart from the staging work_dir; the place is a HostPlace, the installed agent is bound once, and the episode's step param is PROXY_PARAM. Co-Authored-By: Claude Opus 5.5 --- .../agents/agents/proxy/a2a_server.py | 6 +- src/agent_env/task/episode/steps.py | 3 +- src/agent_env/task_step/context.py | 4 +- .../task_step/task_steps/collect_artifacts.py | 2 +- .../task_step/task_steps/install_agent.py | 55 +++++++++++-------- .../task_step/task_steps/load_artifact.py | 8 +-- .../task_step/task_steps/run_code.py | 4 +- .../task_steps/verifiers/rubrics_verifier.py | 2 +- .../task_steps/verifiers/verify_sandbox.py | 2 +- tst/unit/agents/proxy/test_proxy_agent.py | 2 +- tst/unit/task_step/install_agent_test.py | 18 +++--- .../test_container_sandbox_routing.py | 4 +- ...st_load_artifact_agent_container_target.py | 4 +- tst/unit/task_step/test_load_artifact_file.py | 2 +- tst/unit/task_step/test_run_code.py | 2 +- 15 files changed, 64 insertions(+), 54 deletions(-) diff --git a/src/agent_env/agents/agents/proxy/a2a_server.py b/src/agent_env/agents/agents/proxy/a2a_server.py index 0bf79e426..938885ccc 100644 --- a/src/agent_env/agents/agents/proxy/a2a_server.py +++ b/src/agent_env/agents/agents/proxy/a2a_server.py @@ -5,9 +5,9 @@ from agentenv_protocol.a2a_agent.proxy import EXTENSIONS, IDENTITY, ProxyAgent, ProxyConfig, stopping_actions _UNPACK = "mkdir -p {work_dir}/agent && tar -xzf {agent_ctx_tar} -C {work_dir}/agent" -# In the place of the agent it replaces: on a VM's host, its {workdir} as its {user}; in a container, as root in the +# In the place of the agent it replaces: on a VM's host, its {host_workdir} as its {host_user}; in a container, as root in the # container's working directory, as an agent installed there works. -_INSTALL_COMMANDS_HOST = [_UNPACK, "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {workdir} {user}"] +_INSTALL_COMMANDS_HOST = [_UNPACK, "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {host_workdir} {host_user}"] _INSTALL_COMMANDS = [_UNPACK, "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] @@ -19,7 +19,7 @@ enable(INSTALL_V1, install_commands=_INSTALL_COMMANDS, required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port"], install_commands_host=_INSTALL_COMMANDS_HOST, - required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", "workdir", "user"], + required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", "host_workdir", "host_user"], a2a_port=8000), ), lifespan=stopping_actions, diff --git a/src/agent_env/task/episode/steps.py b/src/agent_env/task/episode/steps.py index e363175de..9bcc85816 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -11,6 +11,7 @@ from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_step import TaskStep from agent_env.task_step.task_steps.deploy_agent import bound_deployment, live_mcp_url +from agent_env.task_step.task_steps.install_agent import PROXY_PARAM # The agents bundle's proxy, deployed in the place of an agent an episode plays. PROXY_AGENT_ID = "@local/agentenv-framework/agents/proxy" @@ -55,7 +56,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: await protocol.post_agent_config(deployed.a2a_url + config.endpoint, config.fields) return context await asyncio.to_thread(write_installed, PROXY_AGENT_ID) - override = {"a2a_agent_id": PROXY_AGENT_ID} if self.step.type == "deploy_agent" else {"proxy": PROXY_AGENT_ID} + override = {"a2a_agent_id": PROXY_AGENT_ID} if self.step.type == "deploy_agent" else {PROXY_PARAM: PROXY_AGENT_ID} step_params = context.metadata.setdefault("user_overrides", {}).setdefault("step_params", {}) step_params[self.step.id] = {**step_params.get(self.step.id, {}), **override} context = await self.step.execute(context) diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index 4f97e8abc..b9aa0a0d3 100644 --- a/src/agent_env/task_step/context.py +++ b/src/agent_env/task_step/context.py @@ -80,7 +80,7 @@ class DeployedAgent: # Installed straight onto the sandbox host (install_agent host mode), not into a container. on_host: bool = False # The container install_agent put it in, as the sandbox's host names it; none for a host install or a deploy. - container_name: str | None = None + installed_container: str | None = None @classmethod def from_dict(cls, data: dict) -> DeployedAgent: @@ -95,7 +95,7 @@ def from_dict(cls, data: dict) -> DeployedAgent: role=data.get("role"), network_policy=data.get("network_policy"), on_host=bool(data.get("on_host", False)), - container_name=data.get("container_name"), + installed_container=data.get("installed_container"), ) def url_for(self, sandbox_type: str | None, *, on_host: bool = False) -> str: diff --git a/src/agent_env/task_step/task_steps/collect_artifacts.py b/src/agent_env/task_step/task_steps/collect_artifacts.py index 7ca2e0c8a..5cd475444 100644 --- a/src/agent_env/task_step/task_steps/collect_artifacts.py +++ b/src/agent_env/task_step/task_steps/collect_artifacts.py @@ -781,7 +781,7 @@ async def _collect_via_agent_container(self, context, store, artifact_id, versio # An agent installed on a VM's host has no container: its files are on the host. One installed into a # container works in that one. container = ( - agent.container_name or await find_agent_container(sandbox) + agent.installed_container or await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host else None ) diff --git a/src/agent_env/task_step/task_steps/install_agent.py b/src/agent_env/task_step/task_steps/install_agent.py index 31302740e..dfa2e80a9 100644 --- a/src/agent_env/task_step/task_steps/install_agent.py +++ b/src/agent_env/task_step/task_steps/install_agent.py @@ -17,7 +17,7 @@ import logging import shlex from datetime import datetime, timezone -from typing import Callable, ClassVar, Optional +from typing import Callable, ClassVar, NamedTuple, Optional import httpx @@ -45,6 +45,15 @@ A2A_PORT_KEY = "a2a_port" WORKDIR_HOST_KEY = "workdir_host" USER_HOST_KEY = "user_host" +# The per-step param that installs this agent in place of the step's own, as an episode plays it through a proxy. +PROXY_PARAM = "proxy" + + +class HostPlace(NamedTuple): + """Where an agent works on a VM's host, and as which user.""" + + workdir: str + user: str def _install_problem(agent: A2AAgent) -> str | None: @@ -68,7 +77,7 @@ def _build_param_resolvers( litellm_api_key_override: Optional[str] = None, sandbox_type: Optional[str] = None, trust_lines: str = "", - place: Optional[dict[str, str]] = None, + replaced: Optional[HostPlace] = None, ) -> dict[str, Callable[[], str]]: """Per-install registry: required_params name -> zero-arg resolver returning the raw value. @@ -89,13 +98,13 @@ def _resolve_workspace_dir() -> str: ) return workspace_dir - def _placed(name: str) -> str: - if place is None: + def _replaced(field: str) -> str: + if replaced is None: raise KeyError( - f"agent declares required_param {name!r}, which only an install in another agent's place " - "(the step's `proxy` param) gives" + f"agent declares required_param 'host_{field}', which only an install in another agent's place " + f"(the step's `{PROXY_PARAM}` param) gives" ) - return place[name] + return getattr(replaced, field) def _resolve_container() -> str: if container_name is None: @@ -114,8 +123,8 @@ def _resolve_container() -> str: "agent_name": lambda: agent_name, "workspace_dir": _resolve_workspace_dir, # Where the agent an install replaces works on the VM host, and as which user. - "workdir": lambda: _placed("workdir"), - "user": lambda: _placed("user"), + "host_workdir": lambda: _replaced("workdir"), + "host_user": lambda: _replaced("user"), # Config-sourced (from agent-env config) "litellm_api_key": lambda: (config.resolve_reference(litellm_api_key_override) if litellm_api_key_override else config.get_litellm_api_key()), @@ -245,16 +254,14 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: sandbox = await provider.get_sandbox(ds.sandbox_id) user_overrides = context.metadata.get("user_overrides", {}) - agent = A2AAgent.get(*self._agent_ref(context)) - if problem := _install_problem(agent): + named = A2AAgent.get(*self._agent_ref(context)) + if problem := _install_problem(named): + raise RuntimeError(problem) + proxy = self.step_param_overrides(context).get(PROXY_PARAM) + replaced = await self._host_place_of(sandbox, named) if proxy and self.container_name is None else None + agent = A2AAgent.get(proxy) if proxy else named + if proxy and (problem := _install_problem(agent)): raise RuntimeError(problem) - place = None - if proxy := self.step_param_overrides(context).get("proxy"): - if self.container_name is None: - place = await self._place_of(sandbox, agent) - agent = A2AAgent.get(proxy) - if problem := _install_problem(agent): - raise RuntimeError(problem) logger.info(f"Loading agent image {agent.docker_image_artifact.image_name} onto VM (needed to read agent card)...") await sandbox.load_docker_images([agent.docker_image_artifact]) @@ -313,7 +320,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # A host install on a local sandbox runs on this machine itself, not in a container. sandbox_type=None if host_mode else sandbox.type, trust_lines=trust_lines, - place=place, + replaced=replaced, ) resolved: dict[str, str] = {} for name in required_params: @@ -361,7 +368,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_card=agent_card, instance_id=deployed.instance_id, on_host=host_mode, - container_name=container, + installed_container=container, )) default_model = agent.metadata.get("default_model") or get_config().get_model_for_role("agent") @@ -388,7 +395,7 @@ async def _trust_lines(sandbox, container: Optional[str], work_dir: str) -> str: env = trust_env(directory) return "".join(f"{name}={shlex.quote(value)}\n" for name, value in env.items()) - async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: + async def _host_place_of(self, sandbox, agent: A2AAgent) -> HostPlace: """Where ``agent``, which a proxy replaces, works on the VM host and as which user: its install/v1 params ``workdir_host`` and ``user_host``, templates of ``{agent_name}`` as its install commands are. In a container no place is read: the proxy works there as root in the container's working directory, as an agent does.""" @@ -401,13 +408,13 @@ async def _place_of(self, sandbox, agent: A2AAgent) -> dict[str, str]: await sandbox.load_docker_images([agent.docker_image_artifact]) install_ext = await self._extract_install_extension(sandbox, agent.docker_image_artifact.image_name) params = install_ext.get("params") or {} - place = {"workdir": params.get(WORKDIR_HOST_KEY), "user": params.get(USER_HOST_KEY)} - if not all(isinstance(value, str) and value for value in place.values()): + workdir, user = params.get(WORKDIR_HOST_KEY), params.get(USER_HOST_KEY) + if not all(isinstance(value, str) and value for value in (workdir, user)): raise RuntimeError( f"Agent {agent.id} v{agent.version} install/v1 params don't say where it works on the VM host: a " f"proxy installed in its place needs '{WORKDIR_HOST_KEY}' and '{USER_HOST_KEY}'" ) - return {name: value.format(agent_name=self.agent_name) for name, value in place.items()} + return HostPlace(workdir.format(agent_name=self.agent_name), user.format(agent_name=self.agent_name)) def _register_instance(self, agent, a2a_url: str, ds, agent_card: dict): """Register the installed agent in the instance store, same as A2AAgent.deploy(), diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index 3c98a6327..a4c22878b 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -518,7 +518,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" files = await A2AAgent.load_file_artifact_universe( - deployed_agent, artifact, destination, **into_container(agent.container_name)) + deployed_agent, artifact, destination, **into_container(agent.installed_container)) context.metadata.setdefault("loaded_file_artifact_universes", []).append({ "id": artifact.id, "version": artifact.version, @@ -560,7 +560,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) - container = agent.container_name or ( + container = agent.installed_container or ( sandbox.container_name if isinstance(sandbox, VmSandbox) else None) destination = (destination_path or "/app/files").rstrip("/") or "/" files = await _stage_environment_payload_into_container( @@ -613,7 +613,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: gateway_url = reachable_url(require_gateway_url(deployed, "Installing a CliArtifact"), from_sandbox_type=deployed.sandbox_type, to_sandbox_type=agent.sandbox_type) install_path = await A2AAgent.install_cli( - deployed_agent, artifact, gateway_url, **into_container(agent.container_name)) + deployed_agent, artifact, gateway_url, **into_container(agent.installed_container)) agent_clis = context.metadata.setdefault("installed_clis", {}).setdefault(self.agent_name, {}) agent_clis[artifact.id] = { "install_path": install_path, @@ -636,7 +636,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) - agent_container = agent.container_name + agent_container = agent.installed_container destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" semaphore = asyncio.Semaphore(8) diff --git a/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 0b8e0f7b9..bc8258f1a 100644 --- a/src/agent_env/task_step/task_steps/run_code.py +++ b/src/agent_env/task_step/task_steps/run_code.py @@ -454,11 +454,11 @@ async def _agent_target(self, context: TaskStepContext) -> _ExecTarget: # A host-installed agent has no container: stage and run on the VM host, as root. return _HostTarget(sandbox) prefix = ( - ("sudo", "docker", "exec", "-u", "0", agent.container_name or sandbox.container_name) + ("sudo", "docker", "exec", "-u", "0", agent.installed_container or sandbox.container_name) if sandbox.mode == SANDBOX_MODE_VM else () ) - return _AgentTarget(sandbox, prefix, agent.container_name) + return _AgentTarget(sandbox, prefix, agent.installed_container) # --- errors --- 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 47fdc6d4f..0c952b56d 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 @@ -751,7 +751,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: judge_sandbox_id = agent.sandbox_id judge_agent_card = agent.a2a_card or {} judge_agent = agent - judge_container = agent.container_name + judge_container = agent.installed_container else: from agent_env.a2a_agent import A2AAgent from agent_env.config import get_config diff --git a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py index 35ba0343f..9144fbc58 100644 --- a/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py +++ b/src/agent_env/task_step/task_steps/verifiers/verify_sandbox.py @@ -128,7 +128,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: # An agent installed on a VM's host has no container: its files are on the host. One installed into a # container works in that one. container = ( - agent.container_name or await find_agent_container(sandbox) + agent.installed_container or await find_agent_container(sandbox) if sandbox.mode == SANDBOX_MODE_VM and not agent.on_host else None ) diff --git a/tst/unit/agents/proxy/test_proxy_agent.py b/tst/unit/agents/proxy/test_proxy_agent.py index 9c4f73e3a..39a5fcdc2 100644 --- a/tst/unit/agents/proxy/test_proxy_agent.py +++ b/tst/unit/agents/proxy/test_proxy_agent.py @@ -30,7 +30,7 @@ def test_the_card_offers_the_proxy_and_an_install_on_a_host_or_into_a_container( @pytest.mark.parametrize(("commands", "required", "container", "script"), [ ("install_commands_host", "required_params_host", None, - "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {workdir} {user}"), + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} {host_workdir} {host_user}"), ("install_commands", "required_params", "task", "bash {work_dir}/agent/install.sh container {container} {a2a_port}"), ]) def test_each_install_asks_only_for_params_install_agent_supplies_and_runs_the_folders_script(commands, required, diff --git a/tst/unit/task_step/install_agent_test.py b/tst/unit/task_step/install_agent_test.py index f76f4a9de..981362a91 100644 --- a/tst/unit/task_step/install_agent_test.py +++ b/tst/unit/task_step/install_agent_test.py @@ -15,6 +15,8 @@ from agent_env.task_step.context import DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps import install_agent as install_agent_module from agent_env.task_step.task_steps.install_agent import ( + PROXY_PARAM, + HostPlace, InstallAgentTaskStep, _build_param_resolvers, ) @@ -316,8 +318,8 @@ def test_the_trust_lines_resolve_as_the_trust_env_param(): PROXY = "@local/agentenv-framework/agents/proxy" CODEX_INSTALL = {"params": {"workdir_host": "/opt/agent-{agent_name}/workspace", "user_host": "agent-{agent_name}", "install_commands_host": ["install codex"], "a2a_port": 8000}} -PROXY_INSTALL = {"params": {"install_commands_host": ["serve as {user} in {workdir} on {a2a_port}"], - "required_params_host": ["user", "workdir", "a2a_port"], +PROXY_INSTALL = {"params": {"install_commands_host": ["serve as {host_user} in {host_workdir} on {a2a_port}"], + "required_params_host": ["host_user", "host_workdir", "a2a_port"], "install_commands": ["serve in {container} on {a2a_port}"], "required_params": ["container", "a2a_port"], "a2a_port": 8000}} @@ -341,7 +343,7 @@ def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, cont deployed_sandboxes=[DeployedSandbox( sandbox_name="vm", sandbox_id="sb-1", sandbox_mode="vm", sandbox_type="modal_vm", tunnel_urls={"8000": "https://vm-8000.example"})], - metadata={"user_overrides": {"step_params": {"i": {"proxy": PROXY}}}, + metadata={"user_overrides": {"step_params": {"i": {PROXY_PARAM: PROXY}}}, "deployed_docker_containers": [{"container_name": "task", "sandbox_name": "vm"}]}, ) return step, context, sandbox @@ -355,7 +357,7 @@ async def test_a_proxy_is_installed_where_the_agent_it_replaces_says_it_works(mo assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ "serve as agent-solver in /opt/agent-solver/workspace on 8000"] - assert [(agent.agent_name, agent.on_host, agent.container_name) for agent in context.deployed_agents] == [ + assert [(agent.agent_name, agent.on_host, agent.installed_container) for agent in context.deployed_agents] == [ ("solver", True, None)] @@ -377,7 +379,7 @@ async def test_a_proxy_is_installed_into_the_container_the_agent_it_replaces_wou assert [s for s in sandbox.scripts if s.startswith(("serve", "install"))] == [ f"serve in {sandbox.scoped_name('task')} on 8000"] - assert [(agent.agent_name, agent.on_host, agent.container_name) for agent in context.deployed_agents] == [ + assert [(agent.agent_name, agent.on_host, agent.installed_container) for agent in context.deployed_agents] == [ ("solver", False, sandbox.scoped_name("task"))] @@ -393,9 +395,9 @@ async def test_a_proxy_isnt_installed_on_a_sandbox_that_runs_on_this_machine(mon def test_only_an_install_in_another_agents_place_gives_its_workdir_and_user(): with pytest.raises(KeyError, match="proxy"): - _resolvers(None)["workdir"]() + _resolvers(None)["host_workdir"]() placed = _build_param_resolvers( container_name=None, work_dir="/tmp/w", agent_ctx_tar="/tmp/w/c.tgz", a2a_port=8000, config=object(), - agent_name="solver", place={"workdir": "/opt/agent-solver/workspace", "user": "agent-solver"}, + agent_name="solver", replaced=HostPlace(workdir="/opt/agent-solver/workspace", user="agent-solver"), ) - assert (placed["workdir"](), placed["user"]()) == ("/opt/agent-solver/workspace", "agent-solver") + assert (placed["host_workdir"](), placed["host_user"]()) == ("/opt/agent-solver/workspace", "agent-solver") diff --git a/tst/unit/task_step/test_container_sandbox_routing.py b/tst/unit/task_step/test_container_sandbox_routing.py index 7a58a3848..5daeefb1c 100644 --- a/tst/unit/task_step/test_container_sandbox_routing.py +++ b/tst/unit/task_step/test_container_sandbox_routing.py @@ -40,10 +40,10 @@ def _provide(monkeypatch, sandbox): monkeypatch.setattr(module, "get_sandbox_provider", lambda: provider) -def _agent_context(on_host: bool = False, container_name: str | None = None): +def _agent_context(on_host: bool = False, installed_container: str | None = None): ctx = TaskStepContext() ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1", on_host=on_host, - container_name=container_name)] + installed_container=installed_container)] return ctx diff --git a/tst/unit/task_step/test_load_artifact_agent_container_target.py b/tst/unit/task_step/test_load_artifact_agent_container_target.py index e66d9d600..5edcf03e7 100644 --- a/tst/unit/task_step/test_load_artifact_agent_container_target.py +++ b/tst/unit/task_step/test_load_artifact_agent_container_target.py @@ -77,7 +77,7 @@ async def _fake_stage(sandbox, container_name, environment_artifact, destination return seen -def _context_with_agent(monkeypatch, sandbox, container_name: str | None = None) -> TaskStepContext: +def _context_with_agent(monkeypatch, sandbox, installed_container: str | None = None) -> TaskStepContext: async def _get_sandbox(sandbox_id): return sandbox @@ -92,7 +92,7 @@ async def _get_sandbox(sandbox_id): ctx.deployed_agents = [ DeployedAgent( agent_name="solver", api_url="http://agent", sandbox_id=sandbox.sandbox_id, - sandbox_type=sandbox.type, container_name=container_name, + sandbox_type=sandbox.type, installed_container=installed_container, ) ] return ctx diff --git a/tst/unit/task_step/test_load_artifact_file.py b/tst/unit/task_step/test_load_artifact_file.py index 9d0d777b0..0a26311f7 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -191,7 +191,7 @@ async def _load_into_agent(monkeypatch, installed_in: str | None, **load) -> lis sandbox_id="sb-agent"))) context = TaskStepContext() context.deployed_agents = [DeployedAgent(agent_name="solver", api_url="https://agent", sandbox_id="sb-agent", - instance_id="inst-1", container_name=installed_in)] + instance_id="inst-1", installed_container=installed_in)] await LoadArtifactTaskStep(id="stage", version=None, agent_name="solver", **load).execute(context) return _unstaged(sandbox.scripts) diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index 90f63496f..53c3c50a8 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -663,7 +663,7 @@ async def test_host_installed_agent_on_a_vm_runs_on_the_host(): async def test_an_agent_on_a_vm_runs_the_script_in_the_container_it_was_installed_in(installed_in, expected): ctx = _make_context() ctx.deployed_agents[0].sandbox_type = "modal_vm" - ctx.deployed_agents[0].container_name = installed_in + ctx.deployed_agents[0].installed_container = installed_in sandbox = _vm_sandbox_mock() provider = MagicMock() provider.get_sandbox = AsyncMock(return_value=sandbox) From 1d057c476da4adf85b6ab3e12085be838304bb5a Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:47:57 -0700 Subject: [PATCH 32/39] refactor(agents): one host workdir per agent, checked against its own host install Each agent wrote its host workdir twice, in its host install command and in workdir_host. One _HOST_WORKDIR now feeds both, and a test across the agents checks the declaration is the folder the host install works in. Co-Authored-By: Claude Opus 5.5 --- .../agents/agents/claude-code/a2a_server.py | 6 +++--- .../agents/agents/codex/a2a_server.py | 6 +++--- .../agents/agents/gemini/a2a_server.py | 7 +++---- .../agents/agents/opencode/a2a_server.py | 6 +++--- tst/unit/agents/test_host_workdir.py | 20 +++++++++++++++++++ 5 files changed, 32 insertions(+), 13 deletions(-) create mode 100644 tst/unit/agents/test_host_workdir.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 e241e44a2..082ee207b 100644 --- a/src/agent_env/agents/agents/claude-code/a2a_server.py +++ b/src/agent_env/agents/agents/claude-code/a2a_server.py @@ -117,7 +117,8 @@ _ENDPOINT_PARAMS = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] _INSTALL_COMMANDS = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] # On a host there's no container working directory: the agent works in /app, as in its image. -_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} /app"] +_HOST_WORKDIR = "/app" +_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + f"bash {{work_dir}}/agent/install.sh host {{agent_name}} {{a2a_port}} {_HOST_WORKDIR}"] class ClaudeCodeConfig(AgentConfig): @@ -413,8 +414,7 @@ async def _turns_stop_with_the_server(app: Starlette) -> AsyncIterator[None]: required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], - # Where an install on the host works, and as whom, for a proxy installed in this agent's place. - workdir_host="/app", + workdir_host=_HOST_WORKDIR, user_host="root", ), enable(TRIGGERS_V1, description=( diff --git a/src/agent_env/agents/agents/codex/a2a_server.py b/src/agent_env/agents/agents/codex/a2a_server.py index a91a89711..87b1ccf7d 100644 --- a/src/agent_env/agents/agents/codex/a2a_server.py +++ b/src/agent_env/agents/agents/codex/a2a_server.py @@ -560,7 +560,8 @@ async def _stage_granted_workspace(source: ReadObject | None) -> Path | None: "LITELLM_EXTRA_HEADERS={litellm_extra_headers} AGENT_TRUST_ENV={trust_env} ") _ENDPOINT_PARAMS = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] _INSTALL_COMMANDS = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] -_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} /app"] +_HOST_WORKDIR = "/app" +_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + f"bash {{work_dir}}/agent/install.sh host {{agent_name}} {{a2a_port}} {_HOST_WORKDIR}"] @a2a_agent( @@ -609,8 +610,7 @@ async def _stage_granted_workspace(source: ReadObject | None) -> Path | None: required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], - # Where an install on the host works, and as whom, for a proxy installed in this agent's place. - workdir_host="/app", + workdir_host=_HOST_WORKDIR, user_host="root", ), enable( diff --git a/src/agent_env/agents/agents/gemini/a2a_server.py b/src/agent_env/agents/agents/gemini/a2a_server.py index e70fccd9c..f4a16bdae 100644 --- a/src/agent_env/agents/agents/gemini/a2a_server.py +++ b/src/agent_env/agents/agents/gemini/a2a_server.py @@ -170,8 +170,8 @@ class _Route: "LITELLM_EXTRA_HEADERS={litellm_extra_headers} AGENT_TRUST_ENV={trust_env} ") _ENDPOINT_PARAMS = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] _INSTALL_COMMANDS = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] -_INSTALL_COMMANDS_HOST = [_UNPACK, - _ENDPOINT + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} /workspace"] +_HOST_WORKDIR = "/workspace" +_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + f"bash {{work_dir}}/agent/install.sh host {{agent_name}} {{a2a_port}} {_HOST_WORKDIR}"] def _generate_mcp_name() -> str: @@ -403,8 +403,7 @@ async def _vertex_openai_bridge( install_commands_host=_INSTALL_COMMANDS_HOST, a2a_port=8000, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], - # Where an install on the host works, and as whom, for a proxy installed in this agent's place. - workdir_host="/workspace", + workdir_host=_HOST_WORKDIR, user_host="root", ), enable( diff --git a/src/agent_env/agents/agents/opencode/a2a_server.py b/src/agent_env/agents/agents/opencode/a2a_server.py index f023b833e..b3e4e89ed 100644 --- a/src/agent_env/agents/agents/opencode/a2a_server.py +++ b/src/agent_env/agents/agents/opencode/a2a_server.py @@ -261,7 +261,8 @@ async def lifespan(app: Starlette): "LITELLM_EXTRA_HEADERS={litellm_extra_headers} AGENT_TRUST_ENV={trust_env} ") _ENDPOINT_PARAMS = ["litellm_api_key", "litellm_base_url", "litellm_extra_headers", "trust_env"] _INSTALL_COMMANDS = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] -_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + "bash {work_dir}/agent/install.sh host {agent_name} {a2a_port} /app"] +_HOST_WORKDIR = "/app" +_INSTALL_COMMANDS_HOST = [_UNPACK, _ENDPOINT + f"bash {{work_dir}}/agent/install.sh host {{agent_name}} {{a2a_port}} {_HOST_WORKDIR}"] @a2a_agent( @@ -291,8 +292,7 @@ async def lifespan(app: Starlette): required_params=["container", "work_dir", "agent_ctx_tar", "a2a_port", *_ENDPOINT_PARAMS], install_commands_host=_INSTALL_COMMANDS_HOST, required_params_host=["work_dir", "agent_ctx_tar", "a2a_port", "agent_name", *_ENDPOINT_PARAMS], - # Where an install on the host works, and as whom, for a proxy installed in this agent's place. - workdir_host="/app", + workdir_host=_HOST_WORKDIR, user_host="root", ), # SDK default: the CLI reaches the task's other agents through the SDK's peer_agents MCP tools. diff --git a/tst/unit/agents/test_host_workdir.py b/tst/unit/agents/test_host_workdir.py new file mode 100644 index 000000000..c1a84908b --- /dev/null +++ b/tst/unit/agents/test_host_workdir.py @@ -0,0 +1,20 @@ +"""An agent that installs on a VM's host declares where it works there (install/v1's ``workdir_host``), and that is +the folder its own host install works in, so a proxy installed in its place works where it would.""" + +import pytest + +from tst.unit.agents.claude_code.agent import a2a_server as claude_code +from tst.unit.agents.codex.agent import a2a_server as codex +from tst.unit.agents.gemini.agent import a2a_server as gemini +from tst.unit.agents.opencode.agent import a2a_server as opencode + +INSTALL_V1 = "urn:agentenv:install/v1" + + +@pytest.mark.parametrize("server", [claude_code, codex, gemini, opencode], ids=["claude-code", "codex", "gemini", "opencode"]) +def test_an_agent_declares_the_host_folder_its_own_host_install_works_in(server): + card = server.AGENT_CARD.model_dump(mode="json", by_alias=True, exclude_none=True) + (params,) = [ext["params"] for ext in card["capabilities"]["extensions"] if ext["uri"] == INSTALL_V1] + + assert params["install_commands_host"][-1].split()[-1] == params["workdir_host"] + assert params["user_host"] == "root" From 1156ebb3784baf63ef2d933f0c39456d7bb0638d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:47:59 -0700 Subject: [PATCH 33/39] refactor(protocol): the proxy's bash and tool actions in their own halves The shell and its tool description are chosen together; _run returns a command that can't start as exit code 126 itself and hands back the readers it leaves draining; _perform keeps the span and leaves each kind of action to _bash or _tool. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/proxy.py | 84 ++++++++++--------- 1 file changed, 43 insertions(+), 41 deletions(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index 9b860743c..f9f397de7 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -139,7 +139,7 @@ class ProxyAgent(AgentEnvAgent): def __init__(self, *, shell: bool = True, long_poll_seconds: float = LONG_POLL_SECONDS): self._shell = shell - self._program = "bash" if shutil.which("bash") else "sh" + self._program, self._bash_tool = ("bash", BASH_TOOL) if shutil.which("bash") else ("sh", BASH_TOOL_ON_SH) self._long_poll_seconds = long_poll_seconds self._token: bytes | None = None self._turns: dict[str, _Turn] = {} @@ -292,7 +292,7 @@ async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: name are refused.""" async with turn.listing: if turn.tools is None: - offered = {BASH: BASH_TOOL if self._program == "bash" else BASH_TOOL_ON_SH} if self._shell else {} + offered = {BASH: self._bash_tool} if self._shell else {} for server in turn.servers: try: listed = await asyncio.wait_for(self._listed(turn, server), TOOL_TIMEOUT_SECONDS) @@ -338,37 +338,39 @@ async def _session(self, turn: _Turn, server: str) -> ToolSession: async def _perform(self, turn: _Turn, action: ProxyToolAction | ProxyBashAction) -> ProxyBashResult | ProxyToolResult: started = _now() if isinstance(action, ProxyBashAction): - timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS - try: - result = await _run(self._program, action.command, timeout, self._draining) - except OSError as e: - # As a shell reports a command it can't execute: an argument list too long, say. - result = ProxyBashResult(exit_code=126, stdout="", stderr=f"error: the command couldn't start: {e}\n", - timed_out=False) name, arguments = BASH, {"command": action.command} - output = (f"error: the command timed out after {timeout:g}s" if result.timed_out - else result.stdout + result.stderr) - completion = {"output": output, "exit_code": result.exit_code} - failed = result.timed_out or result.exit_code != 0 + result, completion, failed = await self._bash(action) else: - timeout = action.timeout_seconds or TOOL_TIMEOUT_SECONDS name, arguments = action.name, action.arguments - try: - called = await asyncio.wait_for(self._called(turn, action), timeout) - result = ProxyToolResult(content=called.content, is_error=called.isError) - except asyncio.TimeoutError: - result = _tool_error(f"the tool call timed out after {timeout:g}s") - except (ToolSessionError, httpx.HTTPError) as e: - result = _tool_error(f"{type(e).__name__}: {e}") - text = "\n".join(block["text"] for block in result.content if block.get("type") == "text") - completion = {"error": text} if result.is_error else {"output": text} - failed = result.is_error + result, completion, failed = await self._tool(turn, action) span = _span(name, "execute_tool", started, None, prompt={"input": arguments}, completion=completion, failed=failed) turn.trajectory.append(span) turn.last_event = span["end_time"] return result + async def _bash(self, action: ProxyBashAction) -> tuple[ProxyBashResult, dict[str, Any], bool]: + """Run a command; what it left reading output is drained until the proxy stops.""" + timeout = action.timeout_seconds or BASH_TIMEOUT_SECONDS + result, draining = await _run(self._program, action.command, timeout) + for drain in draining: + self._draining.add(drain) + drain.add_done_callback(self._draining.discard) + output = f"error: the command timed out after {timeout:g}s" if result.timed_out else result.stdout + result.stderr + return result, {"output": output, "exit_code": result.exit_code}, result.timed_out or result.exit_code != 0 + + async def _tool(self, turn: _Turn, action: ProxyToolAction) -> tuple[ProxyToolResult, dict[str, Any], bool]: + timeout = action.timeout_seconds or TOOL_TIMEOUT_SECONDS + try: + called = await asyncio.wait_for(self._called(turn, action), timeout) + result = ProxyToolResult(content=called.content, is_error=called.isError) + except asyncio.TimeoutError: + result = _tool_error(f"the tool call timed out after {timeout:g}s") + except (ToolSessionError, httpx.HTTPError) as e: + result = _tool_error(f"{type(e).__name__}: {e}") + text = "\n".join(block["text"] for block in result.content if block.get("type") == "text") + return result, ({"error": text} if result.is_error else {"output": text}), result.is_error + async def _state(self, action_id: str, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> ProxyActionState: await asyncio.wait((task,), timeout=self._long_poll_seconds) if not task.done(): @@ -534,15 +536,20 @@ def _now() -> str: return datetime.now(timezone.utc).isoformat() -async def _run(program: str, command: str, timeout: float, draining: set[asyncio.Task[None]]) -> ProxyBashResult: +async def _run(program: str, command: str, timeout: float) -> tuple[ProxyBashResult, list[asyncio.Task[None]]]: """Run ``command`` in a process group of its own. It's done when its shell exits: what it started in the background - keeps running, and output that arrives after ``OUTPUT_GRACE_SECONDS`` is read and dropped (in ``draining``), so a - server it started never blocks on a full pipe. A command not done within ``timeout``, or whose call is cancelled, - is stopped with whatever it started.""" - process = await asyncio.create_subprocess_exec( - program, "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, start_new_session=True, - ) + keeps running, and output that arrives after ``OUTPUT_GRACE_SECONDS`` is read and dropped by the readers returned + with the result, so a server it started never blocks on a full pipe. A command not done within ``timeout``, or + whose call is cancelled, is stopped with whatever it started.""" + try: + process = await asyncio.create_subprocess_exec( + program, "-c", command, stdin=asyncio.subprocess.DEVNULL, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, start_new_session=True, + ) + except OSError as e: + # As a shell reports a command it can't execute: an argument list too long, say. + return ProxyBashResult(exit_code=126, stdout="", stderr=f"error: the command couldn't start: {e}\n", + timed_out=False), [] stdout, stderr = _Output(), _Output() readers = [asyncio.create_task(_read(process.stdout, stdout)), asyncio.create_task(_read(process.stderr, stderr))] try: @@ -554,16 +561,11 @@ async def _run(program: str, command: str, timeout: float, draining: set[asyncio for reader in readers: reader.cancel() if isinstance(e, asyncio.TimeoutError): - return ProxyBashResult(exit_code=None, stdout="", stderr="", timed_out=True) + return ProxyBashResult(exit_code=None, stdout="", stderr="", timed_out=True), [] raise await asyncio.wait(readers, timeout=OUTPUT_GRACE_SECONDS) - result = ProxyBashResult(exit_code=process.returncode, stdout=stdout.taken(), stderr=stderr.taken(), - timed_out=False) - for reader in readers: - if not reader.done(): - draining.add(reader) - reader.add_done_callback(draining.discard) - return result + result = ProxyBashResult(exit_code=process.returncode, stdout=stdout.take(), stderr=stderr.take(), timed_out=False) + return result, [reader for reader in readers if not reader.done()] async def _exited(process: asyncio.subprocess.Process) -> None: @@ -582,7 +584,7 @@ def add(self, chunk: bytes) -> None: if self._data is not None: self._data.extend(chunk) - def taken(self) -> str: + def take(self) -> str: data, self._data = self._data or bytearray(), None return data.decode(errors="replace") From e64ed211e287e1a8c7c527b483efda3c1ca69058 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:48:02 -0700 Subject: [PATCH 34/39] refactor(task): the episode reads proxy/v1 as its models ProxyClient returns ProxyNextResponse, ProxyTool and the action results, and takes the actions and reply parts as models, so the episode reads attributes rather than wire keys. Answers are validated with extra='ignore': an agent's protocol isn't pinned, so a newer proxy may add fields. Skill is ProxySkill. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/proxy.py | 56 ++++++++++++++++++++--------- src/agent_env/task/episode/run.py | 53 ++++++++++++++------------- tst/unit/task/episode/proxy_test.py | 18 ++++++++-- tst/unit/task/episode/run_test.py | 45 ++++++++++++++--------- 4 files changed, 110 insertions(+), 62 deletions(-) diff --git a/src/agent_env/task/episode/proxy.py b/src/agent_env/task/episode/proxy.py index 4a4b688f1..bd9bd1e35 100644 --- a/src/agent_env/task/episode/proxy.py +++ b/src/agent_env/task/episode/proxy.py @@ -6,11 +6,23 @@ import asyncio import secrets import threading -from typing import Any +from typing import Any, TypeVar import httpx import uvicorn -from agentenv_protocol.a2a_agent import PROXY_V1 +from agentenv_protocol.a2a_agent import ( + PROXY_V1, + ProxyActionState, + ProxyBashAction, + ProxyBashResult, + ProxyNextResponse, + ProxyTextPart, + ProxyTool, + ProxyToolAction, + ProxyToolResult, + ProxyToolsResponse, +) +from pydantic import BaseModel from agentenv_protocol.a2a_agent.proxy import ProxyAgent from agent_env.a2a_agent.a2a_agent import A2AAgent @@ -21,6 +33,7 @@ _IN_PROCESS_LONG_POLL_SECONDS = 1.0 _ATTEMPTS = 3 _RETRY_SECONDS = 1.0 +_Response = TypeVar("_Response", bound=BaseModel) class ProxyRefused(Exception): @@ -29,7 +42,8 @@ class ProxyRefused(Exception): class ProxyClient: """One proxy, claimed with a token of this episode's, at the endpoints its card names. A call that doesn't reach - the proxy, over a connection a tunnel dropped say, is tried again, except ``record``'s, which isn't idempotent.""" + the proxy, over a connection a tunnel dropped say, is tried again, except ``record``'s, which isn't idempotent. + Answers are read as proxy/v1's models, ignoring fields a newer proxy adds: an agent's protocol isn't pinned.""" def __init__(self, a2a_url: str, card: dict[str, Any]): proxy = A2AAgent.find_extension(card, PROXY_V1.uri) @@ -43,30 +57,36 @@ def __init__(self, a2a_url: str, card: dict[str, Any]): async def claim(self) -> None: await self._call("claim") - async def next(self, after: int) -> dict[str, Any]: - return await self._call("next", after=after) + async def next(self, after: int) -> ProxyNextResponse: + return _read(ProxyNextResponse, await self._call("next", after=after)) - async def tools(self, task_id: str) -> list[dict[str, Any]]: - return (await self._call("tools", task_id=task_id))["tools"] + async def tools(self, task_id: str) -> list[ProxyTool]: + return _read(ProxyToolsResponse, await self._call("tools", task_id=task_id)).tools - async def act(self, task_id: str, action: dict[str, Any]) -> dict[str, Any]: + async def act(self, task_id: str, action: ProxyToolAction | ProxyBashAction) -> ProxyBashResult | ProxyToolResult: """``action``'s result, once it has run in the turn. A call that didn't reach the proxy is tried again: the action's id makes a retry return the action it already started.""" action_id = secrets.token_hex(16) - state = await self._call("act", task_id=task_id, action_id=action_id, action=action) - while not state["done"]: - state = await self._call("result", action_id=action_id) - return state["result"] + body = action.model_dump(mode="json", exclude_none=True) + state = _read(ProxyActionState, await self._call("act", task_id=task_id, action_id=action_id, action=body)) + while not state.done: + state = _read(ProxyActionState, await self._call("result", action_id=action_id)) + return state.result async def record(self, task_id: str, message: str | dict[str, Any], model: str | None, usage: dict[str, Any] | None) -> None: _raise_for_status(await self._post("record", {"task_id": task_id, "message": message, "model": model, "usage": usage}, attempts=1)) - async def reply(self, task_id: str, **reply: Any) -> bool: - """Whether the reply reached an open turn: a turn its sender stopped waiting for is gone, as is one a reply that - reached it but whose answer was lost ended.""" - response = await self._post("reply", {"task_id": task_id, **reply}) + async def reply(self, task_id: str, *, parts: list[ProxyTextPart] | None = None, error: str | None = None) -> bool: + """Whether the reply, its ``parts`` or the turn's ``error``, reached an open turn: a turn its sender stopped + waiting for is gone, as is one a reply that reached it but whose answer was lost ended.""" + body: dict[str, Any] = {"task_id": task_id} + if parts is not None: + body["parts"] = [part.model_dump(mode="json") for part in parts] + if error is not None: + body["error"] = error + response = await self._post("reply", body) if response.status_code == 404: return False _raise_for_status(response) @@ -117,6 +137,10 @@ def stop(self) -> None: self._server.should_exit = True +def _read(model: type[_Response], answer: dict[str, Any]) -> _Response: + return model.model_validate(answer, extra="ignore") + + def _raise_for_status(response: httpx.Response) -> None: if response.status_code in (400, 409): try: diff --git a/src/agent_env/task/episode/run.py b/src/agent_env/task/episode/run.py index e49876d29..44e088e3b 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -3,11 +3,17 @@ import asyncio import logging from collections import deque -from dataclasses import dataclass from typing import Any import httpx from agentenv_protocol import EnvironmentTool, ToolResult +from agentenv_protocol.a2a_agent import ( + ProxyBashAction, + ProxyBashResult, + ProxySkill, + ProxyTextPart, + ProxyToolAction, +) from agent_env.a2a_agent.a2a_agent import A2AAgent from agent_env.providers.sandbox_providers.sandbox_provider import sandbox_request_headers_for_url @@ -30,13 +36,8 @@ class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" -@dataclass(frozen=True) -class Skill: - """A skill the task gave the agent: its instructions are ``path``'s ``SKILL.md``, in the agent's sandbox.""" - - name: str - description: str - path: str +# A skill the task gave the agent: its instructions are ``path``'s ``SKILL.md``, in the agent's sandbox. +Skill = ProxySkill class Turn: @@ -79,7 +80,7 @@ async def end(self, response: str | None = None) -> None: if not self._open: return try: - await self._checked(self._proxy.reply(self._task_id, parts=[{"kind": "text", "text": response or ""}])) + await self._checked(self._proxy.reply(self._task_id, parts=[ProxyTextPart(text=response or "")])) except httpx.HTTPError as e: logger.warning("turn %r of agent %r: its reply didn't reach the proxy: %s", self.step_id, self.agent, e) self._open = False @@ -121,9 +122,8 @@ def __init__(self, turn: Turn, timeout: float): async def __aenter__(self) -> TurnTools: for tool in await self._turn._checked(self._turn._proxy.tools(self._turn._task_id)): - self._tools.append(EnvironmentTool(name=tool["name"], description=tool.get("description"), - inputSchema=tool.get("input_schema"))) - self._servers[tool["name"]] = tool.get("server") + self._tools.append(EnvironmentTool(name=tool.name, description=tool.description, inputSchema=tool.input_schema)) + self._servers[tool.name] = tool.server return self async def __aexit__(self, *exc_info: object) -> None: @@ -138,19 +138,19 @@ async def call(self, name: str, arguments: dict[str, Any] | None = None) -> Tool if (server := self._servers[name]) is None: if not isinstance(command := (arguments or {}).get("command"), str): raise EpisodeError(f"tool {name!r} takes a command string, not {command!r}") - action = {"kind": "bash", "command": command} + action: ProxyBashAction | ProxyToolAction = ProxyBashAction(kind="bash", command=command) else: - action = {"kind": "tool", "server": server, "name": name, "arguments": arguments or {}, - "timeout_seconds": self._timeout} + action = ProxyToolAction(kind="tool", server=server, name=name, arguments=arguments or {}, + timeout_seconds=self._timeout) result = await self._turn._checked(self._turn._proxy.act(self._turn._task_id, action)) - if result["kind"] == "bash": - if result["timed_out"]: + if isinstance(result, ProxyBashResult): + if result.timed_out: return ToolResult(content=[{"type": "text", "text": "error: the command timed out"}], isError=True) - text = result["stdout"] + result["stderr"] - if failed := result["exit_code"] != 0: - text = "\n".join(filter(None, (text.removesuffix("\n"), f"exit code {result['exit_code']}"))) + text = result.stdout + result.stderr + if failed := result.exit_code != 0: + text = "\n".join(filter(None, (text.removesuffix("\n"), f"exit code {result.exit_code}"))) return ToolResult(content=[{"type": "text", "text": text}], isError=failed) - return ToolResult(content=result["content"], isError=result["is_error"]) + return ToolResult(content=result.content, isError=result.is_error) class Episode: @@ -303,8 +303,8 @@ async def _drain(self, agent: str, proxy: ProxyClient) -> None: logger.warning("episode %s: reading agent %r's turns failed, retrying: %s", self.id, agent, e) await asyncio.sleep(_DRAIN_RETRY_SECONDS) continue - after = answer["version"] - listed = {item["task_id"]: item for item in answer["turns"]} + after = answer.version + listed = {item.task_id: item for item in answer.turns} for task_id in [task_id for task_id in open_turns if task_id not in listed]: turn = open_turns.pop(task_id) turn._open = False @@ -312,10 +312,9 @@ async def _drain(self, agent: str, proxy: ProxyClient) -> None: for task_id, item in listed.items(): if task_id not in open_turns: open_turns[task_id] = turn = Turn( - agent=agent, step_id=self._step_of(agent, item["context_id"]), - prompt="\n".join(part["text"] for part in item["parts"]), - system_prompt=item.get("system_prompt"), - skills=[Skill(**skill) for skill in item.get("skills", [])], proxy=proxy, task_id=task_id, + agent=agent, step_id=self._step_of(agent, item.context_id), + prompt="\n".join(part.text for part in item.parts), system_prompt=item.system_prompt, + skills=list(item.skills), proxy=proxy, task_id=task_id, ) self._waiting.append(turn) self._turns.add(turn) diff --git a/tst/unit/task/episode/proxy_test.py b/tst/unit/task/episode/proxy_test.py index 28b97b6bd..a5d4233c2 100644 --- a/tst/unit/task/episode/proxy_test.py +++ b/tst/unit/task/episode/proxy_test.py @@ -2,13 +2,14 @@ import httpx import pytest +from agentenv_protocol.a2a_agent import ProxyTextPart, ProxyTool from agentenv_protocol.a2a_agent.proxy import ProxyAgent from agent_env.task.episode import proxy as client_module from agent_env.task.episode.proxy import ProxyClient CARD = ProxyAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json", by_alias=True, exclude_none=True) -PARTS = [{"kind": "text", "text": "Added a crate."}] +PARTS = [ProxyTextPart(text="Added a crate.")] def _client(monkeypatch, answer): @@ -47,7 +48,7 @@ async def test_a_reply_whose_answer_was_lost_finds_its_turn_ended(monkeypatch): async def test_a_listing_that_doesnt_reach_the_proxy_is_asked_again(monkeypatch): client, sent = _client(monkeypatch, httpx.Response(200, json={"tools": [{"name": "bash"}]})) - assert await client.tools("task-1") == [{"name": "bash"}] + assert await client.tools("task-1") == [ProxyTool(name="bash")] assert sent == ["/ext/proxy/tools", "/ext/proxy/tools"] @@ -58,3 +59,16 @@ async def test_a_recorded_model_call_is_never_sent_twice(monkeypatch): with pytest.raises(httpx.RemoteProtocolError): await client.record("task-1", "Adding one.", None, None) assert sent == ["/ext/proxy/record"] + + +@pytest.mark.asyncio +async def test_fields_a_newer_proxy_adds_are_ignored(): + turn = {"task_id": "task-1", "context_id": "c", "parts": [{"kind": "text", "text": "Add one.", "lang": "en"}], + "priority": 1} + client = ProxyClient("http://proxy", CARD) + client._client = httpx.AsyncClient(transport=httpx.MockTransport( + lambda request: httpx.Response(200, json={"version": 2, "turns": [turn], "paused": False}))) + + answer = await client.next(0) + + assert (answer.version, [t.parts[0].text for t in answer.turns]) == (2, ["Add one."]) diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index 12cec2a5b..4c67d8461 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -10,7 +10,18 @@ import pytest_asyncio import uvicorn from agentenv_protocol import EnvironmentTool, ToolResult, ToolSessionError -from agentenv_protocol.a2a_agent import AgentEnvAgent, AgentIdentity, TaskRequest, TaskResult, a2a_agent +from agentenv_protocol.a2a_agent import ( + AgentEnvAgent, + AgentIdentity, + ProxyBashAction, + ProxyBashResult, + ProxyTextPart, + ProxyTool, + ProxyToolResult, + TaskRequest, + TaskResult, + a2a_agent, +) from agentenv_protocol.a2a_agent import proxy as proxy_module import agent_env.task.episode.run as episode_run @@ -624,17 +635,17 @@ def __init__(self, results=None): self.replies = [] async def tools(self, task_id): - return [{"name": "items_add_item", "server": "items"}, {"name": "bash", "description": "Run a command."}] + return [ProxyTool(name="items_add_item", server="items"), ProxyTool(name="bash", description="Run a command.")] async def act(self, task_id, action): self.actions.append(action) - if action["kind"] == "bash": - return self.results.get(action["command"]) or { - "kind": "bash", "exit_code": 0, "stdout": f"ran {action['command']}\n", "stderr": "", "timed_out": False} - return {"kind": "tool", "content": [{"type": "text", "text": "added"}], "is_error": False} + if isinstance(action, ProxyBashAction): + return self.results.get(action.command) or ProxyBashResult( + exit_code=0, stdout=f"ran {action.command}\n", stderr="", timed_out=False) + return ProxyToolResult(content=[{"type": "text", "text": "added"}], is_error=False) - async def reply(self, task_id, **reply): - self.replies.append((task_id, reply)) + async def reply(self, task_id, *, parts=None, error=None): + self.replies.append((task_id, parts, error)) return True @@ -645,10 +656,9 @@ def _sandboxed_turn(proxy): @pytest.mark.asyncio async def test_a_turn_in_the_agents_sandbox_offers_bash_and_shows_its_output_and_a_failed_exit(): - proxy = _SandboxedProxy({"slow": {"kind": "bash", "stdout": "", "stderr": "", "timed_out": True}, - "both": {"kind": "bash", "exit_code": 2, "stdout": "out\n", "stderr": "err\n", - "timed_out": False}, - "quiet": {"kind": "bash", "exit_code": 1, "stdout": "", "stderr": "", "timed_out": False}}) + proxy = _SandboxedProxy({"slow": ProxyBashResult(stdout="", stderr="", timed_out=True), + "both": ProxyBashResult(exit_code=2, stdout="out\n", stderr="err\n", timed_out=False), + "quiet": ProxyBashResult(exit_code=1, stdout="", stderr="", timed_out=False)}) async with _sandboxed_turn(proxy).tools() as tools: assert [tool.name for tool in tools.list()] == ["items_add_item", "bash"] @@ -663,10 +673,11 @@ async def test_a_turn_in_the_agents_sandbox_offers_bash_and_shows_its_output_and assert [(result.isError, result.text) for result in (ran, both, quiet, timed_out)] == [ (False, "ran ls inputs\n"), (True, "out\nerr\nexit code 2"), (True, "exit code 1"), (True, "error: the command timed out")] - assert proxy.actions == [{"kind": "bash", "command": "ls inputs"}, {"kind": "bash", "command": "both"}, - {"kind": "bash", "command": "quiet"}, {"kind": "bash", "command": "slow"}, - {"kind": "tool", "server": "items", "name": "items_add_item", - "arguments": {"name": "crate"}, "timeout_seconds": 600.0}] + assert [action.model_dump(exclude_none=True) for action in proxy.actions] == [ + {"kind": "bash", "command": "ls inputs"}, {"kind": "bash", "command": "both"}, + {"kind": "bash", "command": "quiet"}, {"kind": "bash", "command": "slow"}, + {"kind": "tool", "server": "items", "name": "items_add_item", "arguments": {"name": "crate"}, + "timeout_seconds": 600.0}] @pytest.mark.asyncio @@ -677,7 +688,7 @@ async def test_a_turn_ends_with_its_reply_as_parts(): await turn.end("Added a crate.") await turn.end("Again.") - assert proxy.replies == [("t-1", {"parts": [{"kind": "text", "text": "Added a crate."}]})] + assert proxy.replies == [("t-1", [ProxyTextPart(text="Added a crate.")], None)] @pytest.mark.asyncio From 6ba355d6be7af52475b1fa551b3964ef3c4db617 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Sun, 25 Oct 2026 19:49:56 -0700 Subject: [PATCH 35/39] fix(protocol): a command cancelled while its late output is awaited stops what it left The wait for output arriving after a command's shell exits sat outside the cleanup, so a turn closing during it left the command's background processes running and their readers buffering, never handed to the drain. The wait is inside it now: a cancelled call stops the command's process group and its readers, as a cancel before the shell exits always did. Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/a2a_agent/proxy.py | 2 +- packages/agentenv-protocol/tests/test_a2a_proxy.py | 13 +++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index f9f397de7..beb75410f 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -554,6 +554,7 @@ async def _run(program: str, command: str, timeout: float) -> tuple[ProxyBashRes readers = [asyncio.create_task(_read(process.stdout, stdout)), asyncio.create_task(_read(process.stderr, stderr))] try: await asyncio.wait_for(_exited(process), timeout) + await asyncio.wait(readers, timeout=OUTPUT_GRACE_SECONDS) except BaseException as e: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGKILL) @@ -563,7 +564,6 @@ async def _run(program: str, command: str, timeout: float) -> tuple[ProxyBashRes if isinstance(e, asyncio.TimeoutError): return ProxyBashResult(exit_code=None, stdout="", stderr="", timed_out=True), [] raise - await asyncio.wait(readers, timeout=OUTPUT_GRACE_SECONDS) result = ProxyBashResult(exit_code=process.returncode, stdout=stdout.take(), stderr=stderr.take(), timed_out=False) return result, [reader for reader in readers if not reader.done()] diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index e48ae72b9..128027f03 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -490,6 +490,19 @@ def test_a_turn_that_closes_stops_its_running_actions_and_forgets_them(client, t assert _proxy(client, "act", task_id=task_id, action_id="new", action=_bash("true")).status_code == 404 +def test_a_turn_that_closes_while_a_done_commands_output_is_awaited_stops_what_it_left(client, monkeypatch, tmp_path): + monkeypatch.setattr(proxy, "OUTPUT_GRACE_SECONDS", 30) + task_id = _send(client, "go") + assert _proxy(client, "act", task_id=task_id, action_id="left", + action=_bash("sleep 60 & echo $! > child")).json()["done"] is False + + _proxy(client, "reply", task_id=task_id, parts=[{"kind": "text", "text": "done"}]) + _settled(client, task_id) + + _stopped(int((tmp_path / "child").read_text())) + assert _proxy(client, "result", action_id="left").status_code == 404 + + def test_a_proxy_without_a_shell_offers_no_bash_and_keeps_no_skills(): with TestClient(ProxyAgent(shell=False, long_poll_seconds=0.2).create_app()) as client: client.post("/ext/proxy/claim", json={"token": TOKEN}) From d276f7180532b8f5d776d271dcda5216762cf39d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 26 Oct 2026 07:51:17 -0700 Subject: [PATCH 36/39] fix(load_artifact): a load aimed at an agent installed on a VM's host lands on the host load_artifact sent every load aimed at an agent by name into a container: the one install_agent put the agent in, or else the VM's default agent container. An agent installed on the VM's host has no container, so its file, URL, environment-payload and CLI loads went to agent-api, which a host install doesn't start, and failed with "No such container: agent-api". It now honours on_host as run_code, verify_sandbox and collect_artifacts do: files and CLI bundles go onto the host through A2AAgent's loaders (on_host), URLs through the host curl a sandbox-aimed load uses, and an environment payload is staged on the host. A CLI on a local sandbox's host, which is this machine, gets the gateway URL as it was issued. Loads aimed at an agent deployed or installed into a container send the same commands as before. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/a2a_agent/a2a_agent.py | 20 +++++++--- .../providers/sandbox_providers/sandbox.py | 17 +++++--- .../task_step/task_steps/load_artifact.py | 34 ++++++++++------ .../task_step/local_container_urls_test.py | 18 +++++---- ...st_load_artifact_agent_container_target.py | 15 ++++--- tst/unit/task_step/test_load_artifact_file.py | 40 +++++++++++++++---- 6 files changed, 101 insertions(+), 43 deletions(-) diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index 89b6720cb..adc1253b5 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -274,11 +274,12 @@ def _skill_extension(deployed: "DeployedA2AAgent") -> dict: @staticmethod async def install_cli( deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", gateway_url: str, container: str | None = None, + on_host: bool = False, ) -> str: """Install a CliArtifact bundle into the agent's filesystem at /opt/cli//, in ``container`` when - install_agent put the agent in one. Writes a sidecar bin/.env with the gateway URL and chmod +x's the - entrypoint. Idempotent. Returns the in-process path to the executable entrypoint.""" - from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container + install_agent put the agent in one, or on the VM host, ``on_host``, when it put it there. Writes a sidecar + bin/.env with the gateway URL and chmod +x's the entrypoint. Idempotent. Returns the path to the entrypoint.""" + from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container, stage_files_into_container from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider provider = build_sandbox_provider(deployed.sandbox_type) if deployed.sandbox_type else get_agent_sandbox_provider() sandbox = await provider.get_sandbox(deployed.sandbox_id) @@ -286,6 +287,14 @@ async def install_cli( entrypoint_path = f"{target_root}/{cli_artifact.entrypoint}" env_dir = posixpath.dirname(entrypoint_path) + if on_host: + await sandbox.write_host_file(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n".encode(), f"{env_dir}/{GATEWAY_ENV_FILE}") + await stage_files_into_container( + sandbox, cli_artifact.get_cli_files().get_file_artifacts(), target_root, on_host=True) + await sandbox.exec_script(f"chmod +x {shlex.quote(entrypoint_path)}") + logger.info(f"Installed CliArtifact '{cli_artifact.id}' v{cli_artifact.version} on the host at {entrypoint_path}") + return entrypoint_path + # write .env first — write_file_from_text creates the parent dir; write_file_from_object does not. await sandbox.write_file_from_text(f"AGENT_ENV_GATEWAY_URL={gateway_url}\n", f"{env_dir}/{GATEWAY_ENV_FILE}", **into_container(container)) @@ -309,9 +318,10 @@ async def load_file_artifact_universe( universe: "FileArtifactUniverse", destination_path: str, container: str | None = None, + on_host: bool = False, ) -> dict[str, str]: """Stage every FileArtifact in `universe` into the agent sandbox under destination_path, in ``container`` when - install_agent put the agent in one.""" + install_agent put the agent in one, or on the VM host, ``on_host``, when it put it there.""" from agent_env.providers.sandbox_providers.sandbox import stage_files_into_container from agent_env.providers.sandbox_providers.sandbox_provider import build_sandbox_provider, get_agent_sandbox_provider @@ -330,7 +340,7 @@ async def load_file_artifact_universe( f"Loading FileArtifactUniverse '{universe.id}' v{universe.version} " f"({len(file_artifacts)} file(s)) into agent at {destination}" ) - return await stage_files_into_container(sandbox, file_artifacts, destination, container) + return await stage_files_into_container(sandbox, file_artifacts, destination, container, on_host=on_host) def __init__( self, diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 29bb4d8e7..cb40c2386 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox.py +++ b/src/agent_env/providers/sandbox_providers/sandbox.py @@ -675,11 +675,12 @@ def into_container(container: str | None) -> dict[str, str]: async def stage_files_into_container( - sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str, container: str | None = None + sandbox: Sandbox, file_artifacts: dict[str, Any], destination: str, container: str | None = None, + on_host: bool = False, ) -> dict[str, str]: """Write each object-store file into the sandbox's container under ``destination``; returns ``{name: path}``. - A VM-backed sandbox stages through its host into ``container``, by default its ``container_name``; a container - sandbox is written directly.""" + A VM-backed sandbox stages through its host into ``container``, by default its ``container_name``, or, ``on_host``, + onto the host itself, for an agent install_agent put there; a container sandbox is written directly.""" loaded: dict[str, str] = {} dirs_to_make: set[str] = {destination} for filename in file_artifacts: @@ -689,7 +690,9 @@ async def stage_files_into_container( loaded[filename] = dest_path for d in sorted(dirs_to_make): - if isinstance(sandbox, VmSandbox): + if on_host: + await sandbox.exec_script(f"mkdir -p {shlex.quote(d)}") + elif isinstance(sandbox, VmSandbox): await sandbox.exec_script( f"docker exec -u 0 {shlex.quote(container or sandbox.container_name)} mkdir -p {shlex.quote(d)}") else: @@ -701,7 +704,11 @@ async def stage_files_into_container( async def _stage(filename: str, file_artifact: Any) -> None: async with sem: logger.info(f" {file_artifact.object_url} -> {loaded[filename]}") - await sandbox.write_file_from_object(file_artifact.object_url, loaded[filename], **into_container(container)) + if on_host: + await sandbox.load_object_file(file_artifact.object_url, loaded[filename]) + else: + await sandbox.write_file_from_object( + file_artifact.object_url, loaded[filename], **into_container(container)) await asyncio.gather(*(_stage(fn, fa) for fn, fa in file_artifacts.items())) return loaded diff --git a/src/agent_env/task_step/task_steps/load_artifact.py b/src/agent_env/task_step/task_steps/load_artifact.py index a4c22878b..185ac73ae 100644 --- a/src/agent_env/task_step/task_steps/load_artifact.py +++ b/src/agent_env/task_step/task_steps/load_artifact.py @@ -13,8 +13,10 @@ from typing import Callable, ClassVar, Optional from urllib.parse import unquote, urlparse +from agent_env.providers.sandbox_providers.local_sandbox import LocalSandbox +from agent_env.providers.sandbox_providers.sandbox import into_container from agent_env.providers.sandbox_providers.sandbox_provider import reachable_url -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import DeployedAgent, TaskStepContext from agent_env.entity.refs import EntityRef from agent_env.entity.scopes import FieldScope from agent_env.task_step.task_step import TaskStep, TaskStepDependency @@ -404,7 +406,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: EnvironmentUniverseArtifact, ) from agent_env.env.env import Env, require_gateway_url - from agent_env.providers.sandbox_providers.sandbox import VmSandbox, into_container + from agent_env.providers.sandbox_providers.sandbox import VmSandbox from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_VM, build_sandbox_provider, @@ -518,7 +520,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" files = await A2AAgent.load_file_artifact_universe( - deployed_agent, artifact, destination, **into_container(agent.installed_container)) + deployed_agent, artifact, destination, **_agent_files(agent)) context.metadata.setdefault("loaded_file_artifact_universes", []).append({ "id": artifact.id, "version": artifact.version, @@ -560,8 +562,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) - container = agent.installed_container or ( - sandbox.container_name if isinstance(sandbox, VmSandbox) else None) + container = None if agent.on_host else (agent.installed_container or ( + sandbox.container_name if isinstance(sandbox, VmSandbox) else None)) destination = (destination_path or "/app/files").rstrip("/") or "/" files = await _stage_environment_payload_into_container( sandbox, container, artifact, destination, @@ -610,10 +612,12 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if agent is None or not agent.instance_id: raise RuntimeError(f"Agent '{self.agent_name}' not found in context.deployed_agents (or missing instance_id)") deployed_agent = get_a2a_agent_instance_store().get(agent.instance_id) - gateway_url = reachable_url(require_gateway_url(deployed, "Installing a CliArtifact"), - from_sandbox_type=deployed.sandbox_type, to_sandbox_type=agent.sandbox_type) - install_path = await A2AAgent.install_cli( - deployed_agent, artifact, gateway_url, **into_container(agent.installed_container)) + gateway_url = require_gateway_url(deployed, "Installing a CliArtifact") + # A local sandbox's host is this machine, which reaches the gateway as it was issued. + if not (agent.on_host and agent.sandbox_type == LocalSandbox.type): + gateway_url = reachable_url(gateway_url, from_sandbox_type=deployed.sandbox_type, + to_sandbox_type=agent.sandbox_type) + install_path = await A2AAgent.install_cli(deployed_agent, artifact, gateway_url, **_agent_files(agent)) agent_clis = context.metadata.setdefault("installed_clis", {}).setdefault(self.agent_name, {}) agent_clis[artifact.id] = { "install_path": install_path, @@ -623,7 +627,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise ValueError(f"Unsupported artifact type for loading: '{artifact.type}'") if downloads: - agent_container = None + agent_container, agent_on_host = None, False if self.sandbox_name is not None: sandbox = target_sandbox else: @@ -636,7 +640,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: else get_agent_sandbox_provider() ) sandbox = await provider.get_sandbox(agent.sandbox_id) - agent_container = agent.installed_container + agent_container, agent_on_host = agent.installed_container, agent.on_host destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" semaphore = asyncio.Semaphore(8) @@ -650,7 +654,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: async def _load_one(url: str, filename: str) -> None: async with semaphore: dest = f"{destination}/{filename}" - if onto_vm_host: + if onto_vm_host or agent_on_host: await self._load_url_onto_vm(sandbox, url, dest) elif self.container_name is not None: await self._load_url_into_container(sandbox, sandbox.scoped_name(self.container_name), url, dest) @@ -670,6 +674,12 @@ async def _load_one(url: str, filename: str) -> None: return context +def _agent_files(agent: DeployedAgent) -> dict: + """How an A2AAgent write reaches ``agent``'s files: on the VM host for an agent install_agent put there, else as + ``into_container`` says.""" + return {"on_host": True} if agent.on_host else into_container(agent.installed_container) + + # Cap on how many staged paths get recorded in TaskStepContext metadata # (persisted, and shipped to external runners) — `file_count` always carries # the real total. diff --git a/tst/unit/task_step/local_container_urls_test.py b/tst/unit/task_step/local_container_urls_test.py index 273634340..9f06b04e6 100644 --- a/tst/unit/task_step/local_container_urls_test.py +++ b/tst/unit/task_step/local_container_urls_test.py @@ -315,20 +315,22 @@ async def test_a_remote_agent_gets_peer_urls_from_beyond_this_machine_as_they_ar @pytest.mark.asyncio -@pytest.mark.parametrize("sandbox_type, gateway_url, expected", [ - ("local", "http://127.0.0.1:41500", "http://host.docker.internal:41500"), - ("modal_vm", "https://gw.modal.host", "https://gw.modal.host"), +@pytest.mark.parametrize("sandbox_type, on_host, gateway_url, expected", [ + ("local", False, "http://127.0.0.1:41500", "http://host.docker.internal:41500"), + ("local", True, "http://127.0.0.1:41500", "http://127.0.0.1:41500"), + ("modal_vm", False, "https://gw.modal.host", "https://gw.modal.host"), + ("modal_vm", True, "https://gw.modal.host", "https://gw.modal.host"), ]) async def test_a_cli_gets_the_gateway_url_its_agent_reaches( - monkeypatch, local_stores, tmp_path, sandbox_type, gateway_url, expected, + monkeypatch, local_stores, tmp_path, sandbox_type, on_host, gateway_url, expected, ): (tmp_path / "cli").mkdir() (tmp_path / "cli" / "run.sh").write_text("#!/bin/sh\n") cli = CliArtifact.put("slack-cli", command_name="slack", entrypoint="run.sh", cli_dir=tmp_path / "cli") installed = [] - async def install_cli(deployed, cli_artifact, gateway_url): - installed.append(gateway_url) + async def install_cli(deployed, cli_artifact, gateway_url, **placed): + installed.append((gateway_url, placed)) return "/opt/cli/slack/run.sh" monkeypatch.setattr(A2AAgent, "install_cli", staticmethod(install_cli)) @@ -338,10 +340,10 @@ async def install_cli(deployed, cli_artifact, gateway_url): deployed_envs=[DeployedGatewayEnv(env_id="env-x", env_version=1, gateway_url=gateway_url, mcp_url="", sandbox_id="sb-env", sandbox_type=sandbox_type)], deployed_agents=[DeployedAgent(agent_name="solver", api_url="http://agent", sandbox_id="sb-agent", - sandbox_type=sandbox_type, instance_id="inst-1")], + sandbox_type=sandbox_type, instance_id="inst-1", on_host=on_host)], ) await LoadArtifactTaskStep(id="cli", version=None, env_id="env-x", artifact_id=cli.id, artifact_version=cli.version, agent_name="solver").execute(context) - assert installed == [expected] + assert installed == [(expected, {"on_host": True} if on_host else {})] diff --git a/tst/unit/task_step/test_load_artifact_agent_container_target.py b/tst/unit/task_step/test_load_artifact_agent_container_target.py index 5edcf03e7..153af421a 100644 --- a/tst/unit/task_step/test_load_artifact_agent_container_target.py +++ b/tst/unit/task_step/test_load_artifact_agent_container_target.py @@ -77,7 +77,9 @@ async def _fake_stage(sandbox, container_name, environment_artifact, destination return seen -def _context_with_agent(monkeypatch, sandbox, installed_container: str | None = None) -> TaskStepContext: +def _context_with_agent( + monkeypatch, sandbox, installed_container: str | None = None, on_host: bool = False, +) -> TaskStepContext: async def _get_sandbox(sandbox_id): return sandbox @@ -92,7 +94,7 @@ async def _get_sandbox(sandbox_id): ctx.deployed_agents = [ DeployedAgent( agent_name="solver", api_url="http://agent", sandbox_id=sandbox.sandbox_id, - sandbox_type=sandbox.type, installed_container=installed_container, + sandbox_type=sandbox.type, installed_container=installed_container, on_host=on_host, ) ] return ctx @@ -137,9 +139,10 @@ async def test_vm_backed_agent_container_keeps_its_per_sandbox_name( @pytest.mark.asyncio -@pytest.mark.parametrize("installed_in,expected", [(None, "agent-api"), ("agent-api", "agent-api"), ("task", "task")]) -async def test_plain_vm_agent_gets_the_default_container_name_or_the_one_it_was_installed_in( - monkeypatch, artifact, staged, installed_in, expected +@pytest.mark.parametrize("installed_in,on_host,expected", [ + (None, False, "agent-api"), ("agent-api", False, "agent-api"), ("task", False, "task"), (None, True, None)]) +async def test_plain_vm_agent_gets_the_default_container_name_the_one_it_was_installed_in_or_its_host( + monkeypatch, artifact, staged, installed_in, on_host, expected ): class _Vm(VmSandbox): type = "platform_vm" @@ -153,7 +156,7 @@ def __init__(self) -> None: async def terminate(self) -> None: # pragma: no cover - never called here pass - ctx = _context_with_agent(monkeypatch, _Vm(), installed_in) + ctx = _context_with_agent(monkeypatch, _Vm(), installed_in, on_host) step = LoadArtifactTaskStep( id="stage", version=None, agent_name="solver", artifact_id=ARTIFACT_ID, ) diff --git a/tst/unit/task_step/test_load_artifact_file.py b/tst/unit/task_step/test_load_artifact_file.py index 0a26311f7..036a70b78 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -15,7 +15,7 @@ from agent_env.env import env as env_module from agent_env.env.env import DeployedGatewayEnv, Env from agent_env.providers.sandbox_providers import sandbox_provider -from agent_env.providers.sandbox_providers.sandbox import VmSandbox +from agent_env.providers.sandbox_providers.sandbox import CURL_RETRY_FLAGS, VmSandbox from agent_env.task_step.context import DeployedAgent, DeployedSandbox, TaskStepContext from agent_env.task_step.task_steps.load_artifact import LoadArtifactTaskStep from tst.util.exec_scripts import script_run @@ -24,6 +24,8 @@ id="check", version=2, description="checker", filename="check.py", content_type="text/x-python", s3_url="s3://bucket/check.py", ) +TOOL = SimpleNamespace(id="tool", version=1, command_name="tool", entrypoint="bin/tool", + get_cli_files=lambda: SimpleNamespace(get_file_artifacts=lambda: {"bin/tool": CHECK})) class _Sandbox: @@ -182,7 +184,7 @@ async def load_object_file(self, object_url, destination_path): self.scripts.append(f"pull {object_url} {destination_path}") -async def _load_into_agent(monkeypatch, installed_in: str | None, **load) -> list[str]: +async def _load_into_agent(monkeypatch, installed_in: str | None, on_host: bool = False, **load) -> list[str]: """The scripts a load aimed at an agent runs on its VM, its staging paths as one name.""" sandbox = _Vm() monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) @@ -191,7 +193,7 @@ async def _load_into_agent(monkeypatch, installed_in: str | None, **load) -> lis sandbox_id="sb-agent"))) context = TaskStepContext() context.deployed_agents = [DeployedAgent(agent_name="solver", api_url="https://agent", sandbox_id="sb-agent", - instance_id="inst-1", installed_container=installed_in)] + instance_id="inst-1", installed_container=installed_in, on_host=on_host)] await LoadArtifactTaskStep(id="stage", version=None, agent_name="solver", **load).execute(context) return _unstaged(sandbox.scripts) @@ -216,14 +218,22 @@ async def test_into_an_agent_installed_into_a_container(monkeypatch, load): @pytest.mark.asyncio -async def test_a_cli_reaches_an_agent_installed_into_a_container(monkeypatch): - cli = SimpleNamespace(id="tool", version=1, command_name="tool", entrypoint="bin/tool", - get_cli_files=lambda: SimpleNamespace(get_file_artifacts=lambda: {"bin/tool": CHECK})) +@pytest.mark.parametrize("load, onto_the_host", [ + ({"artifact_id": "check"}, "pull s3://bucket/check.py /tmp/file_artifacts/check.py"), + ({"urls": ["https://example.com/check.py"]}, + f"curl -fsSL {CURL_RETRY_FLAGS} https://example.com/check.py -o /tmp/file_artifacts/check.py"), +], ids=["artifact", "url"]) +async def test_into_an_agent_installed_on_the_vm_host(monkeypatch, load, onto_the_host): + assert await _load_into_agent(monkeypatch, None, on_host=True, **load) == [ + "mkdir -p /tmp/file_artifacts", onto_the_host] + +@pytest.mark.asyncio +async def test_a_cli_reaches_an_agent_installed_into_a_container(monkeypatch): async def _install(installed_in: str | None) -> list[str]: sandbox = _Vm() monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) - await A2AAgent.install_cli(SimpleNamespace(sandbox_type=None, sandbox_id="sb-agent"), cli, "https://gateway", + await A2AAgent.install_cli(SimpleNamespace(sandbox_type=None, sandbox_id="sb-agent"), TOOL, "https://gateway", installed_in) return _unstaged(sandbox.scripts) @@ -232,3 +242,19 @@ async def _install(installed_in: str | None) -> list[str]: assert in_its_own == deployed assert "docker exec -u 0 task chmod +x /opt/cli/tool/bin/tool" in in_another assert not any("agent-api" in s for s in in_another) + + +@pytest.mark.asyncio +async def test_a_cli_installs_on_the_vm_host_of_an_agent_installed_there(monkeypatch): + sandbox = _Vm() + monkeypatch.setattr(sandbox_provider, "get_agent_sandbox_provider", _provider(sandbox)) + + path = await A2AAgent.install_cli(SimpleNamespace(sandbox_type=None, sandbox_id="sb-agent"), TOOL, "https://gateway", + on_host=True) + + assert path == "/opt/cli/tool/bin/tool" + assert [s for s in sandbox.scripts if not s.startswith("base64 -d")] == [ + "mkdir -p /opt/cli/tool/bin", "mkdir -p /opt/cli/tool", "mkdir -p /opt/cli/tool/bin", + "pull s3://bucket/check.py /opt/cli/tool/bin/tool", "chmod +x /opt/cli/tool/bin/tool"] + assert [s.split("\n")[0] for s in sandbox.scripts if s.startswith("base64 -d")] == [ + "base64 -d <<'ENDB64' > /opt/cli/tool/bin/.env"] From 6ba15e671346bf3aeb46c79d21f66c911f5dd17d Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 26 Oct 2026 09:02:12 -0700 Subject: [PATCH 37/39] fix(episode): a call that waits for a turn's tool listing waits as long as the proxy may take ProxyClient gave every request a 60 s read timeout, but the proxy's tools and record answer only once the turn's tools are listed, which it allows TOOL_TIMEOUT_SECONDS (600 s) for each env server. A slow listing made tools() give up after its three tries, and record_model_call(), which is never sent twice, raise while the proxy went on to record the call. Those two calls now wait the proxy's listing deadline and a minute more; the long-polls and the rest keep 60 s. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/task/episode/proxy.py | 9 +++++++-- tst/unit/task/episode/proxy_test.py | 27 +++++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/agent_env/task/episode/proxy.py b/src/agent_env/task/episode/proxy.py index bd9bd1e35..18e8307be 100644 --- a/src/agent_env/task/episode/proxy.py +++ b/src/agent_env/task/episode/proxy.py @@ -23,11 +23,14 @@ ProxyToolsResponse, ) from pydantic import BaseModel -from agentenv_protocol.a2a_agent.proxy import ProxyAgent +from agentenv_protocol.a2a_agent.proxy import TOOL_TIMEOUT_SECONDS, ProxyAgent from agent_env.a2a_agent.a2a_agent import A2AAgent _CALL_TIMEOUT = httpx.Timeout(60.0, connect=10.0) +# The calls that wait for a turn's tools to be listed, which the proxy allows TOOL_TIMEOUT_SECONDS an env server. +_LISTING_CALLS = frozenset({"tools", "record"}) +_LISTING_TIMEOUT = httpx.Timeout(TOOL_TIMEOUT_SECONDS + 60.0, connect=10.0) _SERVE_TIMEOUT_SECONDS = 30.0 # Short, so a long-poll in flight when the episode ends finishes within the server's shutdown. _IN_PROCESS_LONG_POLL_SECONDS = 1.0 @@ -103,7 +106,9 @@ async def _call(self, operation: str, **body: Any) -> dict[str, Any]: async def _post(self, operation: str, body: dict[str, Any], *, attempts: int = _ATTEMPTS) -> httpx.Response: for attempt in range(1, attempts + 1): try: - return await self._client.post(self._endpoints[operation], json={"token": self._token, **body}) + return await self._client.post( + self._endpoints[operation], json={"token": self._token, **body}, + timeout=_LISTING_TIMEOUT if operation in _LISTING_CALLS else _CALL_TIMEOUT) except httpx.TransportError: if attempt == attempts: raise diff --git a/tst/unit/task/episode/proxy_test.py b/tst/unit/task/episode/proxy_test.py index a5d4233c2..baf850d26 100644 --- a/tst/unit/task/episode/proxy_test.py +++ b/tst/unit/task/episode/proxy_test.py @@ -1,9 +1,10 @@ -"""The episode's client of a proxy: which calls it sends again when one doesn't reach the proxy.""" +"""The episode's client of a proxy: which calls it sends again when one doesn't reach the proxy, and how long it +waits for each.""" import httpx import pytest from agentenv_protocol.a2a_agent import ProxyTextPart, ProxyTool -from agentenv_protocol.a2a_agent.proxy import ProxyAgent +from agentenv_protocol.a2a_agent.proxy import LONG_POLL_SECONDS, TOOL_TIMEOUT_SECONDS, ProxyAgent from agent_env.task.episode import proxy as client_module from agent_env.task.episode.proxy import ProxyClient @@ -72,3 +73,25 @@ async def test_fields_a_newer_proxy_adds_are_ignored(): answer = await client.next(0) assert (answer.version, [t.parts[0].text for t in answer.turns]) == (2, ["Add one."]) + + +@pytest.mark.asyncio +async def test_a_call_that_waits_for_the_turns_listing_waits_as_long_as_the_proxy_may_take_to_list(): + waits: dict[str, float] = {} + + answers = {"tools": {"tools": []}, "next": {"version": 1, "turns": []}} + + def handle(request: httpx.Request) -> httpx.Response: + operation = request.url.path.rsplit("/", 1)[-1] + waits[operation] = request.extensions["timeout"]["read"] + return httpx.Response(200, json=answers.get(operation, {})) + + client = ProxyClient("http://proxy", CARD) + client._client = httpx.AsyncClient(transport=httpx.MockTransport(handle)) + await client.tools("task-1") + await client.record("task-1", "Adding one.", None, None) + await client.next(0) + await client.reply("task-1", parts=PARTS) + + assert waits["tools"] > TOOL_TIMEOUT_SECONDS and waits["record"] > TOOL_TIMEOUT_SECONDS + assert LONG_POLL_SECONDS < waits["next"] == waits["reply"] < TOOL_TIMEOUT_SECONDS From 2f00f10f6863197e16d21843aed9a4a2e9f1fd2c Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 26 Oct 2026 09:02:13 -0700 Subject: [PATCH 38/39] test(proxy): pull the container install test's images from mirror.gcr.io, and show docker's error The test started its alpine and debian containers from public.ecr.aws, which CI runners pull from anonymously and which throttles them: docker run exited 125 at setup in three of the last four runs, before anything was tested. It now pulls from mirror.gcr.io, as the other integration tests that pull on the runner do, and a container that doesn't start fails with docker's own error. Co-Authored-By: Claude Opus 5.5 --- tst/integration/agent/proxy_container_install_test.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tst/integration/agent/proxy_container_install_test.py b/tst/integration/agent/proxy_container_install_test.py index a9323bee8..f2ef19fd4 100644 --- a/tst/integration/agent/proxy_container_install_test.py +++ b/tst/integration/agent/proxy_container_install_test.py @@ -26,8 +26,9 @@ def container(request): image, shell = request.param name = f"proxy-{shell}-{uuid.uuid4().hex[:8]}" - subprocess.run(["docker", "run", "-d", "--name", name, "-p", f"127.0.0.1::{PORT}", - f"public.ecr.aws/docker/library/{image}", "sleep", "infinity"], check=True, capture_output=True) + started = subprocess.run(["docker", "run", "-d", "--name", name, "-p", f"127.0.0.1::{PORT}", + f"mirror.gcr.io/library/{image}", "sleep", "infinity"], capture_output=True, text=True) + assert started.returncode == 0, started.stderr try: yield name, shell finally: From 1194607f01afb20328b54af6d2d21e7a7d078ef8 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Mon, 26 Oct 2026 09:20:29 -0700 Subject: [PATCH 39/39] fix(protocol): a turn's env servers list their tools within one deadline for them all The proxy listed a turn's env servers one after another and gave each TOOL_TIMEOUT_SECONDS, so with several slow servers tools and record could answer well past the deadline the episode's client waits for. The listing as a whole now has TOOL_TIMEOUT_SECONDS: each server gets what is left of it, and one still listing when it passes is a 502. Co-Authored-By: Claude Opus 5.5 --- packages/agentenv-protocol/README.md | 3 ++- .../src/agentenv_protocol/a2a_agent/proxy.py | 8 +++++--- .../agentenv-protocol/tests/test_a2a_proxy.py | 18 ++++++++++++++++++ src/agent_env/task/episode/proxy.py | 2 +- 4 files changed, 26 insertions(+), 5 deletions(-) diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 9311f1ed1..8f9cf6000 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -713,7 +713,8 @@ operations: "skills"}]}`. A turn its sender cancels drops out of the list. - `tools` (`{"token", "task_id"}`) lists what the turn can call, with input schemas: the tools of each env server registered through mcp-config, and `bash`. Two servers offering - one name are refused with 409. + one name are refused with 409, and servers that haven't all listed their tools within + 600 s get 502. `record` waits for the same listing. - `act` (`{"token", "task_id", "action_id", "action"}`) runs one action: `{"kind": "tool", "server", "name", "arguments"}` calls an env server through `ToolSession`, with the URL and headers registered for it, and `{"kind": "bash", "command"}` runs `bash -c` in the diff --git a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py index beb75410f..9c520d30c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py +++ b/packages/agentenv-protocol/src/agentenv_protocol/a2a_agent/proxy.py @@ -288,14 +288,16 @@ def _open_turn(self, token: str, task_id: str) -> _Turn: return turn async def _tools(self, turn: _Turn) -> dict[str, ProxyTool]: - """What the turn can call, listed once: each registered env server's tools, and ``bash``. Two offering one - name are refused.""" + """What the turn can call, listed once and all within ``TOOL_TIMEOUT_SECONDS``: each registered env server's + tools, and ``bash``. Two offering one name are refused.""" async with turn.listing: if turn.tools is None: offered = {BASH: self._bash_tool} if self._shell else {} + loop = asyncio.get_running_loop() + deadline = loop.time() + TOOL_TIMEOUT_SECONDS for server in turn.servers: try: - listed = await asyncio.wait_for(self._listed(turn, server), TOOL_TIMEOUT_SECONDS) + listed = await asyncio.wait_for(self._listed(turn, server), deadline - loop.time()) except asyncio.TimeoutError: raise HTTPException(status_code=502, detail=( f"env server {server!r} didn't list its tools within {TOOL_TIMEOUT_SECONDS}s")) from None diff --git a/packages/agentenv-protocol/tests/test_a2a_proxy.py b/packages/agentenv-protocol/tests/test_a2a_proxy.py index 128027f03..fb30d86ba 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -329,6 +329,24 @@ def test_a_server_that_never_lists_its_tools_is_a_bad_gateway(client, monkeypatc assert "didn't list its tools within 0.5s" in response.text +def test_a_turns_servers_list_their_tools_within_one_deadline_for_them_all(client, monkeypatch): + monkeypatch.setattr(proxy, "TOOL_TIMEOUT_SECONDS", 1.0) + + async def slowly(self, turn, server): + await asyncio.sleep(0.6) + return [] + + monkeypatch.setattr(proxy.ProxyAgent, "_listed", slowly) + _register(client, "http://127.0.0.1:9/mcp", name="items") + _register(client, "http://localhost:9/mcp", name="pantry") + task_id = _send(client, "go") + + response = _proxy(client, "tools", task_id=task_id) + + assert response.status_code == 502 + assert "'pantry' didn't list its tools within 1.0s" in response.text + + def test_a_turns_actions_share_one_connection_to_each_env_server_which_closes_with_the_turn(client, env, monkeypatch): opened, closed, results = [], [], [] open_session, close_session = ToolSession.open, ToolSession.close diff --git a/src/agent_env/task/episode/proxy.py b/src/agent_env/task/episode/proxy.py index 18e8307be..315362218 100644 --- a/src/agent_env/task/episode/proxy.py +++ b/src/agent_env/task/episode/proxy.py @@ -28,7 +28,7 @@ from agent_env.a2a_agent.a2a_agent import A2AAgent _CALL_TIMEOUT = httpx.Timeout(60.0, connect=10.0) -# The calls that wait for a turn's tools to be listed, which the proxy allows TOOL_TIMEOUT_SECONDS an env server. +# The calls that wait for a turn's tools to be listed, which the proxy allows TOOL_TIMEOUT_SECONDS in all. _LISTING_CALLS = frozenset({"tools", "record"}) _LISTING_TIMEOUT = httpx.Timeout(TOOL_TIMEOUT_SECONDS + 60.0, connect=10.0) _SERVE_TIMEOUT_SECONDS = 30.0