diff --git a/.github/dependabot.yml b/.github/dependabot.yml index f7db6f04b..cf07a1f74 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -46,6 +46,7 @@ updates: - "/src/agent_env/agents/agents/opencode" - "/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 c7c955845..7906d2389 100644 --- a/.github/scripts/infra_build.py +++ b/.github/scripts/infra_build.py @@ -37,6 +37,7 @@ "agentenv-opencode": "src/agent_env/agents/agents/opencode/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}$") @@ -55,7 +56,7 @@ # 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, Gemini's, OpenCode's and pi's agents run their CLIs and -# import their servers, and OpenHands' imports its server, which loads the OpenHands SDK. +# 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"], @@ -74,6 +75,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 3cf5e06ba..681a9d3fb 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -223,9 +223,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 +- 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 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 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: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`. diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index 76ad0b97d..8f9cf6000 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -713,20 +713,27 @@ 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 - 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. 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`, 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 +742,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 5d33b6955..9c520d30c 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" @@ -75,6 +77,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,16 +134,19 @@ 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, 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] = {} 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): @@ -164,7 +174,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 +212,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}, 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}") @@ -217,14 +227,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 +243,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 = [(name, _canonical(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}) @@ -277,15 +288,17 @@ 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: BASH_TOOL} if self._shell else {} + 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) - except TimeoutError: + 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 except (ToolSessionError, httpx.HTTPError) as e: @@ -327,33 +340,45 @@ 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 - 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} + 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 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) + 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 _state(self, task: asyncio.Task[ProxyBashResult | ProxyToolResult]) -> ProxyActionState: + 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(): 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: @@ -367,6 +392,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]: @@ -471,8 +498,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], *, 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) + + +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 +521,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,28 +530,67 @@ 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: 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, - ) +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 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: - stdout, stderr = await asyncio.wait_for(process.communicate(), timeout) + 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) - 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) + 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) + 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: + """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 take(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 478bab14b..fb30d86ba 100644 --- a/packages/agentenv-protocol/tests/test_a2a_proxy.py +++ b/packages/agentenv-protocol/tests/test_a2a_proxy.py @@ -2,6 +2,9 @@ import asyncio import json +import os +import shutil +import signal import socket import subprocess import threading @@ -13,7 +16,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,17 +37,38 @@ 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.""" 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")) @@ -198,6 +222,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,8 +262,15 @@ 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["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" @@ -283,6 +315,38 @@ 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_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 @@ -319,9 +383,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): @@ -349,11 +429,70 @@ 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): + 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): @@ -369,6 +508,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}) @@ -386,8 +538,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 @@ -398,22 +552,45 @@ 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_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): diff --git a/src/agent_env/a2a_agent/a2a_agent.py b/src/agent_env/a2a_agent/a2a_agent.py index bf80dd2c8..f43f205f8 100644 --- a/src/agent_env/a2a_agent/a2a_agent.py +++ b/src/agent_env/a2a_agent/a2a_agent.py @@ -65,6 +65,10 @@ _DIND_CONTAINER = "agent-dind" _DIND_NETWORK = "agent-docker-net" _DIND_PORT = 2375 +# 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: @@ -127,6 +131,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" @@ -267,11 +272,14 @@ 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, + 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, 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) @@ -279,14 +287,25 @@ async def install_cli(deployed: "DeployedA2AAgent", cli_artifact: "CliArtifact", 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}") + 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) @@ -298,8 +317,11 @@ async def load_file_artifact_universe( deployed: "DeployedA2AAgent", 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.""" + """Stage every FileArtifact in `universe` into the agent sandbox under destination_path, in ``container`` when + 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 @@ -318,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) + return await stage_files_into_container(sandbox, file_artifacts, destination, container, on_host=on_host) def __init__( self, @@ -377,10 +399,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: @@ -397,16 +419,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: """The agent ``id`` names: the store's agent of that id, else the agent an installed bundle gives the alias @@ -493,16 +525,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"]) @@ -583,12 +606,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 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/agents/README.md b/src/agent_env/agents/README.md index 65b3a34b1..d9a8cd422 100644 --- a/src/agent_env/agents/README.md +++ b/src/agent_env/agents/README.md @@ -76,6 +76,16 @@ A release publishes Codex's, Gemini's, OpenCode's, pi's and OpenHands' images to them as it does Claude Code's: it pulls the image published for the sources installed, by its digest, and builds it from the agent's folder only when it can't. +`agents/proxy/` is the Episode API's, not an agent to pick for a task. It serves the protocol's `ProxyAgent`, +which has no model: an episode deploys it in the place of the agent its caller plays, and the caller's model acts +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; 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. + `agent-env run agents` runs every agent's smoke task, pulling or building each image; `--task claude-code-smoke` runs Claude Code's alone, `--task codex-smoke` Codex's, `--task gemini-smoke` Gemini's, `--task opencode-smoke` OpenCode's, -`--task pi-smoke` pi's, `--task openhands-smoke` OpenHands'. \ No newline at end of file +`--task pi-smoke` pi's, `--task openhands-smoke` OpenHands'. 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 4c9366937..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,6 +414,8 @@ 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], + workdir_host=_HOST_WORKDIR, + user_host="root", ), enable(TRIGGERS_V1, description=( "Register a typed trigger script and decide the stakeholder reply per turn " diff --git a/src/agent_env/agents/agents/codex/a2a_server.py b/src/agent_env/agents/agents/codex/a2a_server.py index dc3701f73..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,6 +610,8 @@ 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], + workdir_host=_HOST_WORKDIR, + user_host="root", ), enable( PEER_AGENTS_V1, diff --git a/src/agent_env/agents/agents/gemini/a2a_server.py b/src/agent_env/agents/agents/gemini/a2a_server.py index 702dab112..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,6 +403,8 @@ 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], + workdir_host=_HOST_WORKDIR, + 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..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,6 +292,8 @@ 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], + 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. PEER_AGENTS_V1, diff --git a/src/agent_env/agents/agents/proxy/Dockerfile b/src/agent_env/agents/agents/proxy/Dockerfile new file mode 100644 index 000000000..a10648162 --- /dev/null +++ b/src/agent_env/agents/agents/proxy/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 PYTHONPATH=/opt/agent +USER 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/proxy/a2a_server.py b/src/agent_env/agents/agents/proxy/a2a_server.py new file mode 100644 index 000000000..938885ccc --- /dev/null +++ b/src/agent_env/agents/agents/proxy/a2a_server.py @@ -0,0 +1,36 @@ +"""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.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 {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} {host_workdir} {host_user}"] +_INSTALL_COMMANDS = [_UNPACK, "bash {work_dir}/agent/install.sh container {container} {a2a_port}"] + + +@a2a_agent( + identity=IDENTITY, + config=ProxyConfig, + extensions=( + *EXTENSIONS, + 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", "host_workdir", "host_user"], + a2a_port=8000), + ), + lifespan=stopping_actions, +) +class InstalledProxy(ProxyAgent): + pass + + +agent = InstalledProxy() +app = agent.create_app() +AGENT_CARD = app.state.agentenv_a2a.card + +if __name__ == "__main__": + agent.serve() diff --git a/src/agent_env/agents/agents/proxy/agent.toml b/src/agent_env/agents/agents/proxy/agent.toml new file mode 100644 index 000000000..fce9bb754 --- /dev/null +++ b/src/agent_env/agents/agents/proxy/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/proxy/install.sh b/src/agent_env/agents/agents/proxy/install.sh new file mode 100644 index 000000000..17e701c3d --- /dev/null +++ b/src/agent_env/agents/agents/proxy/install.sh @@ -0,0 +1,67 @@ +#!/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 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 +# +# 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) + +case $mode in + host) + workdir=$4 user=$5 + 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 &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/agents/agents/proxy/requirements.txt b/src/agent_env/agents/agents/proxy/requirements.txt new file mode 100644 index 000000000..dfa7130a1 --- /dev/null +++ b/src/agent_env/agents/agents/proxy/requirements.txt @@ -0,0 +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.370 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/installed.py b/src/agent_env/bundle/installed.py index 39c7b2184..a04cff749 100644 --- a/src/agent_env/bundle/installed.py +++ b/src/agent_env/bundle/installed.py @@ -11,11 +11,13 @@ from agent_env.plugins._bundles import BUNDLES, CORE, InstalledBundle, installed_bundles -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 -__all__ = ["BUNDLES", "CORE", "InstalledBundle", "checked", "find_bundle", "installed_bundles", "run_name"] +__all__ = ["BUNDLES", "CORE", "InstalledBundle", "checked", "find_bundle", "installed_bundles", "run_name", + "write_installed"] def find_bundle(text: str, bundles: tuple[InstalledBundle, ...] | None = None) -> InstalledBundle: @@ -52,6 +54,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 _unknown(text: str, bundles: tuple[InstalledBundle, ...]) -> str: if not bundles: return f"{text!r} is neither a folder nor an installed bundle, and no bundles are installed" 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 853f7c2dc..4ea3160a6 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 @@ -193,6 +193,13 @@ def plan_bundle(resolved: ResolvedBundle, *, tasks: Sequence[str] = (), evals: S return _Planner(resolved, installed).plan(tasks, evals, agent) +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.""" @@ -493,7 +500,17 @@ 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) - closure = self._closure(selected_tasks + selected_evals) + 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), [], [], {}, None) + + def _planned(self, roots: list[ResolvedEntry], selected_tasks: list[ResolvedEntry], + 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: override = replace(override, version=store_latest.get((EntityKind.AGENT, override.agent_id))) @@ -627,6 +644,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/bundle/run.py b/src/agent_env/bundle/run.py index 9c52a8d76..28e9b1c7d 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 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, }; }); diff --git a/src/agent_env/providers/sandbox_providers/sandbox.py b/src/agent_env/providers/sandbox_providers/sandbox.py index 9ec670735..cb40c2386 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,19 @@ 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, + 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_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: @@ -678,8 +690,11 @@ async def stage_files_into_container(sandbox: Sandbox, file_artifacts: dict[str, loaded[filename] = dest_path 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)}") + 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: await sandbox.exec("mkdir", "-p", d) @@ -689,7 +704,11 @@ 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]) + 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/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/__init__.py b/src/agent_env/task/__init__.py index aa86a0635..a16aa8032 100644 --- a/src/agent_env/task/__init__.py +++ b/src/agent_env/task/__init__.py @@ -1,19 +1,20 @@ """Task module for AgentEnv.""" from .task import Task -from .episode import Episode, EpisodeError, EpisodeSpec, 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", "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/proxy.py b/src/agent_env/task/episode/proxy.py new file mode 100644 index 000000000..315362218 --- /dev/null +++ b/src/agent_env/task/episode/proxy.py @@ -0,0 +1,156 @@ +"""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, TypeVar + +import httpx +import uvicorn +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 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 in all. +_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 +_ATTEMPTS = 3 +_RETRY_SECONDS = 1.0 +_Response = TypeVar("_Response", bound=BaseModel) + + +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. 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. + 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) + 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) -> ProxyNextResponse: + return _read(ProxyNextResponse, await self._call("next", after=after)) + + 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: 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) + 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, *, 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) + return True + + async def close(self) -> None: + await self._client.aclose() + + 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._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 + await asyncio.sleep(_RETRY_SECONDS) + raise AssertionError("unreachable") + + +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 _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: + 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 8d766b428..5090f6efe 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 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. """ from __future__ import annotations @@ -33,10 +34,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,12 +49,22 @@ 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 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") + 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 [ 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), @@ -69,7 +80,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), @@ -94,12 +104,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 +123,17 @@ 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, "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 +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; 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/episode/run.py b/src/agent_env/task/episode/run.py index 8ee4ddfd4..44e088e3b 100644 --- a/src/agent_env/task/episode/run.py +++ b/src/agent_env/task/episode/run.py @@ -3,109 +3,154 @@ import asyncio import logging from collections import deque -from typing import TYPE_CHECKING, Any - -from agentenv_protocol import EnvironmentTool, ToolResult, ToolSession - +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 +from agent_env.store.routing import namespace_routing +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 TaskStepContext - -if TYPE_CHECKING: - from agent_env.task.episode.steps import _EpisodeTurnStep +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__) _live: dict[str, Episode] = {} +_DRAIN_RETRY_SECONDS = 1.0 +_CARD_TIMEOUT_SECONDS = 10.0 class EpisodeError(RuntimeError): """A task that can't run as an episode, or an episode asked for what it can't give.""" +# A skill the task gave the agent: its instructions are ``path``'s ``SKILL.md``, in the agent's sandbox. +Skill = ProxySkill + + 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()``, + 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], + 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._reply: asyncio.Future[str] = asyncio.get_running_loop().create_future() + self.skills = skills + self._proxy = proxy + self._task_id = task_id + 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.""" + """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.""" - 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) + 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=[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 - def _fail(self, error: Exception) -> None: + async def _fail(self, error: Exception) -> None: if self._open: - self._reply.set_exception(error) + self._open = False + 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 _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: + 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's envs, listed once when it opens: ``list()`` names every one, and ``call()`` routes a - call to the env that offers it.""" + """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._step_id = turn.step_id - 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._turn = turn + self._timeout = timeout self._tools: list[EnvironmentTool] = [] - self._offered_by: dict[str, str] = {} + 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) - 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.description, inputSchema=tool.input_schema)) + self._servers[tool.name] = tool.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._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) - - 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: ProxyBashAction | ProxyToolAction = ProxyBashAction(kind="bash", command=command) + else: + 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 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}"))) + return ToolResult(content=[{"type": "text", "text": text}], isError=failed) + return ToolResult(content=result.content, isError=result.is_error) class Episode: @@ -119,18 +164,22 @@ 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._proxies: list[ProxyClient] = [] + self._drains: list[asyncio.Task[None]] = [] + self._played: set[str] = set() + 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 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: @@ -145,14 +194,14 @@ 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" ) 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,10 +209,12 @@ 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 + for turn in self._turns: + turn._open = False self._run.cancel(msg=reason) await asyncio.wait((self._run,)) @@ -181,53 +232,118 @@ 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 proxy, 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_proxies() + 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.""" 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 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: + turn._open = False + def _raise_unless_completed(self) -> None: if self._abandoned is not None: raise EpisodeError(f"episode abandoned: {self._abandoned}") 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 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() + 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) -> 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 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) + 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="\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) + 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 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", {}) + 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 {}) + + async def _release_proxies(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(*(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 0ddddda03..1731edb42 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 = [_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: @@ -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 _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; ``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 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( - 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/steps.py b/src/agent_env/task/episode/steps.py index f6f55d332..9bcc85816 100644 --- a/src/agent_env/task/episode/steps.py +++ b/src/agent_env/task/episode/steps.py @@ -1,80 +1,65 @@ 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.bundle.installed import write_installed 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 +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" 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 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 = (EntityRef.env("env_ids[]"),) - - 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, - ): - 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 - - 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} - - @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"]) - - async def execute(self, context: TaskStepContext) -> TaskStepContext: - 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) - 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, + 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.prompt_id = prompt_id - self.prompt = prompt - self.system_prompt = system_prompt - self.timeout_seconds = timeout_seconds + self.step = step + self.sandboxed = sandboxed 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} + return {**super().to_dict(), "step": self.step.to_dict(), "sandboxed": self.sandboxed} @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"]) + def from_dict(cls, data: dict[str, Any]) -> _EpisodeAgentStep: + 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: - 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, - )) + episode = _episode_of(context) + agent = self.step.agent_name + if not self.sandboxed: + 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, 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) + deployed = next(deployed for deployed in context.deployed_agents if deployed.agent_name == agent) + await episode._attach(agent, deployed) return context diff --git a/src/agent_env/task_step/context.py b/src/agent_env/task_step/context.py index 48a0629ef..b9aa0a0d3 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. + installed_container: 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)), + 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 6b2d788f8..5cd475444 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,13 @@ 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. One installed into a + # container works in that one. + container = ( + agent.installed_container 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 1da700a29..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 @@ -43,6 +43,17 @@ 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" +# 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: @@ -66,6 +77,7 @@ def _build_param_resolvers( litellm_api_key_override: Optional[str] = None, sandbox_type: Optional[str] = None, trust_lines: str = "", + replaced: Optional[HostPlace] = None, ) -> dict[str, Callable[[], str]]: """Per-install registry: required_params name -> zero-arg resolver returning the raw value. @@ -86,6 +98,14 @@ def _resolve_workspace_dir() -> str: ) return workspace_dir + def _replaced(field: str) -> str: + if replaced is None: + raise KeyError( + 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 getattr(replaced, field) + def _resolve_container() -> str: if container_name is None: raise KeyError( @@ -102,6 +122,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. + "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()), @@ -231,8 +254,13 @@ 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) logger.info(f"Loading agent image {agent.docker_image_artifact.image_name} onto VM (needed to read agent card)...") @@ -292,6 +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, + replaced=replaced, ) resolved: dict[str, str] = {} for name in required_params: @@ -339,6 +368,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: a2a_card=agent_card, instance_id=deployed.instance_id, on_host=host_mode, + installed_container=container, )) default_model = agent.metadata.get("default_model") or get_config().get_model_for_role("agent") @@ -365,6 +395,27 @@ 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 _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.""" + 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 " + "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 {} + 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 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(), so store-resolving consumers (AddSkillsTaskStep, the a2a-agent CLIs) can find it.""" 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..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 @@ -517,7 +519,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, **_agent_files(agent)) context.metadata.setdefault("loaded_file_artifact_universes", []).append({ "id": artifact.id, "version": artifact.version, @@ -559,7 +562,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 = 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, @@ -608,9 +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) + 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, @@ -620,6 +627,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: raise ValueError(f"Unsupported artifact type for loading: '{artifact.type}'") if downloads: + agent_container, agent_on_host = None, False if self.sandbox_name is not None: sandbox = target_sandbox else: @@ -632,6 +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_on_host = agent.installed_container, agent.on_host destination = (destination_path or "/tmp/file_artifacts").rstrip("/") or "/" semaphore = asyncio.Semaphore(8) @@ -645,12 +654,12 @@ 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) 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)) @@ -665,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/src/agent_env/task_step/task_steps/run_code.py b/src/agent_env/task_step/task_steps/run_code.py index 450db74b3..bc8258f1a 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 @@ -453,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", 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) + return _AgentTarget(sandbox, prefix, agent.installed_container) # --- 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..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 @@ -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.installed_container 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/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..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 @@ -120,12 +120,18 @@ 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. One installed into a + # container works in that one. + container = ( + agent.installed_container or 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 +202,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_container_install_test.py b/tst/integration/agent/proxy_container_install_test.py new file mode 100644 index 000000000..f2ef19fd4 --- /dev/null +++ b/tst/integration/agent/proxy_container_install_test.py @@ -0,0 +1,100 @@ +"""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]}" + 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: + 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"] diff --git a/tst/integration/agent/proxy_modal_vm_test.py b/tst/integration/agent/proxy_modal_vm_test.py new file mode 100644 index 000000000..d1f8d9fa8 --- /dev/null +++ b/tst/integration/agent/proxy_modal_vm_test.py @@ -0,0 +1,84 @@ +"""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 +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 proxy 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_proxy_installed_in_its_place(): + uid = uuid.uuid4().hex[:8] + 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", + "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$"}]}, + {"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() + 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, "check-agent": 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="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 2481076c0..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,12 +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 and is never deployed; 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") @@ -236,13 +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 [agent["agent_name"] for agent in Task.get_instance(episode.id).context["deployed_agents"]] == ["echo"] + 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() @@ -252,9 +262,12 @@ 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:") + [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() @@ -265,7 +278,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 +312,69 @@ 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_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] + 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": "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, 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: + 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 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, "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, 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-*")) == [] + 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 ed7ef8742..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: @@ -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" @@ -112,7 +111,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): @@ -139,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}", 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): 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/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" diff --git a/tst/unit/agents/proxy/__init__.py b/tst/unit/agents/proxy/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/unit/agents/proxy/agent.py b/tst/unit/agents/proxy/agent.py new file mode 100644 index 000000000..522ab08dc --- /dev/null +++ b/tst/unit/agents/proxy/agent.py @@ -0,0 +1,12 @@ +"""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 +from pathlib import Path + +import agent_env + +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/proxy/test_proxy_agent.py b/tst/unit/agents/proxy/test_proxy_agent.py new file mode 100644 index 000000000..39a5fcdc2 --- /dev/null +++ b/tst/unit/agents/proxy/test_proxy_agent.py @@ -0,0 +1,64 @@ +"""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 + +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 +from agent_env.task_step.task_steps.install_agent import _build_param_resolvers + +from tst.unit.agents.proxy.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_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 params["install_commands"] + + +@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} {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, + container, script): + params = _extensions()["urn:agentenv:install/v1"]["params"] + 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(["bash", "-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(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(PROXY_AGENT_ID) == version 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" diff --git a/tst/unit/bundle/installed_entities_test.py b/tst/unit/bundle/installed_entities_test.py index f26cf2e87..a16bf49df 100644 --- a/tst/unit/bundle/installed_entities_test.py +++ b/tst/unit/bundle/installed_entities_test.py @@ -12,13 +12,14 @@ 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.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.a2a_agent import A2AAgent from agent_env.artifact.artifacts.docker_image import DockerImageArtifact from agent_env.artifact.store import get_artifact_store -from agent_env.bundle.plan import AliasUse, check_bundle, plan_bundle +from agent_env.bundle.plan import AliasUse, check_bundle, plan_bundle, plan_entities from agent_env.bundle.resolve import resolve_bundle from agent_env.plugins._bundles import ALIASES_FILE from agent_env.entity.refs import EntityKind @@ -191,6 +192,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 886a3c3f4..b32571b59 100644 --- a/tst/unit/bundle/materialize_test.py +++ b/tst/unit/bundle/materialize_test.py @@ -349,6 +349,7 @@ def refuse(ref, auth): [metadata] default_model = "claude-sonnet-4-6" min_disk_size_gb = 20 +takes_credentials = false """ @@ -361,7 +362,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, ()) @@ -374,11 +375,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"), ]) @@ -399,8 +405,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/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/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/store/aws/test_shared_credentials.py b/tst/unit/store/aws/test_shared_credentials.py index e67c7390b..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 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 @@ -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={TAKES_CREDENTIALS: False})) diff --git a/tst/unit/task/episode/conftest.py b/tst/unit/task/episode/conftest.py new file mode 100644 index 000000000..6afb74908 --- /dev/null +++ b/tst/unit/task/episode/conftest.py @@ -0,0 +1,12 @@ +"""An episode's turns reach a proxy 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/proxy_test.py b/tst/unit/task/episode/proxy_test.py new file mode 100644 index 000000000..baf850d26 --- /dev/null +++ b/tst/unit/task/episode/proxy_test.py @@ -0,0 +1,97 @@ +"""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 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 + +CARD = ProxyAgent().create_app().state.agentenv_a2a.card.model_dump(mode="json", by_alias=True, exclude_none=True) +PARTS = [ProxyTextPart(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") == [ProxyTool(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"] + + +@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."]) + + +@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 diff --git a/tst/unit/task/episode/run_test.py b/tst/unit/task/episode/run_test.py index 8f3c41590..4c67d8461 100644 --- a/tst/unit/task/episode/run_test.py +++ b/tst/unit/task/episode/run_test.py @@ -1,12 +1,31 @@ """The run of an episode, on fake steps: its turns, its scores, abandoning it, teardown and its record.""" import asyncio +import json +import threading +import time +import httpx import pytest import pytest_asyncio -from agentenv_protocol import EnvironmentTool, ToolResult +import uvicorn +from agentenv_protocol import EnvironmentTool, ToolResult, ToolSessionError +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 +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 @@ -104,21 +123,22 @@ 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): + def __init__(self, offered, raises=None): self.offered = offered + self.raises = raises or {} self.opened = [] 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): @@ -126,6 +146,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): @@ -156,7 +178,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"} @@ -181,9 +203,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} @@ -194,33 +215,25 @@ 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"] == [] - - -@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"}} + assert [(agent["agent_name"], agent["sandbox_id"]) for agent in record.context["deployed_agents"]] == [ + ("alice", None)] 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): - sessions = _ToolSessions(offered) - monkeypatch.setattr(episode_run, "ToolSession", sessions) +def _tool_sessions(monkeypatch, offered, raises=None): + sessions = _ToolSessions(offered, raises) + 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() @@ -228,26 +241,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() @@ -255,7 +269,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 @@ -294,7 +327,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 ."}, @@ -303,9 +337,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 @@ -322,32 +358,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) @@ -568,10 +604,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") - await _until(lambda: teardowns.finished) # the teardown ends a moment after the run is marked 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] @@ -587,3 +623,196 @@ 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 _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, results=None): + self.results = results or {} + self.actions = [] + self.replies = [] + + async def tools(self, task_id): + 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 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, *, parts=None, error=None): + self.replies.append((task_id, parts, error)) + return True + + +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_and_shows_its_output_and_a_failed_exit(): + 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"] + 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 [(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 [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 +async def test_a_turn_ends_with_its_reply_as_parts(): + proxy = _SandboxedProxy() + turn = _sandboxed_turn(proxy) + + await turn.end("Added a crate.") + await turn.end("Again.") + + assert proxy.replies == [("t-1", [ProxyTextPart(text="Added a crate.")], None)] + + +@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_recorded_model_calls_are_held_to_the_calls_that_run(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 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 turn.record_model_call("Added a crate.", model="qwen") + await turn.end("Added a crate.") + assert await episode.result() == {"replies": 1.0} + + [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_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() + + with pytest.raises(EpisodeError, match="OpenAI chat message, an Anthropic message or a string"): + 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"} + 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"]}) + 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..b6f154e23 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,93 +84,112 @@ 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] - 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) == ( + True, True, False, True, True, True, True] + 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_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_with_a_proxy_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_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_proxy_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'"), "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"]}}, "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 +198,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 +240,13 @@ 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},), + "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", "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..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, ) @@ -129,6 +131,9 @@ def test_a2a_port_roundtrip(): class _InstallSandbox: + ON_THIS_MACHINE = False + type = "modal_vm" + def __init__(self): self.scripts: list[str] = [] @@ -310,3 +315,89 @@ 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" +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 {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}} + + +def _installing_in_place(monkeypatch, *, codex=CODEX_INSTALL, sandbox=None, container_name=None): + """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", 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"{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")) + 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": {PROXY_PARAM: PROXY}}}, + "deployed_docker_containers": [{"container_name": "task", "sandbox_name": "vm"}]}, + ) + return step, context, sandbox + + +@pytest.mark.asyncio +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) + + 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.installed_container) for agent in context.deployed_agents] == [ + ("solver", True, None)] + + +@pytest.mark.asyncio +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"}}) + + with pytest.raises(RuntimeError, match="needs 'workdir_host' and 'user_host'"): + await step.execute(context) + + +@pytest.mark.asyncio +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"]}}) + + 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, agent.on_host, agent.installed_container) for agent in context.deployed_agents] == [ + ("solver", False, sandbox.scoped_name("task"))] + + +@pytest.mark.asyncio +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) + + 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="proxy"): + _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", replaced=HostPlace(workdir="/opt/agent-solver/workspace", user="agent-solver"), + ) + assert (placed["host_workdir"](), placed["host_user"]()) == ("/opt/agent-solver/workspace", "agent-solver") 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_container_sandbox_routing.py b/tst/unit/task_step/test_container_sandbox_routing.py index 0c6bb92d2..5daeefb1c 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; one it put in a container works there. """ from __future__ import annotations @@ -39,9 +40,10 @@ def _provide(monkeypatch, sandbox): monkeypatch.setattr(module, "get_sandbox_provider", lambda: provider) -def _agent_context(): +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")] + ctx.deployed_agents = [DeployedAgent(agent_name="solver", api_url="http://x", sandbox_id="sb-1", on_host=on_host, + installed_container=installed_container)] return ctx @@ -81,11 +83,17 @@ 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,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()) == expected + assert await _collect_container(monkeypatch, step, _agent_context(on_host, installed_in)) == expected @pytest.mark.asyncio @@ -102,11 +110,17 @@ 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,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()) == 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_load_artifact_agent_container_target.py b/tst/unit/task_step/test_load_artifact_agent_container_target.py index c71dfe7ab..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) -> 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, + sandbox_type=sandbox.type, installed_container=installed_container, on_host=on_host, ) ] return ctx @@ -137,8 +139,10 @@ 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,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" @@ -152,11 +156,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, on_host) 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..036a70b78 100644 --- a/tst/unit/task_step/test_load_artifact_file.py +++ b/tst/unit/task_step/test_load_artifact_file.py @@ -3,23 +3,29 @@ 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 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 CHECK = FileArtifact( 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: @@ -155,3 +161,100 @@ 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, 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)) + 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", installed_container=installed_in, on_host=on_host)] + 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 +@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"), TOOL, "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) + + +@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"] diff --git a/tst/unit/task_step/test_run_code.py b/tst/unit/task_step/test_run_code.py index ad544628a..53c3c50a8 100644 --- a/tst/unit/task_step/test_run_code.py +++ b/tst/unit/task_step/test_run_code.py @@ -658,6 +658,27 @@ 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].installed_container = installed_in + sandbox = _vm_sandbox_mock() + provider = MagicMock() + 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 async def test_timeout_is_reported(): with _sandbox([(124, "", "")]): # the runner's status at the limit 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 ----------------------- diff --git a/tst/unit/test_infra_images_script.py b/tst/unit/test_infra_images_script.py index 809ded55e..53ad217e3 100644 --- a/tst/unit/test_infra_images_script.py +++ b/tst/unit/test_infra_images_script.py @@ -60,7 +60,9 @@ 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-gemini", "agentenv-opencode", "agentenv-openhands", "agentenv-pi"] + assert "agentenv-gateway" in planned["infra"] and agents == [ + "agentenv-claude-code", "agentenv-codex", "agentenv-gemini", "agentenv-opencode", "agentenv-openhands", + "agentenv-pi", "agentenv-proxy"] def test_plan_missing_lists_the_images_the_registry_lacks_or_refuses_to_show(script, monkeypatch, capsys):