From 96c144e54f7688de3a02a15a61a1d85c8d279523 Mon Sep 17 00:00:00 2001 From: William Loo Date: Tue, 10 Nov 2026 14:27:12 -0800 Subject: [PATCH 1/5] feat(task-steps): native OSWorld initialize and evaluate on computer envs Adds the osworld_deploy_controller, osworld_initialize and osworld_evaluate steps and the osworld-server computer app (agent_env.osworld), so an OSWorld task (classic init_config/evaluator or an OSWorld-V2 task_NNN.py) runs on a framework computer env. - The OSWorld controller runs in a container sandbox of its own beside the computer, in remote mode, so the grader, a task's checks and the judge's key stay out of the agent's reach (a computer's apps run as its session user). - The computer runs only the OSWorld guest server, as the osworld-server app. Its card advertises urn:agentenv:osworld-guest/v1 at the prefix the app forwards to the guest server; the deploy step resolves the guest server's address from the computer's card (stored, else live until the app answers) and refuses a computer without the app before creating a sandbox. - osworld_initialize: OSWorld-V2 via /reset + /initialize_task; classic via /setup with the whole init_config. - osworld_evaluate: OSWorld-V2 via /evaluate; classic sends every func to the controller's /evaluate (async-polled, blocking fallback) and combines by conj; verifications shape matches cua_evaluate's. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/agent_env/osworld/__init__.py | 37 ++ src/agent_env/osworld/__main__.py | 29 ++ .../osworld/apps/osworld-server/env.toml | 5 + .../apps/osworld-server/osworld_server_app.py | 83 +++++ .../apps/osworld-server/requirements.txt | 10 + .../osworld-server/start-osworld-server.sh | 11 + src/agent_env/osworld/controller.py | 149 ++++++++ src/agent_env/osworld/guest.py | 73 ++++ src/agent_env/task/episode/rules.py | 3 +- src/agent_env/task_step/registry.py | 6 + .../task_step/task_steps/osworld/__init__.py | 3 + .../task_steps/osworld/deploy_controller.py | 110 ++++++ .../task_step/task_steps/osworld/evaluate.py | 222 ++++++++++++ .../task_steps/osworld/initialize.py | 74 ++++ tst/unit/osworld/__init__.py | 0 tst/unit/osworld/app_test.py | 104 ++++++ tst/unit/osworld/controller_test.py | 112 ++++++ tst/unit/osworld/guest_test.py | 58 +++ tst/unit/osworld/helpers.py | 62 ++++ tst/unit/osworld/steps_test.py | 339 ++++++++++++++++++ 20 files changed, 1489 insertions(+), 1 deletion(-) create mode 100644 src/agent_env/osworld/__init__.py create mode 100644 src/agent_env/osworld/__main__.py create mode 100644 src/agent_env/osworld/apps/osworld-server/env.toml create mode 100644 src/agent_env/osworld/apps/osworld-server/osworld_server_app.py create mode 100644 src/agent_env/osworld/apps/osworld-server/requirements.txt create mode 100755 src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh create mode 100644 src/agent_env/osworld/controller.py create mode 100644 src/agent_env/osworld/guest.py create mode 100644 src/agent_env/task_step/task_steps/osworld/__init__.py create mode 100644 src/agent_env/task_step/task_steps/osworld/deploy_controller.py create mode 100644 src/agent_env/task_step/task_steps/osworld/evaluate.py create mode 100644 src/agent_env/task_step/task_steps/osworld/initialize.py create mode 100644 tst/unit/osworld/__init__.py create mode 100644 tst/unit/osworld/app_test.py create mode 100644 tst/unit/osworld/controller_test.py create mode 100644 tst/unit/osworld/guest_test.py create mode 100644 tst/unit/osworld/helpers.py create mode 100644 tst/unit/osworld/steps_test.py diff --git a/src/agent_env/osworld/__init__.py b/src/agent_env/osworld/__init__.py new file mode 100644 index 000000000..360576b20 --- /dev/null +++ b/src/agent_env/osworld/__init__.py @@ -0,0 +1,37 @@ +"""OSWorld tasks on agent-env computers: the ``osworld-server`` computer app, which runs the OSWorld guest server on a +computer, and what the ``osworld_*`` steps share to reach it and the OSWorld controller deployed beside it. + +The guest server's source isn't part of agent-env; :func:`stage_server` copies the app's folder with that source into a +folder ``agent-env env put`` takes.""" +from __future__ import annotations + +import shutil +from pathlib import Path + +#: The app's ``environment_name``. +APP_NAME = "osworld-server" +#: The extension the app's card advertises at the prefix its server forwards to the guest server under. A computer +#: nests the app's card under its ``children_environments`` with the endpoint rooted where it serves the app, so a step +#: finds the guest server by this URI in the computer's card. +GUEST_EXTENSION = "urn:agentenv:osworld-guest/v1" +#: The prefix the app's server forwards to the guest server under. +PROXY_PREFIX = "/osworld" +#: The app's folder. +APP_DIR = Path(__file__).parent / "apps" / APP_NAME +#: Where a staged app keeps the guest server's source, and the files of that source it runs. +PAYLOAD_DIR = "payload" +PAYLOAD_FILES = ("main.py", "pyxcursor.py") + + +def stage_server(source: Path, dest: Path) -> Path: + """Copy the app's folder to ``dest`` with the guest server's source from ``source``, a folder holding + ``PAYLOAD_FILES``, under ``payload/``; returns ``dest``.""" + if missing := [name for name in PAYLOAD_FILES if not (source / name).is_file()]: + raise FileNotFoundError(f"{source} is not the OSWorld guest server's source: no {missing}") + if dest.exists(): + raise FileExistsError(f"{dest} exists") + shutil.copytree(APP_DIR, dest, ignore=shutil.ignore_patterns("__pycache__", "*.pyc")) + (dest / PAYLOAD_DIR).mkdir() + for name in PAYLOAD_FILES: + shutil.copy2(source / name, dest / PAYLOAD_DIR / name) + return dest diff --git a/src/agent_env/osworld/__main__.py b/src/agent_env/osworld/__main__.py new file mode 100644 index 000000000..5e898f44f --- /dev/null +++ b/src/agent_env/osworld/__main__.py @@ -0,0 +1,29 @@ +"""``python -m agent_env.osworld stage --source --out ``: the ``osworld-server`` app's +folder with the guest server's source, for ``agent-env env put``.""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from agent_env.osworld import stage_server + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="python -m agent_env.osworld") + commands = parser.add_subparsers(dest="command", required=True) + stage = commands.add_parser("stage", help="stage the osworld-server app with the guest server's source") + stage.add_argument("--source", required=True, type=Path, help="the folder holding the guest server's main.py") + stage.add_argument("--out", required=True, type=Path, help="the folder to write, which must not exist") + args = parser.parse_args(argv) + try: + out = stage_server(args.source, args.out) + except (FileNotFoundError, FileExistsError) as e: + print(e, file=sys.stderr) + return 1 + print(f"Staged {out}; put it with: agent-env env put {out}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/agent_env/osworld/apps/osworld-server/env.toml b/src/agent_env/osworld/apps/osworld-server/env.toml new file mode 100644 index 000000000..2203d2d98 --- /dev/null +++ b/src/agent_env/osworld/apps/osworld-server/env.toml @@ -0,0 +1,5 @@ +type = "computer_app" +environment_name = "osworld-server" +module = "osworld_server_app.py:OSWorldServer" +computer = "ubuntu-24.04" +driver = "reference" diff --git a/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py b/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py new file mode 100644 index 000000000..49a04df7d --- /dev/null +++ b/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py @@ -0,0 +1,83 @@ +"""The ``osworld-server`` computer app: the OSWorld guest server, which a CUA VM runs on port 5000, run on a computer +as the session user from the source staged under ``payload/`` (``agent_env.osworld.stage_server``), so the +``osworld_*`` steps set up and grade OSWorld tasks there. Its own server forwards ``/osworld/`` to the guest +server, and its card advertises that prefix as the endpoint of the ``osworld-guest/v1`` extension, which the computer +roots where it serves the app and guards with the env's control token. The OSWorld controller, deployed beside the +computer by ``osworld_deploy_controller``, reads the endpoint from the computer's card; an agent can't reach it there. + +This file runs on the computer, in the app's folder: it imports only what the app's server has.""" +from __future__ import annotations + +import socket +from typing import Any + +import httpx +from agentenv_protocol import extension +from agentenv_protocol.computer import (ComputerAppEnvironment, Packages, ReadyCheck, Service, computer_app, linux, + ready_check) +from starlette.requests import Request +from starlette.responses import Response + +GUEST_EXTENSION = "urn:agentenv:osworld-guest/v1" +PREFIX = "/osworld" +HOST, PORT = "127.0.0.1", 5000 +UPSTREAM = f"http://{HOST}:{PORT}" +METHODS = ["GET", "POST", "PUT", "PATCH", "DELETE"] +#: Setup downloads task assets and evaluation runs long checks; the controller allows each up to 1800 s. +TIMEOUT_S = 1800.0 +#: Not forwarded: the hop's own headers, and the control token, which is the computer's. +_DROPPED_REQUEST = frozenset({"host", "content-length", "connection", "transfer-encoding", "authorization"}) +_DROPPED_RESPONSE = frozenset({"content-length", "connection", "transfer-encoding", "content-encoding"}) +#: What a CUA VM installs for the guest server that Ubuntu 24.04's desktop computer lacks: AT-SPI's Python bindings +#: for its accessibility tree, Tk for pyautogui, the window and input tools its setup routes run, ffmpeg for recordings. +PACKAGES = ["python3-pyatspi", "python3-gi", "gir1.2-atspi-2.0", "at-spi2-core", "python3-tk", "ffmpeg", "wmctrl", + "xdotool", "x11-xkb-utils", "x11-xserver-utils"] + + +async def forward(request: Request, client: httpx.AsyncClient, upstream: str = UPSTREAM) -> Response: + """``request`` under ``PREFIX``, sent to the guest server with its method, query and body, and its answer + returned as the guest server gave it.""" + headers = {k: v for k, v in request.headers.items() if k.lower() not in _DROPPED_REQUEST} + try: + answer = await client.request(request.method, f"{upstream}/{request.path_params.get('path', '')}", + params=list(request.query_params.multi_items()), headers=headers, + content=await request.body(), timeout=TIMEOUT_S) + except httpx.HTTPError as e: + return Response(f"osworld-server unreachable: {type(e).__name__}", status_code=502) + return Response(answer.content, status_code=answer.status_code, + headers={k: v for k, v in answer.headers.items() if k.lower() not in _DROPPED_RESPONSE}) + + +def guest_server_check() -> ReadyCheck: + try: + with socket.create_connection((HOST, PORT), timeout=2): + return ReadyCheck("osworld-server", True) + except OSError as e: + return ReadyCheck("osworld-server", False, f"{HOST}:{PORT}: {e}") + + +@computer_app(supports=[{"os": {"type": "linux", "display": ["x11"]}, "arch": ["amd64", "arm64"]}]) +@linux(packages=Packages(apt=PACKAGES), python="requirements.txt", + services=[Service(name="osworld-server", command="start-osworld-server.sh")]) +class OSWorldServer(ComputerAppEnvironment): + """The OSWorld guest server, forwarded under ``/osworld/`` to the controller beside the computer.""" + + def create_app(self) -> Any: + app = super().create_app() + client = httpx.AsyncClient() + + async def _forward(request: Request) -> Response: + return await forward(request, client) + + app.custom_route(f"{PREFIX}/{{path:path}}", methods=METHODS)(_forward) + return app + + @extension(GUEST_EXTENSION, method="GET", path=PREFIX, + description="The OSWorld guest server's routes, under this endpoint; a GET says whether it's up.") + async def guest(self) -> dict: + check = guest_server_check() + return {"ready": check.ok, "detail": check.detail} + + @ready_check(os="linux") + def guest_server(self) -> ReadyCheck: + return guest_server_check() diff --git a/src/agent_env/osworld/apps/osworld-server/requirements.txt b/src/agent_env/osworld/apps/osworld-server/requirements.txt new file mode 100644 index 000000000..9cd169f3d --- /dev/null +++ b/src/agent_env/osworld/apps/osworld-server/requirements.txt @@ -0,0 +1,10 @@ +# The app's own server forwards /osworld/* to the guest server with httpx; the rest is what the guest server imports +# on Linux. AT-SPI's bindings come from the OS's packages (python3-pyatspi, python3-gi). +httpx>=0.27.0 +python3-xlib==0.15 +PyAutoGUI==0.9.54 +Pillow==10.1.0 +requests +flask +numpy +lxml diff --git a/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh b/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh new file mode 100755 index 000000000..dcdc4ae10 --- /dev/null +++ b/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +# Runs the OSWorld guest server on 127.0.0.1:5000 with the app's venv, restarting it when it exits. The OS's +# dist-packages go after the venv's, for AT-SPI's bindings, which pip can't install. +cd "$(dirname "$0")/payload" +export DISPLAY="${DISPLAY:-:0}" +export PYTHONUNBUFFERED=1 +while true; do + ../venv/bin/python -c 'import runpy, sys; sys.path.append("/usr/lib/python3/dist-packages"); sys.argv = ["main.py"]; runpy.run_path("main.py", run_name="__main__")' + echo "osworld-server exited ($?); restarting in 5s" >&2 + sleep 5 +done diff --git a/src/agent_env/osworld/controller.py b/src/agent_env/osworld/controller.py new file mode 100644 index 000000000..109ac25de --- /dev/null +++ b/src/agent_env/osworld/controller.py @@ -0,0 +1,149 @@ +"""The OSWorld controller beside a computer, as the ``osworld_*`` steps call it over HTTP: ``/reset``, +``/initialize_task``, ``/setup``, and ``/evaluate``, the last started with ``/evaluate_async`` and polled with short +requests, since a judge can run for minutes and a long silent connection is reaped by what's between.""" +from __future__ import annotations + +import asyncio +import json +import logging +from typing import Any + +import httpx + +from agent_env.config import get_config + +logger = logging.getLogger(__name__) + +#: How long one controller call may take: OSWorld-V2 ``setup()`` downloading assets, or a judge. +DEFAULT_TIMEOUT_S = 1800 +#: The controller has no health route; ``GET /configure`` is a read that answers once its app is up. +READY_PATH = "/configure" +_POLL_INTERVAL_S = 15 +_POLL_READ_TIMEOUT_S = 30 +#: A task ref the controller fetches is signed for this long, covering one step's download window. +TASK_REF_TTL_S = 6 * 60 * 60 + + +class ControllerCallError(RuntimeError): + """A controller call that failed in transport or answered with an HTTP error; carries the status and body.""" + + def __init__(self, message: str, *, status: int | None = None, body: str | None = None): + super().__init__(message) + self.status = status + self.body = body + + +def _parsed(path: str, response: httpx.Response, *, expect_json: bool) -> dict: + if response.status_code >= 400: + raise ControllerCallError(f"controller {path} returned HTTP {response.status_code}: {response.text[:500]}", + status=response.status_code, body=response.text) + try: + return json.loads(response.text) + except json.JSONDecodeError as e: + if not expect_json: + return {} + raise ControllerCallError(f"controller {path} returned non-JSON: {e}: {response.text[:300]}", + status=response.status_code, body=response.text) from e + + +async def post(controller_url: str, path: str, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S, + expect_json: bool = True) -> dict: + """POST ``payload`` to the controller's ``path``. With ``expect_json`` False a 2xx body that isn't JSON (``/reset`` + answers a bare string) returns ``{}``. Raises ControllerCallError on a transport failure or an HTTP error.""" + try: + async with httpx.AsyncClient() as client: + response = await client.post(f"{controller_url.rstrip('/')}{path}", json=payload, timeout=timeout) + except httpx.HTTPError as e: + raise ControllerCallError(f"controller {path} transport failure: {type(e).__name__}: {e}", body=str(e)) from e + return _parsed(path, response, expect_json=expect_json) + + +async def evaluate(controller_url: str, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S) -> dict: + """``/evaluate`` with ``payload``: started with ``/evaluate_async`` and polled at ``/evaluate_result`` until it's + done, or the blocking ``/evaluate`` on a controller without the async routes.""" + base = controller_url.rstrip("/") + try: + async with httpx.AsyncClient() as client: + start = await client.post(f"{base}/evaluate_async", json=payload, timeout=60) + except httpx.HTTPError as e: + raise ControllerCallError(f"controller /evaluate_async transport failure: {type(e).__name__}: {e}", + body=str(e)) from e + if start.status_code == 404: + logger.info("controller has no /evaluate_async; evaluating with the blocking /evaluate") + return await post(controller_url, "/evaluate", payload, timeout=timeout) + job_id = _parsed("/evaluate_async", start, expect_json=True).get("job_id") + if not job_id: + raise ControllerCallError(f"controller /evaluate_async returned no job_id: {start.text[:300]}") + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while True: + if loop.time() >= deadline: + raise ControllerCallError(f"controller /evaluate job {job_id} still running after {timeout}s") + await asyncio.sleep(_POLL_INTERVAL_S) + try: + async with httpx.AsyncClient() as client: + poll = await client.get(f"{base}/evaluate_result", params={"job_id": job_id}, + timeout=_POLL_READ_TIMEOUT_S) + except httpx.HTTPError as e: + logger.warning(f"controller /evaluate_result poll failed, polling again: {type(e).__name__}: {e}") + continue + if poll.status_code >= 400: + raise ControllerCallError(f"controller /evaluate_result returned HTTP {poll.status_code}: " + f"{poll.text[:500]}", status=poll.status_code, body=poll.text) + try: + data = json.loads(poll.text) or {} + except json.JSONDecodeError: + logger.warning(f"controller /evaluate_result returned non-JSON, polling again: {poll.text[:200]}") + continue + status = data.get("status") + if status == "done": + return data.get("result") or {} + if status in ("error", "unknown"): + raise ControllerCallError(f"controller /evaluate job {job_id} {status}: {str(data.get('message', ''))[:300]}", + body=json.dumps(data)[:500]) + + +async def wait_ready(controller_url: str, *, timeout: float = 300.0, interval: float = 5.0) -> None: + """Return once the controller answers ``READY_PATH``; TimeoutError after ``timeout``.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + last = None + while loop.time() < deadline: + try: + async with httpx.AsyncClient() as client: + response = await client.get(f"{controller_url.rstrip('/')}{READY_PATH}", timeout=15) + if response.status_code == 200: + return + last = f"HTTP {response.status_code}" + except httpx.HTTPError as e: + last = repr(e) + await asyncio.sleep(interval) + raise TimeoutError(f"OSWorld controller {controller_url} not ready within {timeout}s (last: {last})") + + +def resolve_task_ref(ref: str) -> str: + """An OSWorld-V2 task ref as the controller fetches it, with no credentials: an http(s) URL as it is, an object of + the configured object store under its artifacts as a signed URL. Anything else raises, rather than hand the + controller a URL it can't fetch or sign an object outside the artifacts.""" + if ref.startswith(("http://", "https://")): + return ref + config = get_config() + store = config.get_object_store_at(ref) + prefix = f"{config.get_artifact_key_prefix()}artifacts/" + if not store.owns(ref) or not store.get_object_key(ref).startswith(prefix): + raise ValueError(f"refusing to sign task ref {ref!r}: not an object under {prefix!r} in the object store") + if (signed := store.signed_get_url(ref, expires_in=TASK_REF_TTL_S)) is None: + raise ValueError(f"the object store can't sign {ref!r} for the controller to fetch") + return signed + + +def normalize_score(score: Any) -> float: + """A controller score as a float: a bool as 1 or 0, a number as it is, a dict as the mean of its numbers, else 0.""" + if isinstance(score, bool): + return 1.0 if score else 0.0 + if isinstance(score, (int, float)): + return float(score) + if isinstance(score, dict): + numbers = [float(v) for v in score.values() if isinstance(v, (int, float)) and not isinstance(v, bool)] + return sum(numbers) / len(numbers) if numbers else 0.0 + return 0.0 diff --git a/src/agent_env/osworld/guest.py b/src/agent_env/osworld/guest.py new file mode 100644 index 000000000..6cb37e09b --- /dev/null +++ b/src/agent_env/osworld/guest.py @@ -0,0 +1,73 @@ +"""Where a computer serves the OSWorld guest server: the endpoint its card advertises for the ``osworld-guest/v1`` +extension, on the card itself or one of its ``children_environments``, resolved against the computer's address. +``osworld_deploy_controller`` resolves it once and records it on the env, with the controller's address, under the +metadata keys here.""" +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +import httpx +from agentenv_protocol import client as protocol_v1 + +from agent_env.osworld import GUEST_EXTENSION + +logger = logging.getLogger(__name__) + +#: The ``DeployedEnv.metadata`` keys ``osworld_deploy_controller`` records on a computer env. +CONTROLLER_URL_KEY = "osworld_controller_url" +CONTROLLER_SANDBOX_ID_KEY = "osworld_controller_sandbox_id" +GUEST_URL_KEY = "osworld_guest_url" +#: A computer's card lists an app without its extensions until the app's server answers: how long resolving waits for +#: the guest extension, and how often it reads the card again. +CARD_WAIT_S = 120.0 +CARD_POLL_S = 5.0 + + +def guest_endpoint(card: dict | None) -> str | None: + """The endpoint ``card`` advertises for the guest extension, its own or a child's; None when nothing does.""" + if not isinstance(card, dict): + return None + for candidate in [card, *(child.card for child in protocol_v1.children(card))]: + params = (protocol_v1.find_extension(candidate, GUEST_EXTENSION) or {}).get("params") or {} + if isinstance(params.get("endpoint"), str): + return params["endpoint"] + return None + + +def _joined(environment_url: str, endpoint: str) -> str: + if "://" in endpoint: + return endpoint.rstrip("/") + return f"{environment_url.rstrip('/')}/{endpoint.lstrip('/')}".rstrip("/") + + +async def resolve_guest_url(deployed: Any, *, wait_s: float | None = None, poll_s: float | None = None) -> str: + """The guest server's base URL on ``deployed``, from the card it was deployed with, else the card it serves now, + read again until the app's server answers or ``wait_s`` passes. Raises RuntimeError for an env with no address or + no app advertising the guest extension.""" + environment_url = deployed.environment_url + if not environment_url: + raise RuntimeError(f"env '{deployed.env_id}' has no address to read its card at") + endpoint = guest_endpoint(deployed.environment_card) + loop = asyncio.get_running_loop() + deadline = loop.time() + (CARD_WAIT_S if wait_s is None else wait_s) + while endpoint is None: + try: + endpoint = guest_endpoint(await protocol_v1.get_card(environment_url)) + except (httpx.HTTPError, ValueError) as e: + logger.info(f"env '{deployed.env_id}': card not readable yet: {type(e).__name__}") + if endpoint is not None or loop.time() >= deadline: + break + await asyncio.sleep(CARD_POLL_S if poll_s is None else poll_s) + if endpoint is None: + raise RuntimeError(f"env '{deployed.env_id}' advertises no {GUEST_EXTENSION} extension on its card: " + f"compose the osworld-server computer app into its computer") + return _joined(environment_url, endpoint) + + +def controller_url(deployed: Any) -> str: + """The controller ``osworld_deploy_controller`` recorded on ``deployed``.""" + if url := (deployed.metadata or {}).get(CONTROLLER_URL_KEY): + return str(url) + raise RuntimeError(f"env '{deployed.env_id}' has no {CONTROLLER_URL_KEY}: run osworld_deploy_controller on it first") diff --git a/src/agent_env/task/episode/rules.py b/src/agent_env/task/episode/rules.py index 0bfb1d51d..6a2002316 100644 --- a/src/agent_env/task/episode/rules.py +++ b/src/agent_env/task/episode/rules.py @@ -153,7 +153,8 @@ def _rubrics_verifier(step: RubricsVerifierTaskStep, outside: Outside) -> list[s "browser_close_tab", "browser_delete_cookies", "browser_open_tab", "browser_set_cookies", "browser_sign_in", "browser_set_local_storage", "build_mcp_cli", "combine_universe_verdicts", "configure_scheduled_checks", "deploy_env", "deploy_sandbox", "end_computer_recording", "start_computer_recording", "env_outcome_verifier", - "modify_env_tool_access", "reset_env", "review", + "modify_env_tool_access", "osworld_deploy_controller", "osworld_evaluate", "osworld_initialize", "reset_env", + "review", "run_container_unit_tests_verifier", "run_docker_container", "snapshot_env", "snapshot_computer_env", "sync_env_clock", "validation_gate_aggregator", "verify_env_card", "verify_env_core_protocol", "verify_mcp_tool_schema", "verify_spec_conformance", "verify_universe_load_export_roundtrip", diff --git a/src/agent_env/task_step/registry.py b/src/agent_env/task_step/registry.py index d428aeb41..9dd693329 100644 --- a/src/agent_env/task_step/registry.py +++ b/src/agent_env/task_step/registry.py @@ -90,6 +90,9 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: from agent_env.task_step.task_steps.snapshot_agent_state import SnapshotAgentStateTaskStep from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep from agent_env.task_step.task_steps.computer.snapshot_computer_env import SnapshotComputerEnvTaskStep + from agent_env.task_step.task_steps.osworld.deploy_controller import OSWorldDeployControllerTaskStep + from agent_env.task_step.task_steps.osworld.evaluate import OSWorldEvaluateTaskStep + from agent_env.task_step.task_steps.osworld.initialize import OSWorldInitializeTaskStep from agent_env.task_step.task_steps.peer_agents import PeerAgentsTaskStep from agent_env.task_step.task_steps.verifiers.verify_sandbox import VerifySandboxTaskStep from agent_env.task_step.task_steps.verifiers.aggregate_verifiers import AggregateVerifiersTaskStep @@ -150,6 +153,9 @@ def _builtin_registry() -> dict[str, type["TaskStep"]]: _get_type(SnapshotAgentStateTaskStep): SnapshotAgentStateTaskStep, _get_type(SnapshotEnvTaskStep): SnapshotEnvTaskStep, _get_type(SnapshotComputerEnvTaskStep): SnapshotComputerEnvTaskStep, + _get_type(OSWorldDeployControllerTaskStep): OSWorldDeployControllerTaskStep, + _get_type(OSWorldInitializeTaskStep): OSWorldInitializeTaskStep, + _get_type(OSWorldEvaluateTaskStep): OSWorldEvaluateTaskStep, _get_type(RunCodeTaskStep): RunCodeTaskStep, _get_type(PeerAgentsTaskStep): PeerAgentsTaskStep, _get_type(VerifyA2ASnapshotStep): VerifyA2ASnapshotStep, diff --git a/src/agent_env/task_step/task_steps/osworld/__init__.py b/src/agent_env/task_step/task_steps/osworld/__init__.py new file mode 100644 index 000000000..a0e783ffd --- /dev/null +++ b/src/agent_env/task_step/task_steps/osworld/__init__.py @@ -0,0 +1,3 @@ +"""OSWorld tasks on agent-env computers: ``osworld_deploy_controller`` deploys the OSWorld controller beside a computer +that runs the ``osworld-server`` app, ``osworld_initialize`` sets a task up through it and ``osworld_evaluate`` grades +it there (``agent_env.osworld``).""" diff --git a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py new file mode 100644 index 000000000..7307af0b1 --- /dev/null +++ b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py @@ -0,0 +1,110 @@ +"""``osworld_deploy_controller``: the OSWorld controller beside a computer env, in a container sandbox of its own. + +The controller holds the grader, a task's checks and the judge's key, and a computer's apps run as its session user, +the agent's, so it runs apart from the computer, in remote mode against the guest server at the endpoint the +computer's card advertises (``agent_env.osworld.guest``), with the env's control token on every call. Once it answers, +the step records the controller's and the guest server's addresses on the env for ``osworld_initialize`` and +``osworld_evaluate``, and the sandbox in ``context.deployed_sandboxes`` for teardown.""" +from __future__ import annotations + +import json +import logging +from typing import ClassVar, Optional + +from agent_env.artifact import Artifact, DockerImageArtifact +from agent_env.artifact.artifacts.docker_image import pull_problem +from agent_env.attribution import deploy_attribution +from agent_env.config import ConfigError, get_config +from agent_env.entity.refs import EntityRef +from agent_env.osworld import controller, guest +from agent_env.providers.sandbox_providers.chained_sandbox_provider import skip_refusing +from agent_env.providers.sandbox_providers.sandbox_provider import get_sandbox_provider +from agent_env.task_step.context import DeployedSandbox, TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency + +logger = logging.getLogger(__name__) + +#: The controller image's artifact id when a step names none. +DEFAULT_CONTROLLER_IMAGE_ARTIFACT_ID = "cua-controller" +CONTROLLER_PORT = 8000 +#: A controller's sandbox name, by its computer env's id. +SANDBOX_NAME_PREFIX = "osworld-controller-" + + +class OSWorldDeployControllerTaskStep(TaskStep): + type: ClassVar[str] = "osworld_deploy_controller" + entity_refs = (EntityRef.env("env_id"), + EntityRef.artifact("controller_image_artifact_id", artifact_type="docker_image")) + + def __init__(self, id: str, version: Optional[int], env_id: str, controller_image_artifact_id: Optional[str] = None, + ttl_seconds: int = TaskStep.DEFAULT_TTL_SECONDS, cpu: float = 2.0, memory_mb: int = 4096, + env_vars: Optional[dict[str, str]] = None, depends_on: Optional[list[TaskStepDependency]] = None, + fail_task_on_error: bool = True): + """``env_vars`` go to the controller beside the ones the step sets, which win: a deployment's own, the + judge's model (``JUDGE_MODEL``) say.""" + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.env_id = env_id + self.controller_image_artifact_id = controller_image_artifact_id + self.ttl_seconds = ttl_seconds + self.cpu = cpu + self.memory_mb = memory_mb + self.env_vars = env_vars + + def to_dict(self) -> dict: + return {**super().to_dict(), "env_id": self.env_id, + "controller_image_artifact_id": self.controller_image_artifact_id, "ttl_seconds": self.ttl_seconds, + "cpu": self.cpu, "memory_mb": self.memory_mb, "env_vars": self.env_vars} + + @classmethod + def from_dict(cls, data: dict) -> OSWorldDeployControllerTaskStep: + return cls(**cls._base_from_dict(data), env_id=data["env_id"], + controller_image_artifact_id=data.get("controller_image_artifact_id"), + ttl_seconds=data.get("ttl_seconds", TaskStep.DEFAULT_TTL_SECONDS), cpu=data.get("cpu", 2.0), + memory_mb=data.get("memory_mb", 4096), env_vars=data.get("env_vars")) + + def controller_env(self, deployed, guest_url: str) -> dict[str, str]: + """Remote mode against the guest server, with the env's control token on every call and the model endpoint's + key for the judge.""" + env = {**(self.env_vars or {}), "PROVIDER": "remote", "OS_TYPE": "Ubuntu", "REMOTE_VM_IP": guest_url, + "REMOTE_SERVER_HEADERS": json.dumps(deployed.control_headers)} + config = get_config() + try: + env["LITELLM_API_KEY"] = config.get_litellm_api_key() + except (KeyError, ConfigError): + logger.warning("no model endpoint key configured; the OSWorld controller's judge will have none") + try: + env["LITELLM_BASE_URL"] = config.get_litellm_base_url() + except ConfigError: + pass + return env + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) + if deployed is None: + raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + sandbox_name = f"{SANDBOX_NAME_PREFIX}{self.env_id}" + if any(s.sandbox_name == sandbox_name for s in context.deployed_sandboxes): + raise RuntimeError(f"Sandbox '{sandbox_name}' is already deployed") + guest_url = await guest.resolve_guest_url(deployed) + image = Artifact.get(self.controller_image_artifact_id or DEFAULT_CONTROLLER_IMAGE_ARTIFACT_ID) + if not isinstance(image, DockerImageArtifact): + raise TypeError(f"the controller image {image.id!r} is a {image.type} artifact, not a docker image") + env = self.controller_env(deployed, guest_url) + logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' (image={image.id} v{image.version}, " + f"env={sorted(env)})") + provider = skip_refusing(get_sandbox_provider(), + lambda link: pull_problem(image.pinned_name, on_this_machine=link.ON_THIS_MACHINE), + what=f"sandbox '{sandbox_name}'") + sandbox = await provider.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, + cpu=self.cpu, memory=self.memory_mb, disk_size_gb=10, + timeout=self.ttl_seconds, attribution=deploy_attribution(self, context)) + context.deployed_sandboxes.append(DeployedSandbox( + sandbox_name=sandbox_name, sandbox_id=sandbox.sandbox_id, sandbox_mode=sandbox.mode, + sandbox_type=sandbox.type, + tunnel_urls={str(p): u for p, u in sandbox.tunnel_urls.items()} if sandbox.tunnel_urls else None)) + controller_url = sandbox.tunnel_urls[CONTROLLER_PORT] + await controller.wait_ready(controller_url) + deployed.metadata = {**(deployed.metadata or {}), guest.CONTROLLER_URL_KEY: controller_url, + guest.CONTROLLER_SANDBOX_ID_KEY: sandbox.sandbox_id, guest.GUEST_URL_KEY: guest_url} + logger.info(f"OSWorld controller for env '{self.env_id}' at {controller_url} (sandbox {sandbox.sandbox_id})") + return context diff --git a/src/agent_env/task_step/task_steps/osworld/evaluate.py b/src/agent_env/task_step/task_steps/osworld/evaluate.py new file mode 100644 index 000000000..796d96f2f --- /dev/null +++ b/src/agent_env/task_step/task_steps/osworld/evaluate.py @@ -0,0 +1,222 @@ +"""``osworld_evaluate``: an OSWorld task graded on a computer by the controller ``osworld_deploy_controller`` deployed +beside it, so no grading runs where the agent can reach it. An OSWorld-V2 task is graded by its ``TASK_CLASS.evaluate()`` +on the controller; a classic one by each of its ``evaluator``'s funcs, sent to the controller's ``/evaluate`` one at a +time with its ``result`` and ``expected`` getters and its ``options``, and combined by its ``conj``. The verdict goes +to ``context.metadata["verifications"][verifier_id]`` as ``{"results", "score"}``.""" +from __future__ import annotations + +import logging +import re +import uuid +from typing import Any, ClassVar, Optional + +from agent_env.entity.refs import EntityRef +from agent_env.osworld import controller, guest +from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency +from agent_env.task_step.task_steps.osworld.initialize import run_setup +from agent_env.task_step.task_steps.verifiers.scoring import ScoreAggregator, aggregate_score + +logger = logging.getLogger(__name__) + +PASS_SCORE, FAIL_SCORE = 1.0, 0.0 + +#: What an agent answers a task it finds infeasible with. +INFEASIBLE_TAG = "[INFEASIBLE]" +INFEASIBLE = "infeasible" +#: The LLM-judge funcs, which take a ``model`` option, a task's ``rubric`` and ``parameters``. A deterministic func is +#: called as ``metric(result, expected, **options)``, which a ``model`` it doesn't take fails. +JUDGE_FUNCS = frozenset({"agent_judge", "agent_judge_multi", "agent_judge_docx", "agent_judge_pptx", "agent_judge_pdf"}) +#: The deterministic funcs that also take a task's ``rubric`` and ``parameters``. +CONTEXT_FUNCS = JUDGE_FUNCS | {"compare_pdfs", "compare_csv", "compare_images", "compare_python_pure_text", + "diff_text_file", "website_verifier"} +# An agent that ran out of turns or time leaves a desktop worth grading; every other stop leaves none. Matched on the +# harness's own sentence, since its error types name both kinds. +_BUDGET_EXHAUSTED = re.compile(r"Max turns \(\d+\) reached|agent timed out after \d+", re.I) + + +def budget_exhaustion(prompt_response: PromptResponse | None) -> str | None: + """The stop reason when the agent stopped on its turn or time budget, else None.""" + if prompt_response is None or not prompt_response.error_type: + return None + for text in (prompt_response.error_message, prompt_response.response): + if text and _BUDGET_EXHAUSTED.search(str(text)): + return str(text) + return None + + +class OSWorldEvaluateTaskStep(TaskStep): + type: ClassVar[str] = "osworld_evaluate" + entity_refs = (EntityRef.env("env_id"),) + + def __init__(self, id: str, version: Optional[int], env_id: str, evaluator: Optional[dict] = None, + verifier_id: Optional[str] = None, osworld_v2_task_path: Optional[str] = None, + osworld_v2_task_url: Optional[str] = None, judge_model: Optional[str] = None, + depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True): + """``evaluator`` is OSWorld's: ``funcs``, a name or a list, with ``result``, ``expected`` and ``options`` of + each, in parallel lists for several, ``conj`` (``and`` or ``or``), ``postconfig``, ``rubric`` and + ``instruction``. ``judge_model`` is the ``model`` option the judge funcs get when theirs set none; without it the + controller's own default applies.""" + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.env_id = env_id + self.evaluator = evaluator or {} + self.verifier_id = verifier_id or uuid.uuid4().hex + self.osworld_v2_task_path = osworld_v2_task_path or osworld_v2_task_url + self.judge_model = judge_model + + def to_dict(self) -> dict: + base = {**super().to_dict(), "env_id": self.env_id, "verifier_id": self.verifier_id, "evaluator": self.evaluator} + if self.osworld_v2_task_path is not None: + base["osworld_v2_task_path"] = self.osworld_v2_task_path + if self.judge_model is not None: + base["judge_model"] = self.judge_model + return base + + @classmethod + def from_dict(cls, data: dict) -> OSWorldEvaluateTaskStep: + return cls(**cls._base_from_dict(data), env_id=data["env_id"], evaluator=data.get("evaluator"), + verifier_id=data.get("verifier_id"), osworld_v2_task_path=data.get("osworld_v2_task_path"), + osworld_v2_task_url=data.get("osworld_v2_task_url"), judge_model=data.get("judge_model")) + + def _record(self, context: TaskStepContext, verification: dict) -> TaskStepContext: + context.metadata.setdefault("verifications", {})[self.verifier_id] = verification + logger.info(f"osworld_evaluate '{self.verifier_id}': score={verification['score']}") + return context + + def _skipped(self, context: TaskStepContext, prompt_response: PromptResponse) -> TaskStepContext: + message = f"Skipped: prompt had error_type={prompt_response.error_type}" + logger.warning(f"osworld_evaluate '{self.verifier_id}': {message}") + return self._record(context, {"results": [{"id": "prompt_error", "score": FAIL_SCORE, "result": False, + "message": message}], "score": FAIL_SCORE}) + + async def _stage_result_files(self, controller_url: str, context: TaskStepContext) -> None: + """A run's earlier outputs, ``step_overrides`` ``stage_files`` ``[{url, path}]``, downloaded onto the computer + before anything reads it, to grade a finished run again on a fresh one. Raises on any failure: an unstaged + computer would be graded on the task's own inputs.""" + files = self.step_param_overrides(context).get("stage_files") + if not files: + return + if not isinstance(files, list) or not all( + isinstance(f, dict) and isinstance(f.get("url"), str) and f["url"].startswith(("http://", "https://")) + and f.get("path") for f in files): + raise ValueError(f"osworld_evaluate '{self.verifier_id}': stage_files must be a list of {{url, path}}, " + f"each url http(s) and fetchable from the computer; got {files!r}") + entries = [{"url": f["url"], "path": f["path"]} for f in files] + logger.info(f"osworld_evaluate '{self.verifier_id}': staging {[e['path'] for e in entries]} before grading") + await run_setup(controller_url, [{"type": "download", "parameters": {"files": entries}}], + what=f"osworld_evaluate '{self.verifier_id}' staging result files") + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) + if deployed is None: + raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + url = guest.controller_url(deployed) + await self._stage_result_files(url, context) + prompt_response = context.prompt_responses[-1] if context.prompt_responses else None + if self.osworld_v2_task_path: + if prompt_response is not None and prompt_response.error_type: + return self._skipped(context, prompt_response) + return self._record(context, await self._evaluate_v2(url)) + + if postconfig := self.evaluator.get("postconfig"): + try: + await run_setup(url, postconfig, what=f"osworld_evaluate '{self.verifier_id}' postconfig") + except Exception as e: + logger.warning(f"osworld_evaluate '{self.verifier_id}': postconfig failed, grading anyway: {e}") + budget = budget_exhaustion(prompt_response) + if prompt_response is not None and prompt_response.error_type and budget is None: + return self._skipped(context, prompt_response) + verification = await self._evaluate_funcs(url, prompt_response) + if budget is not None: + # The score is the end state's, earned after the agent ran out of room; the marker beside the verifications + # keeps such runs separable, since score readers project only each verifier's score. + marker = {"reason": "budget_exhausted", "error_type": str(prompt_response.error_type), + "detail": budget[:300]} + verification["partial_credit"] = marker + existing = context.metadata.get("partial_credit") + context.metadata["partial_credit"] = {**(existing if isinstance(existing, dict) else {}), + self.verifier_id: marker} + return self._record(context, verification) + + async def _evaluate_v2(self, controller_url: str) -> dict: + logger.info(f"OSWorld-V2 evaluate '{self.verifier_id}': task_ref={self.osworld_v2_task_path}") + try: + answer = await controller.evaluate( + controller_url, {"osworld_v2_task_url": controller.resolve_task_ref(self.osworld_v2_task_path)}) + except controller.ControllerCallError as e: + logger.warning(f"OSWorld-V2 evaluate '{self.verifier_id}': controller error: {e}") + return {"results": [{"id": "osworld_v2", "score": FAIL_SCORE, "result": False, "message": str(e)[:500]}], + "score": FAIL_SCORE} + score = controller.normalize_score(answer.get("score", FAIL_SCORE)) + return {"results": [{"id": "osworld_v2", "score": score, "result": not answer.get("error", False), + "message": answer.get("message", "")}], "score": score} + + def _funcs(self) -> list[tuple[str, Any, Any, dict]]: + if "funcs" not in self.evaluator: + raise ValueError(f"osworld_evaluate '{self.verifier_id}': the evaluator names no funcs " + f"(its keys are {sorted(self.evaluator)})") + funcs = self.evaluator["funcs"] + if not isinstance(funcs, list): + return [(funcs, self.evaluator.get("result"), self.evaluator.get("expected"), + self.evaluator.get("options") or {})] + n = len(funcs) + return list(zip(funcs, self.evaluator.get("result") or [None] * n, self.evaluator.get("expected") or [None] * n, + [o or {} for o in self.evaluator.get("options") or [{}] * n])) + + def _payload(self, func: str, result: Any, expected: Any, options: dict, instruction: str | None) -> dict: + options = dict(options) + if func in JUDGE_FUNCS: + if self.judge_model: + options.setdefault("model", self.judge_model) + else: + options.pop("model", None) + payload: dict[str, Any] = {"func": func, "options": options} + if result is not None: + payload["result"] = result + if expected is not None: + payload["expected"] = expected + if func in CONTEXT_FUNCS: + payload.update({key: self.evaluator[key] for key in ("parameters", "rubric") if key in self.evaluator}) + if instruction: + payload["instruction"] = instruction + return payload + + async def _evaluate_funcs(self, controller_url: str, prompt_response: PromptResponse | None) -> dict: + answer_text = prompt_response.response if prompt_response is not None else None + instruction = self.evaluator.get("instruction") or (prompt_response.prompt_text if prompt_response else None) + results, scores, metrics = [], [], [] + funcs = self._funcs() + for func, result, expected, options in funcs: + try: + if func == INFEASIBLE: + found = bool(answer_text) and INFEASIBLE_TAG in answer_text + sub = {"id": INFEASIBLE, "score": PASS_SCORE if found else FAIL_SCORE, "result": found, + "message": "Correctly identified infeasible task" if found + else "Did not identify infeasible task"} + elif answer_text and INFEASIBLE_TAG in answer_text: + sub = {"id": "infeasible_on_feasible", "score": FAIL_SCORE, "result": False, + "message": "Agent declared infeasible on a feasible task"} + else: + answer = await controller.evaluate(controller_url, + self._payload(func, result, expected, options, instruction)) + sub = {"id": func, "score": controller.normalize_score(answer.get("score", FAIL_SCORE)), + "result": not answer.get("error", False), "message": answer.get("message", "")} + passed = True + except Exception as e: + logger.warning(f"osworld_evaluate '{self.verifier_id}': {func} failed: {e}") + sub = {"id": func, "score": FAIL_SCORE, "result": False, "message": f"Exception: {str(e)[:300]}"} + passed = False + results.append({"func": func, "results": [sub]}) + scores.append(float(sub["score"])) + metrics.append({"func": func, "passed": passed, "score": float(sub["score"]), "message": sub["message"]}) + conj = self.evaluator.get("conj", "and") + if len(scores) == 1: + score = scores[0] + else: + score = aggregate_score([{"score": s, "result": s == PASS_SCORE} for s in scores], + ScoreAggregator.ANY_PASS if conj == "or" else ScoreAggregator.ALL_PASS) + errored = any(not m["passed"] for m in metrics) + return {"results": results, "score": score, + "debug": {"overall_score": score, "conjunction": conj if len(funcs) > 1 else None, "metrics": metrics, + "error": errored, "message": "Evaluation completed with errors" if errored + else "Evaluation passed" if score == PASS_SCORE else "Evaluation failed"}} diff --git a/src/agent_env/task_step/task_steps/osworld/initialize.py b/src/agent_env/task_step/task_steps/osworld/initialize.py new file mode 100644 index 000000000..aa91f2dda --- /dev/null +++ b/src/agent_env/task_step/task_steps/osworld/initialize.py @@ -0,0 +1,74 @@ +"""``osworld_initialize``: an OSWorld task set up on a computer through the controller ``osworld_deploy_controller`` +deployed beside it. An OSWorld-V2 task (``osworld_v2_task_path``, a ``task_NNN.py`` whose ``TASK_CLASS.setup()`` the +controller runs) resets the controller and initializes the task; a classic one runs its ``init_config``, OSWorld's +``[{type, parameters}]``, through the controller's ``/setup``, which serves every setup type.""" +from __future__ import annotations + +import logging +from typing import ClassVar, Optional + +from agent_env.entity.refs import EntityRef +from agent_env.osworld import controller, guest +from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.task_step import TaskStep, TaskStepDependency + +logger = logging.getLogger(__name__) + + +class OSWorldInitializeTaskStep(TaskStep): + type: ClassVar[str] = "osworld_initialize" + entity_refs = (EntityRef.env("env_id"),) + + def __init__(self, id: str, version: Optional[int], env_id: str, init_config: Optional[list[dict]] = None, + osworld_v2_task_path: Optional[str] = None, osworld_v2_task_url: Optional[str] = None, + depends_on: Optional[list[TaskStepDependency]] = None, fail_task_on_error: bool = True): + """``osworld_v2_task_path`` is an http(s) URL or an object of the object store under its artifacts, which the + step signs for the controller to fetch; ``osworld_v2_task_url`` is its older name, which it wins over.""" + super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) + self.env_id = env_id + self.init_config = init_config or [] + self.osworld_v2_task_path = osworld_v2_task_path or osworld_v2_task_url + + def to_dict(self) -> dict: + base = {**super().to_dict(), "env_id": self.env_id, "init_config": self.init_config} + if self.osworld_v2_task_path is not None: + base["osworld_v2_task_path"] = self.osworld_v2_task_path + return base + + @classmethod + def from_dict(cls, data: dict) -> OSWorldInitializeTaskStep: + return cls(**cls._base_from_dict(data), env_id=data["env_id"], init_config=data.get("init_config"), + osworld_v2_task_path=data.get("osworld_v2_task_path"), + osworld_v2_task_url=data.get("osworld_v2_task_url")) + + async def execute(self, context: TaskStepContext) -> TaskStepContext: + deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) + if deployed is None: + raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + url = guest.controller_url(deployed) + if self.osworld_v2_task_path: + logger.info(f"OSWorld-V2 initialize '{self.id}': task_ref={self.osworld_v2_task_path}") + # /initialize_task refuses to run before a reset. + await controller.post(url, "/reset", {}, expect_json=False) + answer = await controller.post(url, "/initialize_task", { + "id": self.id, "osworld_v2_task_url": controller.resolve_task_ref(self.osworld_v2_task_path)}) + if answer.get("success") is False or answer.get("error"): + raise RuntimeError(f"OSWorld-V2 initialize failed: {answer.get('error') or answer.get('message')}") + return context + if self.init_config: + answer = await run_setup(url, self.init_config, what=f"osworld_initialize '{self.id}'") + if sites := answer.get("results"): + context.metadata.setdefault("osworld_seeded_sites", []).extend(sites) + return context + + +async def run_setup(controller_url: str, config: list[dict], *, what: str) -> dict: + """``config`` through the controller's ``/setup``; raises when it reports failure or a site it seeded failed, since + a seed that answers 200 with every site failed leaves the apps on their own fixtures.""" + answer = await controller.post(controller_url, "/setup", {"config": config}) + if not answer.get("success", True): + raise RuntimeError(f"{what}: setup failed: {answer.get('error') or answer}") + if failed := [site for site in answer.get("results") or [] if isinstance(site, dict) and not site.get("ok", True)]: + detail = "; ".join(f"{site.get('app')}: {site.get('error')}" for site in failed) + raise RuntimeError(f"{what}: setup failed for {len(failed)} site(s): {detail}") + return answer diff --git a/tst/unit/osworld/__init__.py b/tst/unit/osworld/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tst/unit/osworld/app_test.py b/tst/unit/osworld/app_test.py new file mode 100644 index 000000000..a950b197e --- /dev/null +++ b/tst/unit/osworld/app_test.py @@ -0,0 +1,104 @@ +"""The osworld-server app: its card advertises the osworld-guest/v1 extension at the prefix its server forwards to the +guest server under, that forward keeps the method, path, query and body and drops the computer's control token, an +unreachable guest server is a 502, and staging copies the app's folder with the guest server's source.""" +from __future__ import annotations + +import importlib.util +import json +import sys +import tomllib + +import httpx +import pytest +from agentenv_protocol.agent_env_environment import _EXT_ATTR, _EXT_ROUTE_ATTR +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.testclient import TestClient + +from agent_env import osworld +from agent_env.osworld.__main__ import main + + +@pytest.fixture(scope="module") +def app_module(): + spec = importlib.util.spec_from_file_location("osworld_server_app", osworld.APP_DIR / "osworld_server_app.py") + module = importlib.util.module_from_spec(spec) + sys.modules["osworld_server_app"] = module + try: + spec.loader.exec_module(module) + yield module + finally: + sys.modules.pop("osworld_server_app", None) + + +def _forwarding(module, upstream) -> TestClient: + client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + + async def route(request): + return await module.forward(request, client) + + return TestClient(Starlette(routes=[Route(f"{module.PREFIX}/{{path:path}}", route, methods=module.METHODS)])) + + +def test_the_env_toml_names_the_app_and_its_class(app_module): + toml = tomllib.loads((osworld.APP_DIR / "env.toml").read_text()) + assert toml["type"] == "computer_app" and toml["environment_name"] == osworld.APP_NAME + file, _, cls = toml["module"].partition(":") + assert (osworld.APP_DIR / file).is_file() and hasattr(app_module, cls) + + +def test_the_card_advertises_the_guest_extension_at_the_forward_prefix(app_module): + advertised = getattr(app_module.OSWorldServer.guest, _EXT_ATTR) + assert advertised.uri == osworld.GUEST_EXTENSION == app_module.GUEST_EXTENSION + assert advertised.params["endpoint"] == osworld.PROXY_PREFIX == app_module.PREFIX + assert getattr(app_module.OSWorldServer.guest, _EXT_ROUTE_ATTR) == (osworld.PROXY_PREFIX, "GET") + + +@pytest.mark.asyncio +async def test_the_guest_extension_says_whether_the_guest_server_is_up(app_module, monkeypatch): + def refused(*args, **kwargs): + raise OSError("refused") + + monkeypatch.setattr(app_module.socket, "create_connection", refused) + answer = await app_module.OSWorldServer.guest(None) + assert answer["ready"] is False and "refused" in answer["detail"] + assert app_module.OSWorldServer.guest_server(None).ok is False + + +def test_the_forward_keeps_the_request_and_drops_the_control_token(app_module): + seen = [] + + def upstream(request): + seen.append(request) + return httpx.Response(201, json={"ok": True}, headers={"x-guest": "1"}) + + response = _forwarding(app_module, upstream).post( + "/osworld/setup/launch?a=1&a=2", content=json.dumps({"command": ["gedit"]}), + headers={"Authorization": "Bearer control", "Content-Type": "application/json"}) + assert response.status_code == 201 and response.json() == {"ok": True} and response.headers["x-guest"] == "1" + [request] = seen + assert str(request.url) == "http://127.0.0.1:5000/setup/launch?a=1&a=2" + assert request.method == "POST" and json.loads(request.content) == {"command": ["gedit"]} + assert "authorization" not in request.headers + + +def test_an_unreachable_guest_server_is_a_502(app_module): + def upstream(request): + raise httpx.ConnectError("refused") + + response = _forwarding(app_module, upstream).get("/osworld/platform") + assert response.status_code == 502 and "unreachable" in response.text + + +def test_staging_copies_the_app_with_the_guest_servers_source(tmp_path): + source = tmp_path / "src" + source.mkdir() + for name in osworld.PAYLOAD_FILES: + (source / name).write_text(f"# {name}") + out = tmp_path / "out" + assert main(["stage", "--source", str(source), "--out", str(out)]) == 0 + assert (out / "env.toml").is_file() and (out / "start-osworld-server.sh").is_file() + assert (out / "payload" / "main.py").read_text() == "# main.py" + assert main(["stage", "--source", str(source), "--out", str(out)]) == 1 + with pytest.raises(FileNotFoundError, match="pyxcursor.py"): + (source / "pyxcursor.py").unlink() or osworld.stage_server(source, tmp_path / "again") diff --git a/tst/unit/osworld/controller_test.py b/tst/unit/osworld/controller_test.py new file mode 100644 index 000000000..2e1a54476 --- /dev/null +++ b/tst/unit/osworld/controller_test.py @@ -0,0 +1,112 @@ +"""The controller over HTTP: a call's answer, an HTTP error as ControllerCallError with its status, /evaluate started +and polled, the blocking /evaluate on a controller without the async routes, a score as a float, and a task ref signed +only when it's an object under the object store's artifacts.""" +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import httpx +import pytest + +from agent_env.osworld import controller +from tst.unit.osworld.helpers import CONTROLLER, Controller, route_httpx + + +@pytest.fixture +def ctl(monkeypatch): + stand_in = Controller() + route_httpx(monkeypatch, controller, stand_in) + monkeypatch.setattr(controller.asyncio, "sleep", AsyncMock()) + return stand_in + + +@pytest.mark.asyncio +async def test_a_call_returns_its_answer_and_an_http_error_raises_with_its_status(ctl): + ctl.answers["/setup"] = {"success": True} + assert await controller.post(CONTROLLER, "/setup", {"config": []}) == {"success": True} + ctl.answers["/setup"] = httpx.Response(500, text="boom") + with pytest.raises(controller.ControllerCallError) as raised: + await controller.post(CONTROLLER, "/setup", {}) + assert raised.value.status == 500 and raised.value.body == "boom" + + +@pytest.mark.asyncio +async def test_a_bare_answer_is_empty_when_json_isnt_expected(ctl): + ctl.answers["/reset"] = httpx.Response(200, text="ok") + assert await controller.post(CONTROLLER, "/reset", {}, expect_json=False) == {} + with pytest.raises(controller.ControllerCallError, match="non-JSON"): + await controller.post(CONTROLLER, "/reset", {}) + + +@pytest.mark.asyncio +async def test_a_transport_failure_raises(ctl): + ctl.answers["/setup"] = httpx.ConnectError("refused") + with pytest.raises(controller.ControllerCallError, match="transport failure"): + await controller.post(CONTROLLER, "/setup", {}) + + +@pytest.mark.asyncio +async def test_evaluate_is_started_and_polled_until_done(ctl): + ctl.answers["/evaluate_async"] = {"job_id": "j1"} + ctl.answers["/evaluate_result"] = [httpx.ConnectError("blip"), {"status": "running"}, + httpx.Response(200, text=""), {"status": "done", "result": {"score": 1}}] + assert await controller.evaluate(CONTROLLER, {"func": "compare_table"}) == {"score": 1} + assert ctl.bodies("/evaluate_async") == [{"func": "compare_table"}] + assert {r.url.params["job_id"] for r in ctl.seen if r.url.path == "/evaluate_result"} == {"j1"} + + +@pytest.mark.asyncio +async def test_evaluate_falls_back_to_the_blocking_route(ctl): + ctl.answers["/evaluate_async"] = httpx.Response(404) + ctl.answers["/evaluate"] = {"score": 0.5} + assert await controller.evaluate(CONTROLLER, {"func": "x"}) == {"score": 0.5} + assert ctl.paths() == ["/evaluate_async", "/evaluate"] + + +@pytest.mark.asyncio +async def test_a_failed_job_raises(ctl): + ctl.answers["/evaluate_async"] = {"job_id": "j1"} + ctl.answers["/evaluate_result"] = {"status": "error", "message": "judge died"} + with pytest.raises(controller.ControllerCallError, match="judge died"): + await controller.evaluate(CONTROLLER, {}) + + +def test_scores_normalize_to_floats(): + assert [controller.normalize_score(s) for s in (True, False, 1, 0.25, {"a": 1, "b": 0, "c": "x"}, {}, None)] == \ + [1.0, 0.0, 1.0, 0.25, 0.5, 0.0, 0.0] + + +class _Store: + def __init__(self, signs=True): + self.signs = signs + + def owns(self, url): + return url.startswith("s3://bucket/") + + def get_object_key(self, url): + return url.removeprefix("s3://bucket/") + + def signed_get_url(self, url, expires_in): + return f"https://signed.example/{self.get_object_key(url)}?ttl={expires_in}" if self.signs else None + + +def _config(monkeypatch, store): + monkeypatch.setattr(controller, "get_config", lambda: SimpleNamespace( + get_object_store_at=lambda url: store, get_artifact_key_prefix=lambda: "pre/")) + + +def test_a_task_ref_is_signed_only_under_the_artifacts(monkeypatch): + _config(monkeypatch, _Store()) + assert controller.resolve_task_ref("https://host/task.py") == "https://host/task.py" + assert controller.resolve_task_ref("s3://bucket/pre/artifacts/t.py") == \ + f"https://signed.example/pre/artifacts/t.py?ttl={controller.TASK_REF_TTL_S}" + for outside in ("s3://bucket/pre/other/t.py", "s3://elsewhere/pre/artifacts/t.py"): + with pytest.raises(ValueError, match="refusing to sign"): + controller.resolve_task_ref(outside) + + +def test_a_store_that_cant_sign_is_refused(monkeypatch): + _config(monkeypatch, _Store(signs=False)) + with pytest.raises(ValueError, match="can't sign"): + controller.resolve_task_ref("s3://bucket/pre/artifacts/t.py") diff --git a/tst/unit/osworld/guest_test.py b/tst/unit/osworld/guest_test.py new file mode 100644 index 000000000..61aa1a0ac --- /dev/null +++ b/tst/unit/osworld/guest_test.py @@ -0,0 +1,58 @@ +"""The guest server's address comes from the computer's card: the endpoint of the osworld-guest/v1 extension one of its +apps advertises, resolved against the computer's address, from the stored card or, while the app's server isn't +answering yet, the card the computer serves now; a computer without the app is refused.""" +from __future__ import annotations + +from unittest.mock import AsyncMock + +import httpx +import pytest + +from agent_env.osworld import guest +from tst.unit.osworld.helpers import ENV_URL, card, computer + + +def test_the_endpoint_is_the_one_an_app_advertises_for_the_guest_extension(): + assert guest.guest_endpoint(card("/apps/7/osworld")) == "/apps/7/osworld" + assert guest.guest_endpoint(card()["children_environments"][0]) == "/svc/osworld-server/osworld" + assert guest.guest_endpoint(card(uri="urn:agentenv:other/v1")) is None + assert guest.guest_endpoint(None) is None + + +@pytest.mark.asyncio +async def test_the_stored_card_is_read_first(monkeypatch): + live = AsyncMock() + monkeypatch.setattr(guest.protocol_v1, "get_card", live) + assert await guest.resolve_guest_url(computer(card("/apps/7/osworld/"))) == f"{ENV_URL}/apps/7/osworld" + live.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_the_live_card_is_read_again_until_the_app_answers(monkeypatch): + bare = {"children_environments": [{"name": "osworld-server", "url": "/svc/osworld-server/agentenv"}]} + live = AsyncMock(side_effect=[httpx.ConnectError("down"), bare, card("https://elsewhere.example/g/")]) + monkeypatch.setattr(guest.protocol_v1, "get_card", live) + monkeypatch.setattr(guest.asyncio, "sleep", AsyncMock()) + assert await guest.resolve_guest_url(computer(bare), wait_s=60) == "https://elsewhere.example/g" + assert live.await_count == 3 + + +@pytest.mark.asyncio +async def test_a_computer_without_the_app_is_refused(monkeypatch): + monkeypatch.setattr(guest.protocol_v1, "get_card", AsyncMock(return_value={})) + with pytest.raises(RuntimeError, match="compose the osworld-server computer app"): + await guest.resolve_guest_url(computer({}), wait_s=0) + + +@pytest.mark.asyncio +async def test_an_env_with_no_address_is_refused(): + env = computer() + env.environment_card_url = None + with pytest.raises(RuntimeError, match="no address"): + await guest.resolve_guest_url(env) + + +def test_the_controller_url_is_the_one_recorded(): + assert guest.controller_url(computer(osworld_controller_url="https://c")) == "https://c" + with pytest.raises(RuntimeError, match="run osworld_deploy_controller"): + guest.controller_url(computer()) diff --git a/tst/unit/osworld/helpers.py b/tst/unit/osworld/helpers.py new file mode 100644 index 000000000..269d9f576 --- /dev/null +++ b/tst/unit/osworld/helpers.py @@ -0,0 +1,62 @@ +"""A computer env as the ``osworld_*`` steps see one, and a stand-in for the OSWorld controller's routes.""" +from __future__ import annotations + +import json +from dataclasses import dataclass, field + +import httpx +from agentenv_protocol.types import WELL_KNOWN_PATH + +from agent_env.env.env import DeployedEnv +from agent_env.osworld import GUEST_EXTENSION + +ENV_URL = "https://computer.example" +TOKEN = "s3cret-control-token" +CONTROLLER = "https://controller.example" + + +def card(endpoint: str = "/svc/osworld-server/osworld", *, uri: str = GUEST_EXTENSION) -> dict: + """A computer's card with one app, which advertises ``uri`` at ``endpoint``.""" + app = {"name": "osworld-server", "url": "/svc/osworld-server/agentenv", + "capabilities": {"extensions": [{"uri": uri, "params": {"endpoint": endpoint, + "methods": {"guest": {"method": "GET"}}}}]}} + return {"name": "desk", "capabilities": {"extensions": []}, "children_environments": [app]} + + +def computer(card_: dict | None = None, **metadata) -> DeployedEnv: + return DeployedEnv(env_id="desk", env_version=1, environment_card_url=f"{ENV_URL}{WELL_KNOWN_PATH}", + environment_card=card() if card_ is None else card_, control_token=TOKEN, + metadata=dict(metadata)) + + +@dataclass +class Controller: + """The controller's routes, answered from ``answers`` by path (a list is answered in turn), and the requests made.""" + + answers: dict[str, object] = field(default_factory=dict) + seen: list[httpx.Request] = field(default_factory=list) + + def handler(self, request: httpx.Request) -> httpx.Response: + self.seen.append(request) + answer = self.answers.get(request.url.path, httpx.Response(404, text="no route")) + if isinstance(answer, list): + answer = answer.pop(0) + if isinstance(answer, Exception): + raise answer + return answer if isinstance(answer, httpx.Response) else httpx.Response(200, json=answer) + + def bodies(self, path: str) -> list[dict]: + return [json.loads(r.content) for r in self.seen if r.url.path == path] + + def paths(self) -> list[str]: + return [r.url.path for r in self.seen] + + +def route_httpx(monkeypatch, module, controller: Controller) -> None: + """Send ``module``'s httpx clients to ``controller``.""" + real = httpx.AsyncClient + + def client(*args, **kwargs): + return real(*args, transport=httpx.MockTransport(controller.handler), **kwargs) + + monkeypatch.setattr(module.httpx, "AsyncClient", client) diff --git a/tst/unit/osworld/steps_test.py b/tst/unit/osworld/steps_test.py new file mode 100644 index 000000000..717564bf6 --- /dev/null +++ b/tst/unit/osworld/steps_test.py @@ -0,0 +1,339 @@ +"""The osworld_* steps on a computer env: osworld_deploy_controller deploys the controller beside it, in remote mode +against the guest server its card advertises with the env's control token, and records both addresses; it refuses a +computer without the app before it creates a sandbox, and logs no secret. osworld_initialize sets an OSWorld-V2 task up +with a reset and /initialize_task, and a classic one through /setup. osworld_evaluate grades on the controller: an +OSWorld-V2 task by /evaluate, a classic one func by func combined by its conj, a controller error as a score of 0, a +turn that errored as a skip unless the agent ran out of its budget, result files staged first on a re-grade.""" +from __future__ import annotations + +import json +import logging +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from agent_env.artifact import DockerImageArtifact +from agent_env.osworld import controller, guest +from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.registry import get_task_step_registry +from agent_env.task_step.task_steps.osworld import deploy_controller as deploy_module +from agent_env.task_step.task_steps.osworld.deploy_controller import OSWorldDeployControllerTaskStep +from agent_env.task_step.task_steps.osworld.evaluate import OSWorldEvaluateTaskStep +from agent_env.task_step.task_steps.osworld.initialize import OSWorldInitializeTaskStep +from tst.unit.osworld.helpers import CONTROLLER, ENV_URL, TOKEN, Controller, card, computer, route_httpx + +LITELLM = "sk-model-key" + + +def _context(env, *prompts: PromptResponse) -> TaskStepContext: + context = TaskStepContext() + context.deployed_envs.append(env) + context.prompt_responses.extend(prompts) + return context + + +def _ready(): + return computer(osworld_controller_url=CONTROLLER, osworld_guest_url=f"{ENV_URL}/svc/osworld-server/osworld") + + +@pytest.fixture +def ctl(monkeypatch): + stand_in = Controller() + route_httpx(monkeypatch, controller, stand_in) + monkeypatch.setattr(controller.asyncio, "sleep", AsyncMock()) + return stand_in + + +@pytest.mark.parametrize("data", [ + {"type": "osworld_deploy_controller", "env_id": "desk", "controller_image_artifact_id": "ctl", "ttl_seconds": 60, + "cpu": 1.0, "memory_mb": 2048, "env_vars": {"JUDGE_MODEL": "m"}}, + {"type": "osworld_initialize", "env_id": "desk", "init_config": [{"type": "sleep", "parameters": {"seconds": 1}}], + "osworld_v2_task_path": "https://host/t.py"}, + {"type": "osworld_evaluate", "env_id": "desk", "evaluator": {"funcs": "compare_table"}, "verifier_id": "v", + "osworld_v2_task_path": "https://host/t.py", "judge_model": "m"}, +], ids=lambda data: data["type"]) +def test_the_steps_are_registered_and_round_trip(data): + step = get_task_step_registry()[data["type"]].from_dict({"id": "s", "version": 1, **data}) + assert {k: v for k, v in step.to_dict().items() if k in data} == data + + +def test_the_legacy_task_url_field_is_read_as_the_path(): + step = OSWorldInitializeTaskStep.from_dict({"id": "s", "env_id": "desk", "osworld_v2_task_url": "https://h/t.py"}) + assert step.osworld_v2_task_path == "https://h/t.py" + + +class _Provider: + ON_THIS_MACHINE = False + + def __init__(self): + self.links = [self] + self.create_container = AsyncMock(return_value=SimpleNamespace( + sandbox_id="sb-1", mode="container", type="docker", tunnel_urls={8000: CONTROLLER})) + + +@pytest.fixture +def deploy(monkeypatch): + provider = _Provider() + image = DockerImageArtifact(id="cua-controller", version=3, description="controller", + image_name="ghcr.io/example/controller:1") + config = MagicMock() + config.get_litellm_api_key.return_value = LITELLM + config.get_litellm_base_url.return_value = "https://model.example" + ready = AsyncMock() + get = MagicMock(return_value=image) + monkeypatch.setattr(deploy_module.Artifact, "get", get) + monkeypatch.setattr(deploy_module, "get_config", lambda: config) + monkeypatch.setattr(deploy_module, "get_sandbox_provider", lambda: provider) + monkeypatch.setattr(deploy_module.controller, "wait_ready", ready) + return SimpleNamespace(provider=provider, ready=ready, get=get) + + +@pytest.mark.asyncio +async def test_deploy_runs_the_controller_against_the_card_endpoint_and_records_it(deploy, caplog): + env = computer(card("/apps/7/osworld"), existing="kept") + context = _context(env) + step = OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk", env_vars={"JUDGE_MODEL": "m", + "PROVIDER": "docker"}) + with caplog.at_level(logging.DEBUG): + await step.execute(context) + + deploy.get.assert_called_once_with("cua-controller") + kwargs = deploy.provider.create_container.call_args.kwargs + assert kwargs["image_name"] == "ghcr.io/example/controller:1" and kwargs["port"] == 8000 + assert kwargs["env"] == {"JUDGE_MODEL": "m", "PROVIDER": "remote", "OS_TYPE": "Ubuntu", + "REMOTE_VM_IP": f"{ENV_URL}/apps/7/osworld", + "REMOTE_SERVER_HEADERS": json.dumps({"Authorization": f"Bearer {TOKEN}"}), + "LITELLM_API_KEY": LITELLM, "LITELLM_BASE_URL": "https://model.example"} + deploy.ready.assert_awaited_once_with(CONTROLLER) + assert env.metadata == {"existing": "kept", guest.CONTROLLER_URL_KEY: CONTROLLER, + guest.CONTROLLER_SANDBOX_ID_KEY: "sb-1", guest.GUEST_URL_KEY: f"{ENV_URL}/apps/7/osworld"} + [sandbox] = context.deployed_sandboxes + assert (sandbox.sandbox_name, sandbox.sandbox_id, sandbox.tunnel_urls) == \ + ("osworld-controller-desk", "sb-1", {"8000": CONTROLLER}) + assert TOKEN not in caplog.text and LITELLM not in caplog.text + + +@pytest.mark.asyncio +async def test_deploy_refuses_a_computer_without_the_app_before_creating_a_sandbox(deploy, monkeypatch): + monkeypatch.setattr(guest, "CARD_WAIT_S", 0) + monkeypatch.setattr(guest.protocol_v1, "get_card", AsyncMock(return_value={})) + env = computer({}) + with pytest.raises(RuntimeError, match="osworld-server computer app"): + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(_context(env)) + deploy.provider.create_container.assert_not_called() + assert env.metadata == {} + + +@pytest.mark.asyncio +async def test_deploy_twice_for_one_env_is_refused(deploy): + context = _context(computer()) + await OSWorldDeployControllerTaskStep(id="a", version=1, env_id="desk").execute(context) + with pytest.raises(RuntimeError, match="already deployed"): + await OSWorldDeployControllerTaskStep(id="b", version=1, env_id="desk").execute(context) + + +@pytest.mark.asyncio +async def test_initialize_v2_resets_then_initializes_the_task(ctl): + ctl.answers.update({"/reset": httpx.Response(200, text="ok"), "/initialize_task": {"success": True}}) + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", + osworld_v2_task_path="https://host/t.py").execute(_context(_ready())) + assert ctl.paths() == ["/reset", "/initialize_task"] + assert ctl.bodies("/initialize_task") == [{"id": "init", "osworld_v2_task_url": "https://host/t.py"}] + + +@pytest.mark.asyncio +async def test_initialize_v2_fails_on_a_failed_setup(ctl): + ctl.answers.update({"/reset": {}, "/initialize_task": {"success": False, "error": "no TASK_CLASS"}}) + with pytest.raises(RuntimeError, match="no TASK_CLASS"): + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", + osworld_v2_task_path="https://host/t.py").execute(_context(_ready())) + + +@pytest.mark.asyncio +async def test_initialize_classic_sends_the_whole_config_to_setup(ctl): + config = [{"type": "download", "parameters": {"files": [{"url": "https://h/a", "path": "/home/user/a"}]}}, + {"type": "chrome_open_tabs", "parameters": {"urls_to_open": ["https://h/"]}}, {"type": "launch", "parameters": {"command": ["gedit"]}}] + ctl.answers["/setup"] = {"success": True, "results": [{"app": "mail", "ok": True}]} + context = _context(_ready()) + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", init_config=config).execute(context) + assert ctl.bodies("/setup") == [{"config": config}] + assert context.metadata["osworld_seeded_sites"] == [{"app": "mail", "ok": True}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("answer, match", [({"success": False, "error": "no such type"}, "no such type"), + ({"success": True, "results": [{"app": "mail", "ok": False, + "error": "seed 500"}]}, "seed 500")]) +async def test_initialize_classic_fails_on_a_failed_setup_or_site(ctl, answer, match): + ctl.answers["/setup"] = answer + with pytest.raises(RuntimeError, match=match): + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", + init_config=[{"type": "sleep"}]).execute(_context(_ready())) + + +@pytest.mark.asyncio +async def test_a_step_needs_the_controller_deployed_first(ctl): + with pytest.raises(RuntimeError, match="run osworld_deploy_controller"): + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", + init_config=[{"type": "sleep"}]).execute(_context(computer())) + + +def _prompt(**fields) -> PromptResponse: + return PromptResponse(**{"prompt_id": "p", "response": "done", "prompt_text": "Save the sheet", **fields}) + + +@pytest.mark.asyncio +async def test_evaluate_v2_records_the_controllers_score(ctl): + ctl.answers.update({"/evaluate_async": {"job_id": "j"}, + "/evaluate_result": {"status": "done", "result": {"score": True, "message": "ok"}}}) + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + osworld_v2_task_path="https://host/t.py").execute(context) + assert ctl.bodies("/evaluate_async") == [{"osworld_v2_task_url": "https://host/t.py"}] + assert context.metadata["verifications"]["v"] == { + "results": [{"id": "osworld_v2", "score": 1.0, "result": True, "message": "ok"}], "score": 1.0} + + +@pytest.mark.asyncio +async def test_evaluate_v2_scores_a_controller_error_0(ctl): + ctl.answers.update({"/evaluate_async": httpx.Response(404), "/evaluate": httpx.Response(500, text="boom")}) + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + osworld_v2_task_path="https://host/t.py").execute(context) + [result] = context.metadata["verifications"]["v"]["results"] + assert context.metadata["verifications"]["v"]["score"] == 0.0 + assert result["id"] == "osworld_v2" and result["result"] is False and "HTTP 500" in result["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("v2", [True, False], ids=["v2", "classic"]) +async def test_evaluate_skips_a_turn_that_errored(ctl, v2): + fields = {"osworld_v2_task_path": "https://host/t.py"} if v2 else {"evaluator": {"funcs": "compare_table"}} + context = _context(_ready(), _prompt(error_type="INFRA_ERROR", error_message="sandbox gone")) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", **fields).execute(context) + assert context.metadata["verifications"]["v"] == { + "results": [{"id": "prompt_error", "score": 0.0, "result": False, + "message": "Skipped: prompt had error_type=INFRA_ERROR"}], "score": 0.0} + assert "/evaluate_async" not in ctl.paths() + + +def _scores(ctl, *scores): + ctl.answers["/evaluate_async"] = httpx.Response(404) + ctl.answers["/evaluate"] = [httpx.Response(200, json={"score": s, "message": f"got {s}"}) for s in scores] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conj, scores, expected", [("and", (1, 1), 1.0), ("and", (1, 0), 0.0), ("or", (0, 1), 1.0), + ("or", (0, 0), 0.0)]) +async def test_evaluate_classic_sends_each_func_and_combines_by_conj(ctl, conj, scores, expected): + _scores(ctl, *scores) + evaluator = {"funcs": ["compare_table", "agent_judge"], "conj": conj, + "result": [{"type": "vm_file", "path": "/home/user/a.xlsx"}, {"type": "vm_file", "path": "/b"}], + "expected": [{"type": "cloud_file", "path": "https://h/gold.xlsx"}, None], + "options": [{"model": "copied-by-mistake"}, {}], "rubric": [{"criterion": "x"}]} + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", evaluator=evaluator, + judge_model="judge-m").execute(context) + table, judge = ctl.bodies("/evaluate") + assert table == {"func": "compare_table", "options": {}, "result": evaluator["result"][0], + "expected": evaluator["expected"][0], "instruction": "Save the sheet"} + assert judge == {"func": "agent_judge", "options": {"model": "judge-m"}, "result": evaluator["result"][1], + "rubric": [{"criterion": "x"}], "instruction": "Save the sheet"} + verification = context.metadata["verifications"]["v"] + assert verification["score"] == expected + assert [r["func"] for r in verification["results"]] == ["compare_table", "agent_judge"] + assert verification["results"][0]["results"] == [ + {"id": "compare_table", "score": float(scores[0]), "result": True, "message": f"got {scores[0]}"}] + + +@pytest.mark.asyncio +async def test_evaluate_classic_one_func_keeps_its_partial_score(ctl): + _scores(ctl, 0.4) + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": "agent_judge"}).execute(context) + assert context.metadata["verifications"]["v"]["score"] == 0.4 + + +@pytest.mark.asyncio +async def test_evaluate_classic_scores_a_failing_func_0_and_grades_the_rest(ctl): + ctl.answers["/evaluate_async"] = httpx.Response(404) + ctl.answers["/evaluate"] = [httpx.Response(500, text="getter failed"), httpx.Response(200, json={"score": 1})] + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": ["compare_table", "file_contains"], "conj": "or"}).execute(context) + verification = context.metadata["verifications"]["v"] + assert verification["score"] == 1.0 and verification["debug"]["error"] is True + assert "getter failed" in verification["results"][0]["results"][0]["message"] + + +@pytest.mark.asyncio +async def test_evaluate_infeasible_is_graded_on_the_agents_answer(ctl): + context = _context(_ready(), _prompt(response="[INFEASIBLE] no such app")) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": ["infeasible"]}).execute(context) + assert context.metadata["verifications"]["v"]["score"] == 1.0 + await OSWorldEvaluateTaskStep(id="e2", version=1, env_id="desk", verifier_id="w", + evaluator={"funcs": "compare_table"}).execute(context) + [[result]] = [r["results"] for r in context.metadata["verifications"]["w"]["results"]] + assert result["id"] == "infeasible_on_feasible" and context.metadata["verifications"]["w"]["score"] == 0.0 + assert ctl.paths() == [] + + +@pytest.mark.asyncio +async def test_evaluate_grades_a_run_that_ran_out_of_turns_and_marks_it(ctl): + _scores(ctl, 1) + context = _context(_ready(), _prompt(error_type="AGENT_ERROR", + error_message="Max turns (50) reached without completion.")) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": "compare_table"}).execute(context) + verification = context.metadata["verifications"]["v"] + assert verification["score"] == 1.0 + assert verification["partial_credit"]["reason"] == "budget_exhausted" + assert context.metadata["partial_credit"] == {"v": verification["partial_credit"]} + + +@pytest.mark.asyncio +async def test_evaluate_runs_postconfig_once_and_grades_when_it_fails(ctl): + _scores(ctl, 1) + ctl.answers["/setup"] = {"success": False, "error": "window gone"} + postconfig = [{"type": "activate_window", "parameters": {"window_name": "x"}}] + context = _context(_ready(), _prompt()) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": "compare_table", "postconfig": postconfig}).execute(context) + assert ctl.bodies("/setup") == [{"config": postconfig}] + assert "postconfig" not in ctl.bodies("/evaluate")[0] + assert context.metadata["verifications"]["v"]["score"] == 1.0 + + +def _regrade(context, files): + context.metadata["user_overrides"] = {"step_params": {"e": {"stage_files": files}}} + return context + + +@pytest.mark.asyncio +async def test_evaluate_stages_result_files_before_grading(ctl): + _scores(ctl, 1) + ctl.answers["/setup"] = {"success": True} + files = [{"url": "https://signed.example/out.xlsx", "path": "/home/user/out.xlsx"}] + context = _regrade(_context(_ready(), _prompt()), files) + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": "compare_table"}).execute(context) + assert ctl.paths()[0] == "/setup" + assert ctl.bodies("/setup") == [{"config": [{"type": "download", "parameters": {"files": files}}]}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("files, answer", [([{"url": "s3://b/out.xlsx", "path": "/x"}], None), + ([{"url": "https://h/out", "path": "/x"}], {"success": False})], + ids=["not-http", "download-failed"]) +async def test_evaluate_refuses_to_grade_an_unstaged_computer(ctl, files, answer): + if answer is not None: + ctl.answers["/setup"] = answer + context = _regrade(_context(_ready(), _prompt()), files) + with pytest.raises((ValueError, RuntimeError)): + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator={"funcs": "compare_table"}).execute(context) + assert "/evaluate_async" not in ctl.paths() From 460b019093fa968764707f0f104626d42e10023e Mon Sep 17 00:00:00 2001 From: William Loo Date: Wed, 11 Nov 2026 23:35:28 -0800 Subject: [PATCH 2/5] fix(osworld): run end to end on a local computer env, with an integration test (FD-5492) A computer env with the osworld-server app, the controller beside it, a classic task set up and graded through it: pass scores 1.0, fail 0.0. Running it found: - The controller got the guest server's address as this process reaches it; a local container reaches this machine's loopback by another name. It now gets it, and the model endpoint, through url_from_sandbox, as an agent's deploy does. - supervisord looks a bare command up on PATH, so the guest server's service never started: it runs as `bash start-osworld-server.sh`. - python-xlib refuses to connect without an Xauthority file; the start script makes an empty one. - pyautogui's screenshots need scrot on X11; the app installs it. - A deployed env's record is stored once, when it is deployed, so the controller's address written on it lasted only the process. It is recorded in the run's context metadata too, and read from there. tst/integration/computer/osworld_test.py stages the app from OSWORLD_SERVER_SOURCE and runs the controller image OSWORLD_CONTROLLER_IMAGE (one with scale-cua #360); it skips without them. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apps/osworld-server/osworld_server_app.py | 9 ++-- .../osworld-server/start-osworld-server.sh | 2 + src/agent_env/osworld/guest.py | 11 ++-- .../task_steps/osworld/deploy_controller.py | 16 ++++-- .../task_step/task_steps/osworld/evaluate.py | 2 +- .../task_steps/osworld/initialize.py | 2 +- .../osworld/envs/osworld-desk/env.toml | 4 ++ .../computer/osworld/tasks/classic.json | 18 +++++++ tst/integration/computer/osworld_test.py | 54 +++++++++++++++++++ tst/unit/osworld/steps_test.py | 29 ++++++++++ 10 files changed, 133 insertions(+), 14 deletions(-) create mode 100644 tst/integration/computer/osworld/envs/osworld-desk/env.toml create mode 100644 tst/integration/computer/osworld/tasks/classic.json create mode 100644 tst/integration/computer/osworld_test.py diff --git a/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py b/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py index 49a04df7d..92cd2fcc2 100644 --- a/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py +++ b/src/agent_env/osworld/apps/osworld-server/osworld_server_app.py @@ -29,9 +29,10 @@ _DROPPED_REQUEST = frozenset({"host", "content-length", "connection", "transfer-encoding", "authorization"}) _DROPPED_RESPONSE = frozenset({"content-length", "connection", "transfer-encoding", "content-encoding"}) #: What a CUA VM installs for the guest server that Ubuntu 24.04's desktop computer lacks: AT-SPI's Python bindings -#: for its accessibility tree, Tk for pyautogui, the window and input tools its setup routes run, ffmpeg for recordings. -PACKAGES = ["python3-pyatspi", "python3-gi", "gir1.2-atspi-2.0", "at-spi2-core", "python3-tk", "ffmpeg", "wmctrl", - "xdotool", "x11-xkb-utils", "x11-xserver-utils"] +#: for its accessibility tree, Tk for pyautogui, scrot for pyautogui's screenshots on X11, the window and input tools its +#: setup routes run, ffmpeg for recordings. +PACKAGES = ["python3-pyatspi", "python3-gi", "gir1.2-atspi-2.0", "at-spi2-core", "python3-tk", "scrot", "ffmpeg", + "wmctrl", "xdotool", "x11-xkb-utils", "x11-xserver-utils"] async def forward(request: Request, client: httpx.AsyncClient, upstream: str = UPSTREAM) -> Response: @@ -58,7 +59,7 @@ def guest_server_check() -> ReadyCheck: @computer_app(supports=[{"os": {"type": "linux", "display": ["x11"]}, "arch": ["amd64", "arm64"]}]) @linux(packages=Packages(apt=PACKAGES), python="requirements.txt", - services=[Service(name="osworld-server", command="start-osworld-server.sh")]) + services=[Service(name="osworld-server", command="bash start-osworld-server.sh")]) class OSWorldServer(ComputerAppEnvironment): """The OSWorld guest server, forwarded under ``/osworld/`` to the controller beside the computer.""" diff --git a/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh b/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh index dcdc4ae10..0fb5765c6 100755 --- a/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh +++ b/src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh @@ -4,6 +4,8 @@ cd "$(dirname "$0")/payload" export DISPLAY="${DISPLAY:-:0}" export PYTHONUNBUFFERED=1 +# python-xlib refuses to connect without an Xauthority file; the computer's X server takes connections without one. +[ -e "${XAUTHORITY:-$HOME/.Xauthority}" ] || touch "${XAUTHORITY:-$HOME/.Xauthority}" while true; do ../venv/bin/python -c 'import runpy, sys; sys.path.append("/usr/lib/python3/dist-packages"); sys.argv = ["main.py"]; runpy.run_path("main.py", run_name="__main__")' echo "osworld-server exited ($?); restarting in 5s" >&2 diff --git a/src/agent_env/osworld/guest.py b/src/agent_env/osworld/guest.py index 6cb37e09b..aec58bc6b 100644 --- a/src/agent_env/osworld/guest.py +++ b/src/agent_env/osworld/guest.py @@ -19,6 +19,9 @@ CONTROLLER_URL_KEY = "osworld_controller_url" CONTROLLER_SANDBOX_ID_KEY = "osworld_controller_sandbox_id" GUEST_URL_KEY = "osworld_guest_url" +#: ``context.metadata[CONTROLLERS_KEY][env_id]``: the same three, where the run's stored context keeps them; a deployed +#: env's record is stored once, when it is deployed, so what a later step writes on it lasts only the process. +CONTROLLERS_KEY = "osworld_controllers" #: A computer's card lists an app without its extensions until the app's server answers: how long resolving waits for #: the guest extension, and how often it reads the card again. CARD_WAIT_S = 120.0 @@ -66,8 +69,10 @@ async def resolve_guest_url(deployed: Any, *, wait_s: float | None = None, poll_ return _joined(environment_url, endpoint) -def controller_url(deployed: Any) -> str: - """The controller ``osworld_deploy_controller`` recorded on ``deployed``.""" - if url := (deployed.metadata or {}).get(CONTROLLER_URL_KEY): +def controller_url(deployed: Any, context: Any = None) -> str: + """The controller ``osworld_deploy_controller`` recorded for ``deployed``: in ``context``'s metadata, else on + ``deployed`` itself.""" + recorded = ((getattr(context, "metadata", None) or {}).get(CONTROLLERS_KEY) or {}).get(deployed.env_id) or {} + if url := recorded.get(CONTROLLER_URL_KEY) or (deployed.metadata or {}).get(CONTROLLER_URL_KEY): return str(url) raise RuntimeError(f"env '{deployed.env_id}' has no {CONTROLLER_URL_KEY}: run osworld_deploy_controller on it first") diff --git a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py index 7307af0b1..c81ea1b00 100644 --- a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py +++ b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py @@ -89,12 +89,16 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: image = Artifact.get(self.controller_image_artifact_id or DEFAULT_CONTROLLER_IMAGE_ARTIFACT_ID) if not isinstance(image, DockerImageArtifact): raise TypeError(f"the controller image {image.id!r} is a {image.type} artifact, not a docker image") - env = self.controller_env(deployed, guest_url) - logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' (image={image.id} v{image.version}, " - f"env={sorted(env)})") provider = skip_refusing(get_sandbox_provider(), lambda link: pull_problem(image.pinned_name, on_this_machine=link.ON_THIS_MACHINE), what=f"sandbox '{sandbox_name}'") + # The guest server and the model endpoint as the controller's sandbox reaches them: a local container reaches + # this machine's loopback by another name. + env = self.controller_env(deployed, provider.url_from_sandbox(guest_url)) + if "LITELLM_BASE_URL" in env: + env["LITELLM_BASE_URL"] = provider.url_from_sandbox(env["LITELLM_BASE_URL"]) + logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' (image={image.id} v{image.version}, " + f"env={sorted(env)})") sandbox = await provider.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, cpu=self.cpu, memory=self.memory_mb, disk_size_gb=10, timeout=self.ttl_seconds, attribution=deploy_attribution(self, context)) @@ -104,7 +108,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: tunnel_urls={str(p): u for p, u in sandbox.tunnel_urls.items()} if sandbox.tunnel_urls else None)) controller_url = sandbox.tunnel_urls[CONTROLLER_PORT] await controller.wait_ready(controller_url) - deployed.metadata = {**(deployed.metadata or {}), guest.CONTROLLER_URL_KEY: controller_url, - guest.CONTROLLER_SANDBOX_ID_KEY: sandbox.sandbox_id, guest.GUEST_URL_KEY: guest_url} + recorded = {guest.CONTROLLER_URL_KEY: controller_url, guest.CONTROLLER_SANDBOX_ID_KEY: sandbox.sandbox_id, + guest.GUEST_URL_KEY: guest_url} + deployed.metadata = {**(deployed.metadata or {}), **recorded} + context.metadata.setdefault(guest.CONTROLLERS_KEY, {})[self.env_id] = recorded logger.info(f"OSWorld controller for env '{self.env_id}' at {controller_url} (sandbox {sandbox.sandbox_id})") return context diff --git a/src/agent_env/task_step/task_steps/osworld/evaluate.py b/src/agent_env/task_step/task_steps/osworld/evaluate.py index 796d96f2f..d5f75bee9 100644 --- a/src/agent_env/task_step/task_steps/osworld/evaluate.py +++ b/src/agent_env/task_step/task_steps/osworld/evaluate.py @@ -110,7 +110,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") - url = guest.controller_url(deployed) + url = guest.controller_url(deployed, context) await self._stage_result_files(url, context) prompt_response = context.prompt_responses[-1] if context.prompt_responses else None if self.osworld_v2_task_path: diff --git a/src/agent_env/task_step/task_steps/osworld/initialize.py b/src/agent_env/task_step/task_steps/osworld/initialize.py index aa91f2dda..1978aed75 100644 --- a/src/agent_env/task_step/task_steps/osworld/initialize.py +++ b/src/agent_env/task_step/task_steps/osworld/initialize.py @@ -45,7 +45,7 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") - url = guest.controller_url(deployed) + url = guest.controller_url(deployed, context) if self.osworld_v2_task_path: logger.info(f"OSWorld-V2 initialize '{self.id}': task_ref={self.osworld_v2_task_path}") # /initialize_task refuses to run before a reset. diff --git a/tst/integration/computer/osworld/envs/osworld-desk/env.toml b/tst/integration/computer/osworld/envs/osworld-desk/env.toml new file mode 100644 index 000000000..7bd5c8571 --- /dev/null +++ b/tst/integration/computer/osworld/envs/osworld-desk/env.toml @@ -0,0 +1,4 @@ +type = "computer" +computer = "ubuntu-24.04" +driver = "reference" +apps = ["osworld-server"] diff --git a/tst/integration/computer/osworld/tasks/classic.json b/tst/integration/computer/osworld/tasks/classic.json new file mode 100644 index 000000000..0bcad1248 --- /dev/null +++ b/tst/integration/computer/osworld/tasks/classic.json @@ -0,0 +1,18 @@ +[ + {"id": "env", "type": "deploy_env", "env_id": "osworld-desk"}, + {"id": "controller", "type": "osworld_deploy_controller", "env_id": "osworld-desk", "controller_image_artifact_id": "controller", + "depends_on": ["env"]}, + {"id": "init", "type": "osworld_initialize", "env_id": "osworld-desk", "depends_on": ["controller"], + "init_config": [{"type": "execute", + "parameters": {"command": ["bash", "-c", "echo osworld-e2e > /tmp/osworld-e2e.txt"]}}]}, + {"id": "pass", "type": "osworld_evaluate", "env_id": "osworld-desk", "verifier_id": "pass", + "depends_on": ["init"], + "evaluator": {"funcs": "check_include_exclude", + "result": {"type": "vm_command_line", "command": ["cat", "/tmp/osworld-e2e.txt"]}, + "expected": {"type": "rule", "rules": {"include": ["osworld-e2e"], "exclude": []}}}}, + {"id": "fail", "type": "osworld_evaluate", "env_id": "osworld-desk", "verifier_id": "fail", + "depends_on": ["pass"], + "evaluator": {"funcs": "check_include_exclude", + "result": {"type": "vm_command_line", "command": ["cat", "/tmp/osworld-e2e.txt"]}, + "expected": {"type": "rule", "rules": {"include": ["not-there"], "exclude": []}}}} +] diff --git a/tst/integration/computer/osworld_test.py b/tst/integration/computer/osworld_test.py new file mode 100644 index 000000000..536486e37 --- /dev/null +++ b/tst/integration/computer/osworld_test.py @@ -0,0 +1,54 @@ +"""OSWorld on a computer env, end to end: the osworld bundle's desk, the ubuntu-24.04 base with the reference driver and +the osworld-server app staged from the guest server's source, deployed; the OSWorld controller deployed beside it by +osworld_deploy_controller, in a container of its own that reaches the guest server only through the endpoint the +computer's card advertises, with the env's control token; osworld_initialize running a classic init_config's execute +on the computer through the controller; and osworld_evaluate grading what that wrote, by the controller's own +check_include_exclude over a vm_command_line getter, 1.0 for a rule it meets and 0.0 for one it doesn't. + +Needs Docker, the local registry on localhost:5000, the guest server's source (scale-cua's +``environments/linux/src/vm/setup/osworld_server``) at ``OSWORLD_SERVER_SOURCE``, and a controller image that takes a +remote base URL and headers (scale-cua #360) at ``OSWORLD_CONTROLLER_IMAGE``; skipped without the last two.""" + +import os +import shutil +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from agent_env.cli import cli +from agent_env.osworld import stage_server +from tst.integration.computer.conftest import run_context + +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] + +OSWORLD = Path(__file__).parent / "osworld" +SOURCE = os.environ.get("OSWORLD_SERVER_SOURCE") +CONTROLLER_IMAGE = os.environ.get("OSWORLD_CONTROLLER_IMAGE") + + +@pytest.fixture +def bundle(tmp_path) -> Path: + """The osworld bundle with the app staged from the guest server's source and the controller image named.""" + if not SOURCE or not CONTROLLER_IMAGE: + pytest.skip("needs OSWORLD_SERVER_SOURCE and OSWORLD_CONTROLLER_IMAGE") + root = shutil.copytree(OSWORLD, tmp_path / "osworld") + stage_server(Path(SOURCE), root / "envs" / "osworld-server") + (root / "artifacts" / "controller").mkdir(parents=True) + (root / "artifacts" / "controller" / "artifact.toml").write_text( + f'type = "docker_image"\nref = "{CONTROLLER_IMAGE}"\n') + return root + + +def test_a_classic_task_is_set_up_and_graded_on_a_computer_through_the_controller_beside_it(state, bundle): + ran = CliRunner().invoke(cli, ["run", str(bundle), "--task", "classic", "--no-emulate"]) + + assert ran.exit_code == 0, ran.output + context = run_context(ran.output) + verifications = context["metadata"]["verifications"] + assert verifications["pass"]["score"] == 1.0, verifications["pass"] + assert verifications["fail"]["score"] == 0.0, verifications["fail"] + (deployed,) = context["deployed_envs"] + recorded = context["metadata"]["osworld_controllers"][deployed["env_id"]] + assert recorded["osworld_guest_url"].endswith("/svc/osworld-server/osworld") + assert recorded["osworld_controller_sandbox_id"] in {s["sandbox_id"] for s in context["deployed_sandboxes"]} diff --git a/tst/unit/osworld/steps_test.py b/tst/unit/osworld/steps_test.py index 717564bf6..408974cf7 100644 --- a/tst/unit/osworld/steps_test.py +++ b/tst/unit/osworld/steps_test.py @@ -72,6 +72,9 @@ def __init__(self): self.create_container = AsyncMock(return_value=SimpleNamespace( sandbox_id="sb-1", mode="container", type="docker", tunnel_urls={8000: CONTROLLER})) + def url_from_sandbox(self, url): + return url + @pytest.fixture def deploy(monkeypatch): @@ -109,12 +112,29 @@ async def test_deploy_runs_the_controller_against_the_card_endpoint_and_records_ deploy.ready.assert_awaited_once_with(CONTROLLER) assert env.metadata == {"existing": "kept", guest.CONTROLLER_URL_KEY: CONTROLLER, guest.CONTROLLER_SANDBOX_ID_KEY: "sb-1", guest.GUEST_URL_KEY: f"{ENV_URL}/apps/7/osworld"} + # The run's stored context keeps them too: the env's record is stored once, when it is deployed. + assert context.metadata[guest.CONTROLLERS_KEY] == {"desk": { + guest.CONTROLLER_URL_KEY: CONTROLLER, guest.CONTROLLER_SANDBOX_ID_KEY: "sb-1", + guest.GUEST_URL_KEY: f"{ENV_URL}/apps/7/osworld"}} [sandbox] = context.deployed_sandboxes assert (sandbox.sandbox_name, sandbox.sandbox_id, sandbox.tunnel_urls) == \ ("osworld-controller-desk", "sb-1", {"8000": CONTROLLER}) assert TOKEN not in caplog.text and LITELLM not in caplog.text +@pytest.mark.asyncio +async def test_deploy_gives_the_controller_the_guest_server_and_model_endpoint_as_its_sandbox_reaches_them(deploy): + deploy.provider.url_from_sandbox = lambda url: url.replace("https://", "https://from-sandbox.") + env = computer(card("/apps/7/osworld")) + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(_context(env)) + + sent = deploy.provider.create_container.call_args.kwargs["env"] + assert sent["REMOTE_VM_IP"] == f"{ENV_URL}/apps/7/osworld".replace("https://", "https://from-sandbox.") + assert sent["LITELLM_BASE_URL"] == "https://from-sandbox.model.example" + # Later steps call the guest server from this process, by the address it reaches. + assert env.metadata[guest.GUEST_URL_KEY] == f"{ENV_URL}/apps/7/osworld" + + @pytest.mark.asyncio async def test_deploy_refuses_a_computer_without_the_app_before_creating_a_sandbox(deploy, monkeypatch): monkeypatch.setattr(guest, "CARD_WAIT_S", 0) @@ -337,3 +357,12 @@ async def test_evaluate_refuses_to_grade_an_unstaged_computer(ctl, files, answer await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", evaluator={"funcs": "compare_table"}).execute(context) assert "/evaluate_async" not in ctl.paths() + + +def test_the_controller_is_found_in_the_runs_context_when_the_env_record_was_stored_without_it(): + env = computer(card("/apps/7/osworld")) + context = _context(env) + context.metadata[guest.CONTROLLERS_KEY] = {"desk": {guest.CONTROLLER_URL_KEY: CONTROLLER}} + assert guest.controller_url(env, context) == CONTROLLER + with pytest.raises(RuntimeError, match="run osworld_deploy_controller"): + guest.controller_url(env, _context(computer(card("/apps/7/osworld")))) From 7ba4a3954ec8a4786a5324ab739c9abf9e3a39bb Mon Sep 17 00:00:00 2001 From: William Loo Date: Wed, 11 Nov 2026 23:39:57 -0800 Subject: [PATCH 3/5] test(osworld): the end-to-end test skips as the osworld_sources capability (FD-5492) It needs scale-cua's guest server source and a controller image, which this repository doesn't carry, so it skips without them. CI's skip policy fails a skip that isn't a declared capability gap, so the skip names osworld_sources, the slow tier allows it as it does mcp_server_sources, and CONTRIBUTING asks a change to the OSWorld code to run the test with both set and say so. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/local-backends.yml | 2 +- CONTRIBUTING.md | 5 ++++- tst/integration/computer/osworld_test.py | 8 ++++---- tst/util/capabilities.py | 16 ++++++++++++++++ 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/.github/workflows/local-backends.yml b/.github/workflows/local-backends.yml index 71e12e08e..a7816cd0e 100644 --- a/.github/workflows/local-backends.yml +++ b/.github/workflows/local-backends.yml @@ -248,7 +248,7 @@ jobs: run: | .venv/bin/python .github/scripts/check_skip_policy.py junit-integration-slow.xml \ --allow model_endpoint_configured remote_sandbox default_a2a_agent mcp_server_sources windows_media \ - omarchy_image + omarchy_image osworld_sources # The tests' times, for .github/scripts/slow_test_durations.py to rebalance the shards with. - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 if: always() diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2c6e81d47..3f1826f7b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -52,7 +52,10 @@ all four must pass, and so must `plugin-api` (see "Making changes"). The `instal the lockfile, and must pass when it runs. The gateway's virtual-clock tests build MCP servers from sources outside this repository, found through `AGENT_ENV_TEST_MCP_SERVERS_DIR` (one `/Dockerfile` per server), and skip without them; when your change touches the virtual -clock, run `make int-test-slow` with it set, and say so in the pull request. A test may skip only when a capability is genuinely unavailable, and the skip reason +clock, run `make int-test-slow` with it set, and say so in the pull request. The OSWorld end-to-end test +(`tst/integration/computer/osworld_test.py`) needs scale-cua's OSWorld guest server source at `OSWORLD_SERVER_SOURCE` +and a controller image at `OSWORLD_CONTROLLER_IMAGE`, and skips without them; when your change touches +`src/agent_env/osworld` or the `osworld_*` steps, run it with them set, and say so in the pull request. A test may skip only when a capability is genuinely unavailable, and the skip reason must name it (`agentenv-capability-missing: `, see `tst/util/capabilities.py`); the `check_skip_policy` step fails the run otherwise. The `clean-install` job runs the same script as `make clean-install-test` on every pull request, and must pass too, as must the `ui` job, which diff --git a/tst/integration/computer/osworld_test.py b/tst/integration/computer/osworld_test.py index 536486e37..209b44a39 100644 --- a/tst/integration/computer/osworld_test.py +++ b/tst/integration/computer/osworld_test.py @@ -7,7 +7,8 @@ Needs Docker, the local registry on localhost:5000, the guest server's source (scale-cua's ``environments/linux/src/vm/setup/osworld_server``) at ``OSWORLD_SERVER_SOURCE``, and a controller image that takes a -remote base URL and headers (scale-cua #360) at ``OSWORLD_CONTROLLER_IMAGE``; skipped without the last two.""" +remote base URL and headers (scale-cua #360) at ``OSWORLD_CONTROLLER_IMAGE``; skipped without them, as the +``osworld_sources`` capability.""" import os import shutil @@ -19,8 +20,9 @@ from agent_env.cli import cli from agent_env.osworld import stage_server from tst.integration.computer.conftest import run_context +from tst.util.capabilities import skip_without_osworld_sources -pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow] +pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow, skip_without_osworld_sources()] OSWORLD = Path(__file__).parent / "osworld" SOURCE = os.environ.get("OSWORLD_SERVER_SOURCE") @@ -30,8 +32,6 @@ @pytest.fixture def bundle(tmp_path) -> Path: """The osworld bundle with the app staged from the guest server's source and the controller image named.""" - if not SOURCE or not CONTROLLER_IMAGE: - pytest.skip("needs OSWORLD_SERVER_SOURCE and OSWORLD_CONTROLLER_IMAGE") root = shutil.copytree(OSWORLD, tmp_path / "osworld") stage_server(Path(SOURCE), root / "envs" / "osworld-server") (root / "artifacts" / "controller").mkdir(parents=True) diff --git a/tst/util/capabilities.py b/tst/util/capabilities.py index 1d4d46d17..2c6231489 100644 --- a/tst/util/capabilities.py +++ b/tst/util/capabilities.py @@ -35,6 +35,11 @@ #: found through ``AGENT_ENV_TEST_MCP_SERVERS_DIR`` (one ``/Dockerfile`` per server). MCP_SERVER_SOURCES = "mcp_server_sources" +#: The OSWorld end-to-end test (``tst/integration/computer/osworld_test.py``) stages the osworld-server app from the +#: OSWorld guest server's source at ``OSWORLD_SERVER_SOURCE`` and runs the controller image ``OSWORLD_CONTROLLER_IMAGE``, +#: both from scale-cua, which this repository doesn't carry. +OSWORLD_SOURCES = "osworld_sources" + #: The TRL contract tier (``tst/trl``) trains with TRL on CPU; trl and torch aren't in the dev extra, and only that #: tier's CI job installs them. TRL = "trl_installed" @@ -254,3 +259,14 @@ def skip_without_windows_media() -> pytest.MarkDecorator: def skip_without_kvm() -> pytest.MarkDecorator: """Collection-time ``skipif`` for tests that boot a computer's disk form on local Docker.""" return pytest.mark.skipif(not kvm_is_available(), reason=missing_capability_reason(KVM)) + + +def osworld_sources_are_available() -> bool: + """Whether the OSWorld guest server's source and a controller image are named for the OSWorld end-to-end test.""" + source = os.environ.get("OSWORLD_SERVER_SOURCE") + return bool(source and (Path(source) / "main.py").is_file() and os.environ.get("OSWORLD_CONTROLLER_IMAGE")) + + +def skip_without_osworld_sources() -> pytest.MarkDecorator: + """Collection-time ``skipif`` for the OSWorld end-to-end test.""" + return pytest.mark.skipif(not osworld_sources_are_available(), reason=missing_capability_reason(OSWORLD_SOURCES)) From 2ba780639ef7a575cb40253ca84f6cfc7189dbe2 Mon Sep 17 00:00:00 2001 From: William Loo Date: Thu, 12 Nov 2026 12:30:38 -0800 Subject: [PATCH 4/5] fix(osworld): the controller requires a token, each link gets its own addresses, and review fixes (FD-5492) - osworld_deploy_controller starts the controller with CONTROLLER_AUTH_TOKEN, an HMAC of the env's control token, and refuses one that answers a call without it: its port can be reachable from the computer, and so by the agent. Every controller call carries the token (controller.Endpoint). - The controller's sandbox is created link by link, each given the guest server and model endpoint as its own sandboxes reach them; a chain's url_from_sandbox can't say for the link it ends up on. - osworld_evaluate refuses a result/expected/options list whose length isn't the funcs'. - A failed setup is reported by its error alone, bounded, never the controller's whole answer. - Context failures raise StepContextError; the controller's disk is a named constant. - The osworld-server app pins Pillow 13.0.0, past every Pillow advisory. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../apps/osworld-server/requirements.txt | 2 +- src/agent_env/osworld/controller.py | 56 +++++++++--- src/agent_env/osworld/guest.py | 23 ++++- .../task_steps/osworld/deploy_controller.py | 63 +++++++++---- .../task_step/task_steps/osworld/evaluate.py | 51 +++++++---- .../task_steps/osworld/initialize.py | 24 +++-- tst/integration/computer/osworld_test.py | 4 +- tst/unit/osworld/controller_test.py | 48 ++++++++-- tst/unit/osworld/guest_test.py | 14 ++- tst/unit/osworld/helpers.py | 13 ++- tst/unit/osworld/steps_test.py | 91 +++++++++++++++++-- 11 files changed, 303 insertions(+), 86 deletions(-) diff --git a/src/agent_env/osworld/apps/osworld-server/requirements.txt b/src/agent_env/osworld/apps/osworld-server/requirements.txt index 9cd169f3d..40747b17f 100644 --- a/src/agent_env/osworld/apps/osworld-server/requirements.txt +++ b/src/agent_env/osworld/apps/osworld-server/requirements.txt @@ -3,7 +3,7 @@ httpx>=0.27.0 python3-xlib==0.15 PyAutoGUI==0.9.54 -Pillow==10.1.0 +Pillow==13.0.0 requests flask numpy diff --git a/src/agent_env/osworld/controller.py b/src/agent_env/osworld/controller.py index 109ac25de..a79df63ce 100644 --- a/src/agent_env/osworld/controller.py +++ b/src/agent_env/osworld/controller.py @@ -1,11 +1,13 @@ """The OSWorld controller beside a computer, as the ``osworld_*`` steps call it over HTTP: ``/reset``, ``/initialize_task``, ``/setup``, and ``/evaluate``, the last started with ``/evaluate_async`` and polled with short -requests, since a judge can run for minutes and a long silent connection is reaped by what's between.""" +requests, since a judge can run for minutes and a long silent connection is reaped by what's between. Every call +carries the bearer token ``osworld_deploy_controller`` started the controller with (``Endpoint``).""" from __future__ import annotations import asyncio import json import logging +from dataclasses import dataclass, field from typing import Any import httpx @@ -24,6 +26,21 @@ TASK_REF_TTL_S = 6 * 60 * 60 +@dataclass(frozen=True) +class Endpoint: + """A controller's address and what every call to it sends: the bearer token it requires.""" + + url: str + headers: dict[str, str] = field(default_factory=dict, repr=False) + + def at(self, path: str) -> str: + return f"{self.url.rstrip('/')}{path}" + + +def bearer_headers(token: str) -> dict[str, str]: + return {"Authorization": f"Bearer {token}"} + + class ControllerCallError(RuntimeError): """A controller call that failed in transport or answered with an HTTP error; carries the status and body.""" @@ -46,31 +63,31 @@ def _parsed(path: str, response: httpx.Response, *, expect_json: bool) -> dict: status=response.status_code, body=response.text) from e -async def post(controller_url: str, path: str, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S, +async def post(endpoint: Endpoint, path: str, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S, expect_json: bool = True) -> dict: """POST ``payload`` to the controller's ``path``. With ``expect_json`` False a 2xx body that isn't JSON (``/reset`` answers a bare string) returns ``{}``. Raises ControllerCallError on a transport failure or an HTTP error.""" try: async with httpx.AsyncClient() as client: - response = await client.post(f"{controller_url.rstrip('/')}{path}", json=payload, timeout=timeout) + response = await client.post(endpoint.at(path), json=payload, headers=endpoint.headers, timeout=timeout) except httpx.HTTPError as e: raise ControllerCallError(f"controller {path} transport failure: {type(e).__name__}: {e}", body=str(e)) from e return _parsed(path, response, expect_json=expect_json) -async def evaluate(controller_url: str, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S) -> dict: +async def evaluate(endpoint: Endpoint, payload: dict, *, timeout: float = DEFAULT_TIMEOUT_S) -> dict: """``/evaluate`` with ``payload``: started with ``/evaluate_async`` and polled at ``/evaluate_result`` until it's done, or the blocking ``/evaluate`` on a controller without the async routes.""" - base = controller_url.rstrip("/") try: async with httpx.AsyncClient() as client: - start = await client.post(f"{base}/evaluate_async", json=payload, timeout=60) + start = await client.post(endpoint.at("/evaluate_async"), json=payload, headers=endpoint.headers, + timeout=60) except httpx.HTTPError as e: raise ControllerCallError(f"controller /evaluate_async transport failure: {type(e).__name__}: {e}", body=str(e)) from e if start.status_code == 404: logger.info("controller has no /evaluate_async; evaluating with the blocking /evaluate") - return await post(controller_url, "/evaluate", payload, timeout=timeout) + return await post(endpoint, "/evaluate", payload, timeout=timeout) job_id = _parsed("/evaluate_async", start, expect_json=True).get("job_id") if not job_id: raise ControllerCallError(f"controller /evaluate_async returned no job_id: {start.text[:300]}") @@ -82,8 +99,8 @@ async def evaluate(controller_url: str, payload: dict, *, timeout: float = DEFAU await asyncio.sleep(_POLL_INTERVAL_S) try: async with httpx.AsyncClient() as client: - poll = await client.get(f"{base}/evaluate_result", params={"job_id": job_id}, - timeout=_POLL_READ_TIMEOUT_S) + poll = await client.get(endpoint.at("/evaluate_result"), params={"job_id": job_id}, + headers=endpoint.headers, timeout=_POLL_READ_TIMEOUT_S) except httpx.HTTPError as e: logger.warning(f"controller /evaluate_result poll failed, polling again: {type(e).__name__}: {e}") continue @@ -103,7 +120,7 @@ async def evaluate(controller_url: str, payload: dict, *, timeout: float = DEFAU body=json.dumps(data)[:500]) -async def wait_ready(controller_url: str, *, timeout: float = 300.0, interval: float = 5.0) -> None: +async def wait_ready(endpoint: Endpoint, *, timeout: float = 300.0, interval: float = 5.0) -> None: """Return once the controller answers ``READY_PATH``; TimeoutError after ``timeout``.""" loop = asyncio.get_running_loop() deadline = loop.time() + timeout @@ -111,14 +128,29 @@ async def wait_ready(controller_url: str, *, timeout: float = 300.0, interval: f while loop.time() < deadline: try: async with httpx.AsyncClient() as client: - response = await client.get(f"{controller_url.rstrip('/')}{READY_PATH}", timeout=15) + response = await client.get(endpoint.at(READY_PATH), headers=endpoint.headers, timeout=15) if response.status_code == 200: return last = f"HTTP {response.status_code}" except httpx.HTTPError as e: last = repr(e) await asyncio.sleep(interval) - raise TimeoutError(f"OSWorld controller {controller_url} not ready within {timeout}s (last: {last})") + raise TimeoutError(f"OSWorld controller {endpoint.url} not ready within {timeout}s (last: {last})") + + +async def require_auth(endpoint: Endpoint) -> None: + """Raise unless the controller refuses a call without its token, as one whose image predates + ``CONTROLLER_AUTH_TOKEN`` doesn't: its port may be reachable from the computer, and so by the agent.""" + try: + async with httpx.AsyncClient() as client: + response = await client.get(endpoint.at(READY_PATH), timeout=15) + except httpx.HTTPError as e: + raise ControllerCallError(f"controller {READY_PATH} transport failure: {type(e).__name__}: {e}") from e + if response.status_code != 401: + raise ControllerCallError( + f"the OSWorld controller at {endpoint.url} answered a call without its token with HTTP " + f"{response.status_code}, not 401: its image doesn't enforce CONTROLLER_AUTH_TOKEN, so the agent could " + f"call its grader; use a controller image that does", status=response.status_code) def resolve_task_ref(ref: str) -> str: diff --git a/src/agent_env/osworld/guest.py b/src/agent_env/osworld/guest.py index aec58bc6b..e2a6826d0 100644 --- a/src/agent_env/osworld/guest.py +++ b/src/agent_env/osworld/guest.py @@ -1,17 +1,20 @@ """Where a computer serves the OSWorld guest server: the endpoint its card advertises for the ``osworld-guest/v1`` extension, on the card itself or one of its ``children_environments``, resolved against the computer's address. ``osworld_deploy_controller`` resolves it once and records it on the env, with the controller's address, under the -metadata keys here.""" +metadata keys here. The controller requires a bearer token derived from the env's control token +(``controller_token``), so every step that calls it can rebuild it, and none is stored.""" from __future__ import annotations import asyncio +import hashlib +import hmac import logging from typing import Any import httpx from agentenv_protocol import client as protocol_v1 -from agent_env.osworld import GUEST_EXTENSION +from agent_env.osworld import GUEST_EXTENSION, controller logger = logging.getLogger(__name__) @@ -26,6 +29,8 @@ #: the guest extension, and how often it reads the card again. CARD_WAIT_S = 120.0 CARD_POLL_S = 5.0 +#: What the controller's token is derived for, so it is never the control token itself. +CONTROLLER_TOKEN_PURPOSE = b"osworld-controller" def guest_endpoint(card: dict | None) -> str | None: @@ -76,3 +81,17 @@ def controller_url(deployed: Any, context: Any = None) -> str: if url := recorded.get(CONTROLLER_URL_KEY) or (deployed.metadata or {}).get(CONTROLLER_URL_KEY): return str(url) raise RuntimeError(f"env '{deployed.env_id}' has no {CONTROLLER_URL_KEY}: run osworld_deploy_controller on it first") + + +def controller_token(deployed: Any) -> str: + """The bearer token the controller beside ``deployed`` requires: an HMAC of the env's control token, which only + the framework holds. Raises RuntimeError for an env with none, whose controller would be open to anyone.""" + if not deployed.control_token: + raise RuntimeError(f"env '{deployed.env_id}' has no control token to derive its OSWorld controller's from") + return hmac.new(deployed.control_token.encode(), CONTROLLER_TOKEN_PURPOSE + b":" + deployed.env_id.encode(), + hashlib.sha256).hexdigest() + + +def controller_endpoint(deployed: Any, context: Any = None) -> controller.Endpoint: + """The controller recorded for ``deployed`` (``controller_url``), with its token on every call.""" + return controller.Endpoint(controller_url(deployed, context), controller.bearer_headers(controller_token(deployed))) diff --git a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py index c81ea1b00..8f2de27e3 100644 --- a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py +++ b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py @@ -2,8 +2,9 @@ The controller holds the grader, a task's checks and the judge's key, and a computer's apps run as its session user, the agent's, so it runs apart from the computer, in remote mode against the guest server at the endpoint the -computer's card advertises (``agent_env.osworld.guest``), with the env's control token on every call. Once it answers, -the step records the controller's and the guest server's addresses on the env for ``osworld_initialize`` and +computer's card advertises (``agent_env.osworld.guest``), with the env's control token on every call. Its own port may be +reachable from the computer, so it requires a bearer token derived from that control token (``CONTROLLER_AUTH_TOKEN``) +and the step refuses a controller that answers without it. Once it answers, the step records the controller's and the guest server's addresses on the env for ``osworld_initialize`` and ``osworld_evaluate``, and the sandbox in ``context.deployed_sandboxes`` for teardown.""" from __future__ import annotations @@ -18,8 +19,9 @@ from agent_env.entity.refs import EntityRef from agent_env.osworld import controller, guest from agent_env.providers.sandbox_providers.chained_sandbox_provider import skip_refusing -from agent_env.providers.sandbox_providers.sandbox_provider import get_sandbox_provider -from agent_env.task_step.context import DeployedSandbox, TaskStepContext +from agent_env.providers.sandbox_providers.sandbox import Sandbox +from agent_env.providers.sandbox_providers.sandbox_provider import SandboxProvider, get_sandbox_provider +from agent_env.task_step.context import DeployedSandbox, StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency logger = logging.getLogger(__name__) @@ -27,6 +29,10 @@ #: The controller image's artifact id when a step names none. DEFAULT_CONTROLLER_IMAGE_ARTIFACT_ID = "cua-controller" CONTROLLER_PORT = 8000 +#: The disk the controller's sandbox gets: its image, the grader's scratch and the files it fetches to grade. +CONTROLLER_DISK_GB = 10 +#: The env var the controller reads the bearer token every call to it must carry from. +CONTROLLER_TOKEN_ENV = "CONTROLLER_AUTH_TOKEN" #: A controller's sandbox name, by its computer env's id. SANDBOX_NAME_PREFIX = "osworld-controller-" @@ -63,10 +69,11 @@ def from_dict(cls, data: dict) -> OSWorldDeployControllerTaskStep: memory_mb=data.get("memory_mb", 4096), env_vars=data.get("env_vars")) def controller_env(self, deployed, guest_url: str) -> dict[str, str]: - """Remote mode against the guest server, with the env's control token on every call and the model endpoint's - key for the judge.""" + """Remote mode against the guest server, with the env's control token on every call, the token every call to + the controller must carry, and the model endpoint's key for the judge.""" env = {**(self.env_vars or {}), "PROVIDER": "remote", "OS_TYPE": "Ubuntu", "REMOTE_VM_IP": guest_url, - "REMOTE_SERVER_HEADERS": json.dumps(deployed.control_headers)} + "REMOTE_SERVER_HEADERS": json.dumps(deployed.control_headers), + CONTROLLER_TOKEN_ENV: guest.controller_token(deployed)} config = get_config() try: env["LITELLM_API_KEY"] = config.get_litellm_api_key() @@ -81,10 +88,10 @@ def controller_env(self, deployed, guest_url: str) -> dict[str, str]: async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") sandbox_name = f"{SANDBOX_NAME_PREFIX}{self.env_id}" if any(s.sandbox_name == sandbox_name for s in context.deployed_sandboxes): - raise RuntimeError(f"Sandbox '{sandbox_name}' is already deployed") + raise StepContextError(f"Sandbox '{sandbox_name}' is already deployed") guest_url = await guest.resolve_guest_url(deployed) image = Artifact.get(self.controller_image_artifact_id or DEFAULT_CONTROLLER_IMAGE_ARTIFACT_ID) if not isinstance(image, DockerImageArtifact): @@ -92,25 +99,41 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: provider = skip_refusing(get_sandbox_provider(), lambda link: pull_problem(image.pinned_name, on_this_machine=link.ON_THIS_MACHINE), what=f"sandbox '{sandbox_name}'") - # The guest server and the model endpoint as the controller's sandbox reaches them: a local container reaches - # this machine's loopback by another name. - env = self.controller_env(deployed, provider.url_from_sandbox(guest_url)) - if "LITELLM_BASE_URL" in env: - env["LITELLM_BASE_URL"] = provider.url_from_sandbox(env["LITELLM_BASE_URL"]) - logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' (image={image.id} v{image.version}, " - f"env={sorted(env)})") - sandbox = await provider.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, - cpu=self.cpu, memory=self.memory_mb, disk_size_gb=10, - timeout=self.ttl_seconds, attribution=deploy_attribution(self, context)) + sandbox = await self._create(provider, deployed, guest_url, image, context) context.deployed_sandboxes.append(DeployedSandbox( sandbox_name=sandbox_name, sandbox_id=sandbox.sandbox_id, sandbox_mode=sandbox.mode, sandbox_type=sandbox.type, tunnel_urls={str(p): u for p, u in sandbox.tunnel_urls.items()} if sandbox.tunnel_urls else None)) controller_url = sandbox.tunnel_urls[CONTROLLER_PORT] - await controller.wait_ready(controller_url) + endpoint = controller.Endpoint(controller_url, controller.bearer_headers(guest.controller_token(deployed))) + await controller.wait_ready(endpoint) + await controller.require_auth(endpoint) recorded = {guest.CONTROLLER_URL_KEY: controller_url, guest.CONTROLLER_SANDBOX_ID_KEY: sandbox.sandbox_id, guest.GUEST_URL_KEY: guest_url} deployed.metadata = {**(deployed.metadata or {}), **recorded} context.metadata.setdefault(guest.CONTROLLERS_KEY, {})[self.env_id] = recorded logger.info(f"OSWorld controller for env '{self.env_id}' at {controller_url} (sandbox {sandbox.sandbox_id})") return context + + async def _create(self, provider: SandboxProvider, deployed, guest_url: str, image: DockerImageArtifact, + context: TaskStepContext) -> Sandbox: + """The controller's sandbox on the first of ``provider``'s links that creates it, each given the guest server and + the model endpoint as its own sandboxes reach them: a local container reaches this machine's loopback by + another name, which a chain of providers can't say for the link it ends up on.""" + errors = [] + for link in provider.links: + env = self.controller_env(deployed, link.url_from_sandbox(guest_url)) + if "LITELLM_BASE_URL" in env: + env["LITELLM_BASE_URL"] = link.url_from_sandbox(env["LITELLM_BASE_URL"]) + logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' on {type(link).__name__} " + f"(image={image.id} v{image.version}, env={sorted(env)})") + try: + return await link.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, + cpu=self.cpu, memory=self.memory_mb, disk_size_gb=CONTROLLER_DISK_GB, + timeout=self.ttl_seconds, + attribution=deploy_attribution(self, context)) + except Exception as e: + logger.warning(f"{type(link).__name__}: the OSWorld controller's sandbox failed: " + f"{type(e).__name__}: {str(e)[:300]}; trying the next provider") + errors.append(f"{type(link).__name__}: {type(e).__name__}: {str(e)[:300]}") + raise RuntimeError(f"no provider created the OSWorld controller's sandbox: {'; '.join(errors)}") diff --git a/src/agent_env/task_step/task_steps/osworld/evaluate.py b/src/agent_env/task_step/task_steps/osworld/evaluate.py index d5f75bee9..78e037b90 100644 --- a/src/agent_env/task_step/task_steps/osworld/evaluate.py +++ b/src/agent_env/task_step/task_steps/osworld/evaluate.py @@ -12,9 +12,9 @@ from agent_env.entity.refs import EntityRef from agent_env.osworld import controller, guest -from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.context import PromptResponse, StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency -from agent_env.task_step.task_steps.osworld.initialize import run_setup +from agent_env.task_step.task_steps.osworld.initialize import SETUP_ERROR_CHARS, run_setup from agent_env.task_step.task_steps.verifiers.scoring import ScoreAggregator, aggregate_score logger = logging.getLogger(__name__) @@ -89,7 +89,7 @@ def _skipped(self, context: TaskStepContext, prompt_response: PromptResponse) -> return self._record(context, {"results": [{"id": "prompt_error", "score": FAIL_SCORE, "result": False, "message": message}], "score": FAIL_SCORE}) - async def _stage_result_files(self, controller_url: str, context: TaskStepContext) -> None: + async def _stage_result_files(self, endpoint: controller.Endpoint, context: TaskStepContext) -> None: """A run's earlier outputs, ``step_overrides`` ``stage_files`` ``[{url, path}]``, downloaded onto the computer before anything reads it, to grade a finished run again on a fresh one. Raises on any failure: an unstaged computer would be graded on the task's own inputs.""" @@ -103,30 +103,31 @@ async def _stage_result_files(self, controller_url: str, context: TaskStepContex f"each url http(s) and fetchable from the computer; got {files!r}") entries = [{"url": f["url"], "path": f["path"]} for f in files] logger.info(f"osworld_evaluate '{self.verifier_id}': staging {[e['path'] for e in entries]} before grading") - await run_setup(controller_url, [{"type": "download", "parameters": {"files": entries}}], + await run_setup(endpoint, [{"type": "download", "parameters": {"files": entries}}], what=f"osworld_evaluate '{self.verifier_id}' staging result files") async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") - url = guest.controller_url(deployed, context) - await self._stage_result_files(url, context) + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") + endpoint = guest.controller_endpoint(deployed, context) + await self._stage_result_files(endpoint, context) prompt_response = context.prompt_responses[-1] if context.prompt_responses else None if self.osworld_v2_task_path: if prompt_response is not None and prompt_response.error_type: return self._skipped(context, prompt_response) - return self._record(context, await self._evaluate_v2(url)) + return self._record(context, await self._evaluate_v2(endpoint)) if postconfig := self.evaluator.get("postconfig"): try: - await run_setup(url, postconfig, what=f"osworld_evaluate '{self.verifier_id}' postconfig") + await run_setup(endpoint, postconfig, what=f"osworld_evaluate '{self.verifier_id}' postconfig") except Exception as e: - logger.warning(f"osworld_evaluate '{self.verifier_id}': postconfig failed, grading anyway: {e}") + logger.warning(f"osworld_evaluate '{self.verifier_id}': postconfig failed, grading anyway: " + f"{type(e).__name__}: {str(e)[:SETUP_ERROR_CHARS]}") budget = budget_exhaustion(prompt_response) if prompt_response is not None and prompt_response.error_type and budget is None: return self._skipped(context, prompt_response) - verification = await self._evaluate_funcs(url, prompt_response) + verification = await self._evaluate_funcs(endpoint, prompt_response) if budget is not None: # The score is the end state's, earned after the agent ran out of room; the marker beside the verifications # keeps such runs separable, since score readers project only each verifier's score. @@ -138,11 +139,11 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: self.verifier_id: marker} return self._record(context, verification) - async def _evaluate_v2(self, controller_url: str) -> dict: + async def _evaluate_v2(self, endpoint: controller.Endpoint) -> dict: logger.info(f"OSWorld-V2 evaluate '{self.verifier_id}': task_ref={self.osworld_v2_task_path}") try: answer = await controller.evaluate( - controller_url, {"osworld_v2_task_url": controller.resolve_task_ref(self.osworld_v2_task_path)}) + endpoint, {"osworld_v2_task_url": controller.resolve_task_ref(self.osworld_v2_task_path)}) except controller.ControllerCallError as e: logger.warning(f"OSWorld-V2 evaluate '{self.verifier_id}': controller error: {e}") return {"results": [{"id": "osworld_v2", "score": FAIL_SCORE, "result": False, "message": str(e)[:500]}], @@ -159,9 +160,20 @@ def _funcs(self) -> list[tuple[str, Any, Any, dict]]: if not isinstance(funcs, list): return [(funcs, self.evaluator.get("result"), self.evaluator.get("expected"), self.evaluator.get("options") or {})] - n = len(funcs) - return list(zip(funcs, self.evaluator.get("result") or [None] * n, self.evaluator.get("expected") or [None] * n, - [o or {} for o in self.evaluator.get("options") or [{}] * n])) + return list(zip(funcs, self._per_func("result", len(funcs), None), self._per_func("expected", len(funcs), None), + [o or {} for o in self._per_func("options", len(funcs), {})], strict=True)) + + def _per_func(self, key: str, n: int, absent: Any) -> list: + """The evaluator's ``key``, one entry per func, or ``absent`` for each when it has none; a list of another + length is refused, since zipping it would grade on fewer checks than the task names.""" + value = self.evaluator.get(key) + if not value: + return [absent] * n + if not isinstance(value, list) or len(value) != n: + got = f"a list of {len(value)}" if isinstance(value, list) else type(value).__name__ + raise ValueError(f"osworld_evaluate '{self.verifier_id}': the evaluator's {key!r} must be a list of {n}, " + f"one per func, got {got}") + return value def _payload(self, func: str, result: Any, expected: Any, options: dict, instruction: str | None) -> dict: options = dict(options) @@ -181,7 +193,7 @@ def _payload(self, func: str, result: Any, expected: Any, options: dict, instruc payload["instruction"] = instruction return payload - async def _evaluate_funcs(self, controller_url: str, prompt_response: PromptResponse | None) -> dict: + async def _evaluate_funcs(self, endpoint: controller.Endpoint, prompt_response: PromptResponse | None) -> dict: answer_text = prompt_response.response if prompt_response is not None else None instruction = self.evaluator.get("instruction") or (prompt_response.prompt_text if prompt_response else None) results, scores, metrics = [], [], [] @@ -197,13 +209,14 @@ async def _evaluate_funcs(self, controller_url: str, prompt_response: PromptResp sub = {"id": "infeasible_on_feasible", "score": FAIL_SCORE, "result": False, "message": "Agent declared infeasible on a feasible task"} else: - answer = await controller.evaluate(controller_url, + answer = await controller.evaluate(endpoint, self._payload(func, result, expected, options, instruction)) sub = {"id": func, "score": controller.normalize_score(answer.get("score", FAIL_SCORE)), "result": not answer.get("error", False), "message": answer.get("message", "")} passed = True except Exception as e: - logger.warning(f"osworld_evaluate '{self.verifier_id}': {func} failed: {e}") + logger.warning(f"osworld_evaluate '{self.verifier_id}': {func} failed: {type(e).__name__}: " + f"{str(e)[:SETUP_ERROR_CHARS]}") sub = {"id": func, "score": FAIL_SCORE, "result": False, "message": f"Exception: {str(e)[:300]}"} passed = False results.append({"func": func, "results": [sub]}) diff --git a/src/agent_env/task_step/task_steps/osworld/initialize.py b/src/agent_env/task_step/task_steps/osworld/initialize.py index 1978aed75..14368eed6 100644 --- a/src/agent_env/task_step/task_steps/osworld/initialize.py +++ b/src/agent_env/task_step/task_steps/osworld/initialize.py @@ -9,11 +9,14 @@ from agent_env.entity.refs import EntityRef from agent_env.osworld import controller, guest -from agent_env.task_step.context import TaskStepContext +from agent_env.task_step.context import StepContextError, TaskStepContext from agent_env.task_step.task_step import TaskStep, TaskStepDependency logger = logging.getLogger(__name__) +#: How much of a controller's setup error an exception or a log line carries. +SETUP_ERROR_CHARS = 300 + class OSWorldInitializeTaskStep(TaskStep): type: ClassVar[str] = "osworld_initialize" @@ -44,31 +47,32 @@ def from_dict(cls, data: dict) -> OSWorldInitializeTaskStep: async def execute(self, context: TaskStepContext) -> TaskStepContext: deployed = next((d for d in context.deployed_envs if d.env_id == self.env_id), None) if deployed is None: - raise RuntimeError(f"Env '{self.env_id}' not found in context.deployed_envs") - url = guest.controller_url(deployed, context) + raise StepContextError(f"Env '{self.env_id}' not found in context.deployed_envs") + endpoint = guest.controller_endpoint(deployed, context) if self.osworld_v2_task_path: logger.info(f"OSWorld-V2 initialize '{self.id}': task_ref={self.osworld_v2_task_path}") # /initialize_task refuses to run before a reset. - await controller.post(url, "/reset", {}, expect_json=False) - answer = await controller.post(url, "/initialize_task", { + await controller.post(endpoint, "/reset", {}, expect_json=False) + answer = await controller.post(endpoint, "/initialize_task", { "id": self.id, "osworld_v2_task_url": controller.resolve_task_ref(self.osworld_v2_task_path)}) if answer.get("success") is False or answer.get("error"): raise RuntimeError(f"OSWorld-V2 initialize failed: {answer.get('error') or answer.get('message')}") return context if self.init_config: - answer = await run_setup(url, self.init_config, what=f"osworld_initialize '{self.id}'") + answer = await run_setup(endpoint, self.init_config, what=f"osworld_initialize '{self.id}'") if sites := answer.get("results"): context.metadata.setdefault("osworld_seeded_sites", []).extend(sites) return context -async def run_setup(controller_url: str, config: list[dict], *, what: str) -> dict: +async def run_setup(endpoint: controller.Endpoint, config: list[dict], *, what: str) -> dict: """``config`` through the controller's ``/setup``; raises when it reports failure or a site it seeded failed, since a seed that answers 200 with every site failed leaves the apps on their own fixtures.""" - answer = await controller.post(controller_url, "/setup", {"config": config}) + answer = await controller.post(endpoint, "/setup", {"config": config}) if not answer.get("success", True): - raise RuntimeError(f"{what}: setup failed: {answer.get('error') or answer}") + # The error alone, bounded: the rest of the answer can be large and carry the task's content. + raise RuntimeError(f"{what}: setup failed: {str(answer.get('error') or 'no error given')[:SETUP_ERROR_CHARS]}") if failed := [site for site in answer.get("results") or [] if isinstance(site, dict) and not site.get("ok", True)]: - detail = "; ".join(f"{site.get('app')}: {site.get('error')}" for site in failed) + detail = "; ".join(f"{site.get('app')}: {str(site.get('error'))[:SETUP_ERROR_CHARS]}" for site in failed) raise RuntimeError(f"{what}: setup failed for {len(failed)} site(s): {detail}") return answer diff --git a/tst/integration/computer/osworld_test.py b/tst/integration/computer/osworld_test.py index 209b44a39..c4c8a2728 100644 --- a/tst/integration/computer/osworld_test.py +++ b/tst/integration/computer/osworld_test.py @@ -7,8 +7,8 @@ Needs Docker, the local registry on localhost:5000, the guest server's source (scale-cua's ``environments/linux/src/vm/setup/osworld_server``) at ``OSWORLD_SERVER_SOURCE``, and a controller image that takes a -remote base URL and headers (scale-cua #360) at ``OSWORLD_CONTROLLER_IMAGE``; skipped without them, as the -``osworld_sources`` capability.""" +remote base URL and headers and requires ``CONTROLLER_AUTH_TOKEN`` (scale-cua #360) at ``OSWORLD_CONTROLLER_IMAGE``; +skipped without them, as the ``osworld_sources`` capability.""" import os import shutil diff --git a/tst/unit/osworld/controller_test.py b/tst/unit/osworld/controller_test.py index 2e1a54476..732e5b4de 100644 --- a/tst/unit/osworld/controller_test.py +++ b/tst/unit/osworld/controller_test.py @@ -10,12 +10,12 @@ import pytest from agent_env.osworld import controller -from tst.unit.osworld.helpers import CONTROLLER, Controller, route_httpx +from tst.unit.osworld.helpers import CONTROLLER, CONTROLLER_TOKEN, ENDPOINT, Controller, route_httpx @pytest.fixture def ctl(monkeypatch): - stand_in = Controller() + stand_in = Controller(token=CONTROLLER_TOKEN) route_httpx(monkeypatch, controller, stand_in) monkeypatch.setattr(controller.asyncio, "sleep", AsyncMock()) return stand_in @@ -24,26 +24,26 @@ def ctl(monkeypatch): @pytest.mark.asyncio async def test_a_call_returns_its_answer_and_an_http_error_raises_with_its_status(ctl): ctl.answers["/setup"] = {"success": True} - assert await controller.post(CONTROLLER, "/setup", {"config": []}) == {"success": True} + assert await controller.post(ENDPOINT, "/setup", {"config": []}) == {"success": True} ctl.answers["/setup"] = httpx.Response(500, text="boom") with pytest.raises(controller.ControllerCallError) as raised: - await controller.post(CONTROLLER, "/setup", {}) + await controller.post(ENDPOINT, "/setup", {}) assert raised.value.status == 500 and raised.value.body == "boom" @pytest.mark.asyncio async def test_a_bare_answer_is_empty_when_json_isnt_expected(ctl): ctl.answers["/reset"] = httpx.Response(200, text="ok") - assert await controller.post(CONTROLLER, "/reset", {}, expect_json=False) == {} + assert await controller.post(ENDPOINT, "/reset", {}, expect_json=False) == {} with pytest.raises(controller.ControllerCallError, match="non-JSON"): - await controller.post(CONTROLLER, "/reset", {}) + await controller.post(ENDPOINT, "/reset", {}) @pytest.mark.asyncio async def test_a_transport_failure_raises(ctl): ctl.answers["/setup"] = httpx.ConnectError("refused") with pytest.raises(controller.ControllerCallError, match="transport failure"): - await controller.post(CONTROLLER, "/setup", {}) + await controller.post(ENDPOINT, "/setup", {}) @pytest.mark.asyncio @@ -51,7 +51,7 @@ async def test_evaluate_is_started_and_polled_until_done(ctl): ctl.answers["/evaluate_async"] = {"job_id": "j1"} ctl.answers["/evaluate_result"] = [httpx.ConnectError("blip"), {"status": "running"}, httpx.Response(200, text=""), {"status": "done", "result": {"score": 1}}] - assert await controller.evaluate(CONTROLLER, {"func": "compare_table"}) == {"score": 1} + assert await controller.evaluate(ENDPOINT, {"func": "compare_table"}) == {"score": 1} assert ctl.bodies("/evaluate_async") == [{"func": "compare_table"}] assert {r.url.params["job_id"] for r in ctl.seen if r.url.path == "/evaluate_result"} == {"j1"} @@ -60,7 +60,7 @@ async def test_evaluate_is_started_and_polled_until_done(ctl): async def test_evaluate_falls_back_to_the_blocking_route(ctl): ctl.answers["/evaluate_async"] = httpx.Response(404) ctl.answers["/evaluate"] = {"score": 0.5} - assert await controller.evaluate(CONTROLLER, {"func": "x"}) == {"score": 0.5} + assert await controller.evaluate(ENDPOINT, {"func": "x"}) == {"score": 0.5} assert ctl.paths() == ["/evaluate_async", "/evaluate"] @@ -69,7 +69,35 @@ async def test_a_failed_job_raises(ctl): ctl.answers["/evaluate_async"] = {"job_id": "j1"} ctl.answers["/evaluate_result"] = {"status": "error", "message": "judge died"} with pytest.raises(controller.ControllerCallError, match="judge died"): - await controller.evaluate(CONTROLLER, {}) + await controller.evaluate(ENDPOINT, {}) + + +@pytest.mark.asyncio +async def test_every_call_carries_the_controllers_token(ctl): + ctl.answers.update({"/setup": {"success": True}, "/configure": {}, "/evaluate_async": {"job_id": "j1"}, + "/evaluate_result": {"status": "done", "result": {"score": 1}}}) + await controller.post(ENDPOINT, "/setup", {}) + await controller.wait_ready(ENDPOINT) + await controller.evaluate(ENDPOINT, {}) + assert {r.headers["authorization"] for r in ctl.seen} == {f"Bearer {CONTROLLER_TOKEN}"} + with pytest.raises(controller.ControllerCallError) as raised: + await controller.post(controller.Endpoint(CONTROLLER), "/setup", {}) + assert raised.value.status == 401 + + +@pytest.mark.asyncio +async def test_a_controller_that_answers_without_its_token_is_refused(ctl): + ctl.answers["/configure"] = {} + await controller.require_auth(ENDPOINT) + ctl.token = None + with pytest.raises(controller.ControllerCallError, match="doesn't enforce CONTROLLER_AUTH_TOKEN") as raised: + await controller.require_auth(ENDPOINT) + assert raised.value.status == 200 + assert all("authorization" not in r.headers for r in ctl.seen) + + +def test_the_token_is_kept_out_of_an_endpoints_repr(): + assert CONTROLLER_TOKEN not in repr(ENDPOINT) def test_scores_normalize_to_floats(): diff --git a/tst/unit/osworld/guest_test.py b/tst/unit/osworld/guest_test.py index 61aa1a0ac..aac9d7b19 100644 --- a/tst/unit/osworld/guest_test.py +++ b/tst/unit/osworld/guest_test.py @@ -9,7 +9,7 @@ import pytest from agent_env.osworld import guest -from tst.unit.osworld.helpers import ENV_URL, card, computer +from tst.unit.osworld.helpers import ENV_URL, TOKEN, card, computer def test_the_endpoint_is_the_one_an_app_advertises_for_the_guest_extension(): @@ -56,3 +56,15 @@ def test_the_controller_url_is_the_one_recorded(): assert guest.controller_url(computer(osworld_controller_url="https://c")) == "https://c" with pytest.raises(RuntimeError, match="run osworld_deploy_controller"): guest.controller_url(computer()) + + +def test_the_controllers_token_is_derived_from_the_envs_control_token_and_is_not_it(): + token = guest.controller_token(computer()) + assert token == guest.controller_token(computer()) and len(token) == 64 + assert TOKEN not in token + other = computer() + other.control_token = "another-instance" + assert guest.controller_token(other) != token + other.control_token = None + with pytest.raises(RuntimeError, match="no control token"): + guest.controller_token(other) diff --git a/tst/unit/osworld/helpers.py b/tst/unit/osworld/helpers.py index 269d9f576..3e002f329 100644 --- a/tst/unit/osworld/helpers.py +++ b/tst/unit/osworld/helpers.py @@ -8,7 +8,7 @@ from agentenv_protocol.types import WELL_KNOWN_PATH from agent_env.env.env import DeployedEnv -from agent_env.osworld import GUEST_EXTENSION +from agent_env.osworld import GUEST_EXTENSION, controller, guest ENV_URL = "https://computer.example" TOKEN = "s3cret-control-token" @@ -31,13 +31,17 @@ def computer(card_: dict | None = None, **metadata) -> DeployedEnv: @dataclass class Controller: - """The controller's routes, answered from ``answers`` by path (a list is answered in turn), and the requests made.""" + """The controller's routes, answered from ``answers`` by path (a list is answered in turn), and the requests made; + with a ``token``, a request without it as its bearer token is answered 401, as the controller's inbound auth does.""" answers: dict[str, object] = field(default_factory=dict) seen: list[httpx.Request] = field(default_factory=list) + token: str | None = None def handler(self, request: httpx.Request) -> httpx.Response: self.seen.append(request) + if self.token is not None and request.headers.get("authorization") != f"Bearer {self.token}": + return httpx.Response(401, json={"detail": "unauthorized"}) answer = self.answers.get(request.url.path, httpx.Response(404, text="no route")) if isinstance(answer, list): answer = answer.pop(0) @@ -60,3 +64,8 @@ def client(*args, **kwargs): return real(*args, transport=httpx.MockTransport(controller.handler), **kwargs) monkeypatch.setattr(module.httpx, "AsyncClient", client) + + +#: The token the controller beside ``computer()`` requires, and that controller as a step calls it. +CONTROLLER_TOKEN = guest.controller_token(computer()) +ENDPOINT = controller.Endpoint(CONTROLLER, controller.bearer_headers(CONTROLLER_TOKEN)) diff --git a/tst/unit/osworld/steps_test.py b/tst/unit/osworld/steps_test.py index 408974cf7..74764167c 100644 --- a/tst/unit/osworld/steps_test.py +++ b/tst/unit/osworld/steps_test.py @@ -16,15 +16,18 @@ from agent_env.artifact import DockerImageArtifact from agent_env.osworld import controller, guest -from agent_env.task_step.context import PromptResponse, TaskStepContext +from agent_env.task_step.context import PromptResponse, StepContextError, TaskStepContext from agent_env.task_step.registry import get_task_step_registry from agent_env.task_step.task_steps.osworld import deploy_controller as deploy_module from agent_env.task_step.task_steps.osworld.deploy_controller import OSWorldDeployControllerTaskStep from agent_env.task_step.task_steps.osworld.evaluate import OSWorldEvaluateTaskStep from agent_env.task_step.task_steps.osworld.initialize import OSWorldInitializeTaskStep -from tst.unit.osworld.helpers import CONTROLLER, ENV_URL, TOKEN, Controller, card, computer, route_httpx +from tst.unit.osworld.helpers import (CONTROLLER, CONTROLLER_TOKEN, ENDPOINT, ENV_URL, TOKEN, Controller, card, computer, + route_httpx) LITELLM = "sk-model-key" +#: The real check, which the deploy fixture stands in for. +REQUIRE_AUTH = controller.require_auth def _context(env, *prompts: PromptResponse) -> TaskStepContext: @@ -40,7 +43,7 @@ def _ready(): @pytest.fixture def ctl(monkeypatch): - stand_in = Controller() + stand_in = Controller(token=CONTROLLER_TOKEN) route_httpx(monkeypatch, controller, stand_in) monkeypatch.setattr(controller.asyncio, "sleep", AsyncMock()) return stand_in @@ -85,12 +88,14 @@ def deploy(monkeypatch): config.get_litellm_api_key.return_value = LITELLM config.get_litellm_base_url.return_value = "https://model.example" ready = AsyncMock() + auth = AsyncMock() get = MagicMock(return_value=image) monkeypatch.setattr(deploy_module.Artifact, "get", get) monkeypatch.setattr(deploy_module, "get_config", lambda: config) monkeypatch.setattr(deploy_module, "get_sandbox_provider", lambda: provider) monkeypatch.setattr(deploy_module.controller, "wait_ready", ready) - return SimpleNamespace(provider=provider, ready=ready, get=get) + monkeypatch.setattr(deploy_module.controller, "require_auth", auth) + return SimpleNamespace(provider=provider, ready=ready, auth=auth, get=get) @pytest.mark.asyncio @@ -108,8 +113,11 @@ async def test_deploy_runs_the_controller_against_the_card_endpoint_and_records_ assert kwargs["env"] == {"JUDGE_MODEL": "m", "PROVIDER": "remote", "OS_TYPE": "Ubuntu", "REMOTE_VM_IP": f"{ENV_URL}/apps/7/osworld", "REMOTE_SERVER_HEADERS": json.dumps({"Authorization": f"Bearer {TOKEN}"}), + "CONTROLLER_AUTH_TOKEN": CONTROLLER_TOKEN, "LITELLM_API_KEY": LITELLM, "LITELLM_BASE_URL": "https://model.example"} - deploy.ready.assert_awaited_once_with(CONTROLLER) + # It's called with the token, and must refuse a call without it. + deploy.ready.assert_awaited_once_with(ENDPOINT) + deploy.auth.assert_awaited_once_with(ENDPOINT) assert env.metadata == {"existing": "kept", guest.CONTROLLER_URL_KEY: CONTROLLER, guest.CONTROLLER_SANDBOX_ID_KEY: "sb-1", guest.GUEST_URL_KEY: f"{ENV_URL}/apps/7/osworld"} # The run's stored context keeps them too: the env's record is stored once, when it is deployed. @@ -119,7 +127,7 @@ async def test_deploy_runs_the_controller_against_the_card_endpoint_and_records_ [sandbox] = context.deployed_sandboxes assert (sandbox.sandbox_name, sandbox.sandbox_id, sandbox.tunnel_urls) == \ ("osworld-controller-desk", "sb-1", {"8000": CONTROLLER}) - assert TOKEN not in caplog.text and LITELLM not in caplog.text + assert TOKEN not in caplog.text and CONTROLLER_TOKEN not in caplog.text and LITELLM not in caplog.text @pytest.mark.asyncio @@ -135,6 +143,46 @@ async def test_deploy_gives_the_controller_the_guest_server_and_model_endpoint_a assert env.metadata[guest.GUEST_URL_KEY] == f"{ENV_URL}/apps/7/osworld" +@pytest.mark.asyncio +async def test_deploy_on_a_chain_gives_each_link_the_addresses_its_own_sandboxes_reach(deploy): + """A chain can't say how the link it ends up on reaches this machine, so each link tried is given its own.""" + failing, local = _Provider(), _Provider() + failing.create_container.side_effect = RuntimeError("no capacity") + local.url_from_sandbox = lambda url: url.replace("https://", "https://from-local.") + deploy.provider.links = [failing, local] + env = computer(card("/apps/7/osworld")) + context = _context(env) + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(context) + + assert failing.create_container.call_args.kwargs["env"]["REMOTE_VM_IP"] == f"{ENV_URL}/apps/7/osworld" + sent = local.create_container.call_args.kwargs["env"] + assert sent["REMOTE_VM_IP"] == f"{ENV_URL}/apps/7/osworld".replace("https://", "https://from-local.") + assert sent["LITELLM_BASE_URL"] == "https://from-local.model.example" + assert sent["CONTROLLER_AUTH_TOKEN"] == CONTROLLER_TOKEN + deploy.provider.create_container.assert_not_called() + assert [s.sandbox_id for s in context.deployed_sandboxes] == ["sb-1"] + + +@pytest.mark.asyncio +async def test_deploy_fails_naming_each_link_when_none_creates_the_sandbox(deploy): + deploy.provider.create_container.side_effect = RuntimeError("no capacity") + with pytest.raises(RuntimeError, match="_Provider: RuntimeError: no capacity"): + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(_context(computer())) + + +@pytest.mark.asyncio +async def test_deploy_refuses_a_controller_that_answers_without_its_token(deploy, monkeypatch): + stand_in = Controller(answers={"/configure": {}}) + route_httpx(monkeypatch, controller, stand_in) + deploy.auth.side_effect = REQUIRE_AUTH + context = _context(computer()) + with pytest.raises(controller.ControllerCallError, match="doesn't enforce CONTROLLER_AUTH_TOKEN"): + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(context) + # Its sandbox is in the context, for teardown, but nothing records it as the env's controller. + assert [s.sandbox_id for s in context.deployed_sandboxes] == ["sb-1"] + assert guest.CONTROLLERS_KEY not in context.metadata + + @pytest.mark.asyncio async def test_deploy_refuses_a_computer_without_the_app_before_creating_a_sandbox(deploy, monkeypatch): monkeypatch.setattr(guest, "CARD_WAIT_S", 0) @@ -150,7 +198,7 @@ async def test_deploy_refuses_a_computer_without_the_app_before_creating_a_sandb async def test_deploy_twice_for_one_env_is_refused(deploy): context = _context(computer()) await OSWorldDeployControllerTaskStep(id="a", version=1, env_id="desk").execute(context) - with pytest.raises(RuntimeError, match="already deployed"): + with pytest.raises(StepContextError, match="already deployed"): await OSWorldDeployControllerTaskStep(id="b", version=1, env_id="desk").execute(context) @@ -268,6 +316,18 @@ async def test_evaluate_classic_sends_each_func_and_combines_by_conj(ctl, conj, {"id": "compare_table", "score": float(scores[0]), "result": True, "message": f"got {scores[0]}"}] +@pytest.mark.asyncio +@pytest.mark.parametrize("key, value", [("options", [{}]), ("result", [{"type": "vm_file", "path": "/a"}] * 3), + ("expected", {"type": "cloud_file"})]) +async def test_evaluate_classic_refuses_per_func_lists_that_dont_match_its_funcs(ctl, key, value): + """Zipped, a shorter list would grade on fewer checks than the task names.""" + evaluator = {"funcs": ["compare_table", "agent_judge"], key: value} + with pytest.raises(ValueError, match=f"'{key}' must be a list of 2, one per func"): + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", + evaluator=evaluator).execute(_context(_ready(), _prompt())) + assert "/evaluate_async" not in ctl.paths() + + @pytest.mark.asyncio async def test_evaluate_classic_one_func_keeps_its_partial_score(ctl): _scores(ctl, 0.4) @@ -328,6 +388,23 @@ async def test_evaluate_runs_postconfig_once_and_grades_when_it_fails(ctl): assert context.metadata["verifications"]["v"]["score"] == 1.0 +@pytest.mark.asyncio +async def test_a_failed_setup_is_reported_by_its_error_alone_bounded(ctl, caplog): + """The rest of a setup answer can be large and carry the task's content, so neither the step nor its log has it.""" + _scores(ctl, 1) + ctl.answers["/setup"] = {"success": False, "echo": "TASK-CONTENT " * 100} + with caplog.at_level(logging.WARNING): + await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", evaluator={ + "funcs": "compare_table", "postconfig": [{"type": "sleep"}]}).execute(_context(_ready(), _prompt())) + assert "postconfig failed, grading anyway: RuntimeError" in caplog.text and "no error given" in caplog.text + assert "TASK-CONTENT" not in caplog.text + ctl.answers["/setup"] = {"success": False, "error": "x" * 5000} + with pytest.raises(RuntimeError) as raised: + await OSWorldInitializeTaskStep(id="init", version=1, env_id="desk", + init_config=[{"type": "sleep"}]).execute(_context(_ready())) + assert len(str(raised.value)) < 400 + + def _regrade(context, files): context.metadata["user_overrides"] = {"step_params": {"e": {"stage_files": files}}} return context From b1e420c680fd04eaba37f7ab8752511e36a4bac1 Mon Sep 17 00:00:00 2001 From: William Loo Date: Thu, 12 Nov 2026 12:40:42 -0800 Subject: [PATCH 5/5] fix(osworld): prepare the controller image on each link, and read every shipped agent's budget stop (FD-5492) - osworld_deploy_controller calls the link's prepare_image before creating its container, so an image stored as its build context alone is built rather than pulled, and lets go a sandbox the link says can't run the image for the next link's, as the provider chain does. - osworld_evaluate reads Claude Code's and Codex's "Reached maximum number of turns (N)" and each CLI's " timed out after Ns" as a budget stop, grading the desktop and marking partial credit, beside the harness's own sentences. A stall, a tool's timeout or any other error still isn't one. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../task_steps/osworld/deploy_controller.py | 23 +++++++-- .../task_step/task_steps/osworld/evaluate.py | 8 +++- tst/unit/osworld/steps_test.py | 48 +++++++++++++++++-- 3 files changed, 69 insertions(+), 10 deletions(-) diff --git a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py index 8f2de27e3..3ad62ce9c 100644 --- a/src/agent_env/task_step/task_steps/osworld/deploy_controller.py +++ b/src/agent_env/task_step/task_steps/osworld/deploy_controller.py @@ -127,13 +127,28 @@ async def _create(self, provider: SandboxProvider, deployed, guest_url: str, ima env["LITELLM_BASE_URL"] = link.url_from_sandbox(env["LITELLM_BASE_URL"]) logger.info(f"Deploying the OSWorld controller for env '{self.env_id}' on {type(link).__name__} " f"(image={image.id} v{image.version}, env={sorted(env)})") + attribution = deploy_attribution(self, context) try: - return await link.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, - cpu=self.cpu, memory=self.memory_mb, disk_size_gb=CONTROLLER_DISK_GB, - timeout=self.ttl_seconds, - attribution=deploy_attribution(self, context)) + # An image stored as its build context alone is built here, where the link's sandboxes get it. + await link.prepare_image(image, attribution=attribution) + sandbox = await link.create_container(image_name=image.pinned_name, port=CONTROLLER_PORT, env=env, + cpu=self.cpu, memory=self.memory_mb, + disk_size_gb=CONTROLLER_DISK_GB, timeout=self.ttl_seconds, + attribution=attribution) except Exception as e: logger.warning(f"{type(link).__name__}: the OSWorld controller's sandbox failed: " f"{type(e).__name__}: {str(e)[:300]}; trying the next provider") errors.append(f"{type(link).__name__}: {type(e).__name__}: {str(e)[:300]}") + continue + # A sandbox that can't run the image, a VM that can't load it say, is let go for the next link's. + if problem := link.sandbox_problem(sandbox, image): + logger.warning(f"{type(link).__name__}: {problem}; trying the next provider") + errors.append(f"{type(link).__name__}: {problem}") + try: + await sandbox.terminate() + except Exception as e: + logger.warning(f"{type(link).__name__}: terminating sandbox {sandbox.sandbox_id} failed: " + f"{type(e).__name__}") + continue + return sandbox raise RuntimeError(f"no provider created the OSWorld controller's sandbox: {'; '.join(errors)}") diff --git a/src/agent_env/task_step/task_steps/osworld/evaluate.py b/src/agent_env/task_step/task_steps/osworld/evaluate.py index 78e037b90..a4f168374 100644 --- a/src/agent_env/task_step/task_steps/osworld/evaluate.py +++ b/src/agent_env/task_step/task_steps/osworld/evaluate.py @@ -31,8 +31,12 @@ CONTEXT_FUNCS = JUDGE_FUNCS | {"compare_pdfs", "compare_csv", "compare_images", "compare_python_pure_text", "diff_text_file", "website_verifier"} # An agent that ran out of turns or time leaves a desktop worth grading; every other stop leaves none. Matched on the -# harness's own sentence, since its error types name both kinds. -_BUDGET_EXHAUSTED = re.compile(r"Max turns \(\d+\) reached|agent timed out after \d+", re.I) +# sentences the harness and the shipped agents stop with, since their error types name both kinds: the harness's +# "Max turns (N) reached" and "agent timed out after N"; Claude Code's and Codex's "Reached maximum number of turns +# (N)"; each CLI's " timed out after Ns" (claude CLI, codex CLI, gemini CLI, Letta Code, hermes, opencode). +_BUDGET_EXHAUSTED = re.compile( + r"Max turns \(\d+\) reached|Reached maximum number of turns \(\d+\)" + r"|\b(?:agent|CLI|Letta Code|hermes|opencode) timed out after \d+", re.I) def budget_exhaustion(prompt_response: PromptResponse | None) -> str | None: diff --git a/tst/unit/osworld/steps_test.py b/tst/unit/osworld/steps_test.py index 74764167c..2edb09ce0 100644 --- a/tst/unit/osworld/steps_test.py +++ b/tst/unit/osworld/steps_test.py @@ -72,12 +72,17 @@ class _Provider: def __init__(self): self.links = [self] + self.problem = None + self.prepare_image = AsyncMock() self.create_container = AsyncMock(return_value=SimpleNamespace( - sandbox_id="sb-1", mode="container", type="docker", tunnel_urls={8000: CONTROLLER})) + sandbox_id="sb-1", mode="container", type="docker", tunnel_urls={8000: CONTROLLER}, terminate=AsyncMock())) def url_from_sandbox(self, url): return url + def sandbox_problem(self, sandbox, image): + return self.problem + @pytest.fixture def deploy(monkeypatch): @@ -163,6 +168,29 @@ async def test_deploy_on_a_chain_gives_each_link_the_addresses_its_own_sandboxes assert [s.sandbox_id for s in context.deployed_sandboxes] == ["sb-1"] +@pytest.mark.asyncio +async def test_deploy_prepares_the_image_on_each_link_before_creating_its_sandbox(deploy): + """An image stored as its build context alone is built by the link's prepare_image, never pulled.""" + calls = [] + deploy.provider.prepare_image.side_effect = lambda image, **_: calls.append(("prepare", image.id)) + deploy.provider.create_container.side_effect = lambda **_: calls.append(("create", None)) or SimpleNamespace( + sandbox_id="sb-1", mode="container", type="docker", tunnel_urls={8000: CONTROLLER}) + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(_context(computer())) + assert calls == [("prepare", "cua-controller"), ("create", None)] + + +@pytest.mark.asyncio +async def test_deploy_lets_go_a_sandbox_that_cant_run_the_image_for_the_next_links(deploy): + refusing, running = _Provider(), _Provider() + refusing.problem = "the VM can't load the image" + deploy.provider.links = [refusing, running] + context = _context(computer()) + await OSWorldDeployControllerTaskStep(id="ctl", version=1, env_id="desk").execute(context) + refusing.create_container.return_value.terminate.assert_awaited_once() + running.create_container.assert_awaited_once() + assert len(context.deployed_sandboxes) == 1 + + @pytest.mark.asyncio async def test_deploy_fails_naming_each_link_when_none_creates_the_sandbox(deploy): deploy.provider.create_container.side_effect = RuntimeError("no capacity") @@ -363,10 +391,15 @@ async def test_evaluate_infeasible_is_graded_on_the_agents_answer(ctl): @pytest.mark.asyncio -async def test_evaluate_grades_a_run_that_ran_out_of_turns_and_marks_it(ctl): +@pytest.mark.parametrize("message", [ + "Max turns (50) reached without completion.", "agent timed out after 1800s", + "Reached maximum number of turns (50)", "codex CLI timed out after 600s", "claude CLI timed out after 900s", + "gemini CLI timed out after 600s", "Letta Code timed out after 600s", "hermes timed out after 600s", + "opencode timed out after 600s"]) +async def test_evaluate_grades_a_run_that_ran_out_of_turns_and_marks_it(ctl, message): + """Each shipped agent's own words for running out of turns or time.""" _scores(ctl, 1) - context = _context(_ready(), _prompt(error_type="AGENT_ERROR", - error_message="Max turns (50) reached without completion.")) + context = _context(_ready(), _prompt(error_type="AGENT_ERROR", error_message=message)) await OSWorldEvaluateTaskStep(id="e", version=1, env_id="desk", verifier_id="v", evaluator={"funcs": "compare_table"}).execute(context) verification = context.metadata["verifications"]["v"] @@ -443,3 +476,10 @@ def test_the_controller_is_found_in_the_runs_context_when_the_env_record_was_sto assert guest.controller_url(env, context) == CONTROLLER with pytest.raises(RuntimeError, match="run osworld_deploy_controller"): guest.controller_url(env, _context(computer(card("/apps/7/osworld")))) + + +@pytest.mark.parametrize("message", ["codex CLI produced no output for 300s (stalled)", "command timed out after 30s", + "Tool 'click' timed out after 60s", "model returned 500"]) +def test_a_stop_that_isnt_the_agents_budget_isnt_read_as_one(message): + from agent_env.task_step.task_steps.osworld.evaluate import budget_exhaustion + assert budget_exhaustion(_prompt(error_type="AGENT_ERROR", error_message=message)) is None