feat(task-steps): native OSWorld initialize and evaluate on computer envs (FD-5492)

.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()

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
 `<server>/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: <name>`, 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

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

src/agent_env/osworld/__main__.py

@@ -0,0 +1,29 @@
+"""``python -m agent_env.osworld stage --source <guest server source> --out <folder>``: 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())

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"

src/agent_env/osworld/apps/osworld-server/osworld_server_app.py

@@ -0,0 +1,84 @@
+"""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/<route>`` 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, 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:
+    """``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="bash 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()

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==13.0.0
+requests
+flask
+numpy
+lxml

src/agent_env/osworld/apps/osworld-server/start-osworld-server.sh

@@ -0,0 +1,13 @@
+#!/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
+# 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
+    sleep 5
+done

src/agent_env/osworld/controller.py

@@ -0,0 +1,181 @@
+"""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. 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
+
+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
+
+
+@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."""
+
+    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(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(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(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."""
+    try:
+        async with httpx.AsyncClient() as client:
+            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(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]}")
+    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(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
+        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(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
+    last = None
+    while loop.time() < deadline:
+        try:
+            async with httpx.AsyncClient() as client:
+                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 {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:
+    """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