feat(gateway): relay a website browser's CDP addresses at /browser/<key>/cdp/<token>/

src/agent_env/env/gateway/constants.py

@@ -80,6 +80,8 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int:
 EXT_CHANGES_URI = "urn:agentenv:changes/v1"
 EXT_NOTIFICATIONS_URI = "urn:agentenv:notifications/v1"
 EXT_WEBSITE_URI = "urn:agentenv:website/v1"
+EXT_BROWSER_URI = "urn:agentenv:browser/v1"
+EXT_COMPUTER_APP_URI = "urn:agentenv:computer-app/v1"
 WEBSITE_READY_OP = "ready"
 
 # Rows per GET /changes page: the default, and the most a client may ask for.

src/agent_env/env/gateway/gateway.py

@@ -32,6 +32,7 @@
 from pydantic import create_model
 from starlette.requests import Request
 from starlette.responses import Response, StreamingResponse
+from starlette.routing import WebSocketRoute
 from anyio import ClosedResourceError
 
 if TYPE_CHECKING:
@@ -67,6 +68,7 @@
     GatewayMode,
     tool_name_clashes_text,
 )
+from .browser_proxy import PING_INTERVAL_S, BrowserRoutes, TokenRedaction
 from .caller import Caller
 from .changes import ChangeAttribution
 from .child_transport import child_http_client
@@ -80,6 +82,11 @@
 
 # RFC 9110's safe methods: a request using one asks for no change, so the proxy measures no write window for it.
 _SAFE_METHODS = frozenset({"GET", "HEAD", "OPTIONS", "TRACE"})
+# A CDP address at the gateway: the website browser behind the /svc child <key>, discovery and its WebSocket.
+BROWSER_ROUTE = "/browser/{key}/cdp/{token}/{path:path}"
+# The loggers that write a request's path or a forwarded URL: uvicorn's, for HTTP requests and WebSocket handshakes,
+# and httpx's, for each request the gateway forwards.
+_TOKEN_LOGGERS = ("uvicorn.access", "uvicorn.error", "httpx")
 _STEP_API = {"source": "step_api"}
 
 
@@ -127,6 +134,12 @@ def __init__(
         self.website_urls = website_urls or {}
         self.rest_proxy_urls = rest_proxy_urls or {}
         self._site_reserved: dict[str, frozenset[str]] = {}
+        self._browsers = BrowserRoutes(self.rest_proxy_urls, self.REST_PROXY_TIMEOUT, self._log_event,
+                                       self._attribution, self._request_role)
+        for name in _TOKEN_LOGGERS:  # each names a CDP address's path, its token in it, at INFO
+            log = logging.getLogger(name)
+            if not any(isinstance(f, TokenRedaction) for f in log.filters):
+                log.addFilter(TokenRedaction())
         self.gateway_mode = gateway_mode
         self._service_db_url = service_db_url
         self._db_conn: psycopg2.extensions.connection | None = None
@@ -188,6 +201,10 @@ async def list_website_urls() -> list[str]:
             async def proxy_to_service(request: Request) -> Response:
                 return await self._proxy_rest_request(request)
 
+            @self._mcp.custom_route(BROWSER_ROUTE, methods=["GET", "HEAD"])
+            async def browser_discovery(request: Request) -> Response:
+                return await self._browsers.discovery(request)
+
         if self.website_urls:
             @self._mcp.custom_route(
                 "/website/{service_name}/{path:path}",
@@ -396,6 +413,8 @@ def run(self) -> None:
             host=self._mcp.settings.host,
             port=self._mcp.settings.port,
             log_level=self._mcp.settings.log_level.lower(),
+            ws_ping_interval=PING_INTERVAL_S,
+            ws_ping_timeout=PING_INTERVAL_S,
         )
         server = uvicorn.Server(config)
         if self._resolve_i6pn_host():
@@ -418,6 +437,8 @@ def _asgi_app(self):
         lifespan wrapped around the MCP session manager's."""
         inner = self._mcp.streamable_http_app()
         inner.router.lifespan_context = self._process_lifespan(inner.router.lifespan_context)
+        if self.rest_proxy_urls:
+            inner.router.routes.append(WebSocketRoute(BROWSER_ROUTE, self._browsers.relay))
 
         async def app(scope, receive, send):
             if scope["type"] != "http":
@@ -602,6 +623,8 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] |
             return self._json_response({"error": f"Unknown service: {service_name}"}, 404)
         role = self._request_role(request)
         mcp_server = url_map is self.rest_proxy_urls and service_name.startswith("mcp-")
+        if url_map is self.rest_proxy_urls and (cdp := await self._browsers.cdp_answer(request, base_url, role)) is not None:
+            return cdp
         if request.method in _SAFE_METHODS:
             return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server)
         changelog_before = self._query_changelog_id()

src/agent_env/env/gateway/requirements.txt

@@ -3,3 +3,4 @@ httpx
 mcp>=1.25,<2
 pydantic>=2.0
 tzdata
+websockets>=13

tst/integration/computer/browser_computer_test.py

@@ -18,8 +18,6 @@
 import platform
 import re
 import shlex
-import shutil
-import subprocess
 import time
 import urllib.parse
 import uuid
@@ -54,8 +52,6 @@
 ARCH = {"x86_64": "amd64", "amd64": "amd64", "aarch64": "arm64", "arm64": "arm64"}[platform.machine().lower()]
 #: How long an effect may take to reach the page, or the computer to be ready again; how often it's asked meanwhile.
 WAIT_S, POLL_S = 120.0, 0.5
-#: Playwright and Puppeteer, which connect to a cdp address from Node.
-CDP_CLIENTS = Path(__file__).parent / "cdp_clients"
 #: A WebSocket handshake for the browser's relay sent straight to its app server's socket, as code in the computer's
 #: session can, with the token it's given; prints the status line it gets.
 #: The lines of the computer's logs that hold either of two tokens, each shown masked, then one path the driver
@@ -72,22 +68,6 @@
 """
 
 
-@pytest.fixture(scope="module")
-def cdp_clients(tmp_path_factory) -> Path:
-    """The CDP clients installed in a folder of their own, or the test skipped without Node 20 or later and npm."""
-    node, npm = shutil.which("node"), shutil.which("npm")
-    version = subprocess.run([node, "--version"], capture_output=True, text=True).stdout if node else ""
-    if not npm or not (match := re.match(r"v(\d+)", version)) or int(match.group(1)) < 20:
-        pytest.skip("the CDP clients need Node 20 or later and npm")
-    folder = tmp_path_factory.mktemp("cdp-clients")
-    for name in ("package.json", "connect.mjs"):
-        shutil.copy(CDP_CLIENTS / name, folder / name)
-    installed = subprocess.run([npm, "install", "--no-audit", "--no-fund", "--loglevel=error"], cwd=folder,
-                               capture_output=True, text=True, timeout=600)
-    assert installed.returncode == 0, installed.stderr
-    return folder
-
-
 # A page whose title says, as JSON, what reached it: clicks and double clicks on its red box, its input's value, how far
 # it has scrolled and the last keys pressed.
 PAGE = """<title>loading</title><style>body{margin:0;height:4000px;background:rgb(255,255,255)}

tst/integration/computer/browser_gateway_test.py

@@ -0,0 +1,178 @@
+"""agentenv's chromium, deployed alone as a browser computer, behind the gateway: the gateway's image built from its
+folder and run with the computer as its /svc child ``web``. Playwright and Puppeteer reach the browser through the
+gateway's ``/browser/web/cdp/<token>/`` route by the ``http_url`` and the ``ws_url`` the browser's ``cdp`` answer
+gives through the gateway, an address the browser didn't mint is refused, and each connection is a trajectory event
+that, like the gateway's log, never shows its token. A gateway restart or a browser relaunch closes a live connection
+with a code, and a new one then works. Needs Docker, and Node 20 or later with npm for the CDP clients."""
+
+import asyncio
+import json
+import platform
+import socket
+import subprocess
+import uuid
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+import httpx
+import pytest
+from agentenv_protocol import client as protocol_v1
+from agentenv_protocol.website_browser import BROWSER
+from websockets.asyncio.client import connect
+
+import agent_env
+from agent_env.env import Env
+from agent_env.env.envs.computer.utils import inside
+from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT
+from agent_env.providers.sandbox_providers.local_docker.provider import host_url_for
+from agent_env.store.routing import namespace_routing
+from agent_env.task.teardown import teardown_run
+from agent_env.task_step.context import TaskStepContext
+from tst.integration.computer.browser_computer_test import DATA_PAGE, SERVE_SITE, SERVE_START_S, SITE, _clients
+
+pytestmark = [pytest.mark.integration, pytest.mark.int_test_slow]
+
+GATEWAY = Path(agent_env.__file__).parent / "env" / "gateway"
+EXTRA_HOSTS = ["--add-host", "host.docker.internal:host-gateway"] if platform.system() == "Linux" else []
+#: What the browser's relay closes its connections with when it relaunches.
+RELAUNCHED = 1012
+
+
+@pytest.fixture(scope="module")
+def gateway_image():
+    tag = f"agentenv-test-gateway:{uuid.uuid4().hex[:8]}"
+    subprocess.run(["docker", "build", "-q", "-t", tag, str(GATEWAY)], check=True, capture_output=True)
+    yield tag
+    subprocess.run(["docker", "rmi", "-f", tag], capture_output=True)
+
+
+@pytest.mark.asyncio
+async def test_playwright_and_puppeteer_reach_the_browser_through_the_gateway_and_no_trace_shows_a_token(
+        browser_computer, cdp_clients, gateway_image):
+    with namespace_routing():
+        deployed = await Env.get(browser_computer.id).deploy(
+            env_params={"app_settings": {"chromium": {"start_urls": [SITE]}}})
+        context = TaskStepContext(deployed_envs=[deployed])
+        try:
+            serving = asyncio.ensure_future(inside.run(deployed, ["python3", "-c", SERVE_SITE, "900"], timeout_s=960))
+            await asyncio.sleep(SERVE_START_S)
+            async with _gateway(gateway_image, deployed) as (gateway, container):
+                address = await _cdp(gateway, deployed, "ada")
+                together = await _clients(cdp_clients, "together", address)
+                refused = await _clients(cdp_clients, "refused", address)
+                events = await _connections(gateway, at_least=4)
+                logs = subprocess.run(["docker", "logs", container], capture_output=True, text=True)
+            serving.cancel()
+        finally:
+            report = await teardown_run(context)
+
+    token = address["http_url"].split("/cdp/")[1].strip("/")
+    assert address == {"http_url": f"{gateway}/browser/web/cdp/{token}/",
+                       "ws_url": f"ws{gateway.removeprefix('http')}/browser/web/cdp/{token}/devtools/browser"}
+    for client in ("playwright_http", "puppeteer_ws", "playwright_ws"):
+        assert together[client].get("ok", {}).get("title") == "Home", together
+    assert "error" in refused["playwright"] and "error" in refused["puppeteer"], refused
+    opened = [event for event in events if event["opened_at"]]
+    assert len(opened) == 3 and {event["closed_by"] for event in opened} == {"client"}, events
+    assert {(event["closed_by"], event["status"]) for event in events if not event["opened_at"]} == {
+        ("browser", 403)}, events
+    assert token not in json.dumps(events), events
+    assert "/cdp/<redacted>/" in logs.stderr + logs.stdout and token not in logs.stderr + logs.stdout
+    assert not report.failed
+
+
+@pytest.mark.asyncio
+async def test_a_gateway_restart_or_a_browser_relaunch_closes_a_live_connection_with_a_code_and_a_new_one_works(
+        browser_computer, gateway_image):
+    with namespace_routing():
+        deployed = await Env.get(browser_computer.id).deploy(
+            env_params={"app_settings": {"chromium": {"start_urls": [DATA_PAGE]}}})
+        context = TaskStepContext(deployed_envs=[deployed])
+        try:
+            async with _gateway(gateway_image, deployed) as (gateway, container):
+                address = await _cdp(gateway, deployed, "ada")
+                held, first = await _version(address["ws_url"])
+                subprocess.run(["docker", "restart", "-t", "5", container], check=True, capture_output=True)
+                at_restart = await _closed(held)
+                await _until_up(gateway)
+                held, second = await _version(address["ws_url"])
+                await protocol_v1.reset_data(f"{deployed.environment_url}/svc/chromium",
+                                             headers=deployed.control_headers)
+                at_relaunch = await _closed(held)
+                held, third = await _version(address["ws_url"])
+                await held.close()
+        finally:
+            report = await teardown_run(context)
+
+    assert first.startswith("Chrome/") and first == second == third
+    assert at_restart is not None and at_relaunch == RELAUNCHED
+    assert not report.failed
+
+
+@asynccontextmanager
+async def _gateway(image: str, deployed):
+    """The gateway ``image`` run with ``deployed`` as its /svc child ``web``, at a fixed port a restart keeps; yields
+    its address and its container."""
+    with socket.socket() as probe:
+        probe.bind(("127.0.0.1", 0))
+        port = probe.getsockname()[1]
+    name = f"agentenv-test-gateway-{uuid.uuid4().hex[:8]}"
+    browser = host_url_for(deployed.environment_url.rstrip("/"), "local")
+    subprocess.run(["docker", "run", "-d", "--name", name, *EXTRA_HOSTS, "-e", f"REST_PROXY_URLS=web={browser}",
+                    "-p", f"127.0.0.1:{port}:{AGENT_ENV_GATEWAY_MCP_PORT}", image], check=True, capture_output=True)
+    try:
+        url = f"http://127.0.0.1:{port}"
+        await _until_up(url)
+        yield url, name
+    finally:
+        subprocess.run(["docker", "rm", "-f", name], capture_output=True)
+
+
+async def _until_up(url: str, timeout_s: float = 90) -> None:
+    deadline = asyncio.get_running_loop().time() + timeout_s
+    async with httpx.AsyncClient(timeout=5) as http:
+        while asyncio.get_running_loop().time() < deadline:
+            try:
+                if (await http.get(f"{url}/.well-known/agent-env.json")).status_code == 200:
+                    return
+            except httpx.HTTPError:
+                pass
+            await asyncio.sleep(1)
+    raise TimeoutError(f"the gateway at {url} never came up")
+
+
+async def _cdp(gateway: str, deployed, user: str) -> dict:
+    """``user``'s CDP address, asked of the browser through the gateway's /svc route, with the computer's control
+    token, as agentenv asks it when the browser is a gateway's child."""
+    app = next(child for child in deployed.environment_card["children_environments"] if child["name"] == "chromium")
+    path = protocol_v1.find_extension_method(app, BROWSER, "cdp")["endpoint"]
+    async with httpx.AsyncClient(timeout=180) as http:
+        answer = await http.post(f"{gateway}/svc/web{path}", json={"user": user}, headers=deployed.control_headers)
+    assert answer.status_code == 200, answer.text
+    return answer.json()
+
+
+async def _connections(gateway: str, at_least: int, timeout_s: float = 30) -> list[dict]:
+    """The gateway's browser_connection events, once there are ``at_least``: each is written as its connection ends."""
+    deadline = asyncio.get_running_loop().time() + timeout_s
+    async with httpx.AsyncClient(timeout=10) as http:
+        while True:
+            lines = (await http.get(f"{gateway}/trajectory")).text.splitlines()
+            events = [event for event in map(json.loads, lines) if event.get("event_type") == "browser_connection"]
+            if len(events) >= at_least or asyncio.get_running_loop().time() > deadline:
+                return events
+            await asyncio.sleep(1)
+
+
+async def _version(ws_url: str):
+    """A connection to ``ws_url`` and the browser's product, which it answered ``Browser.getVersion`` with."""
+    held = await connect(ws_url, max_size=None, open_timeout=30)
+    await held.send(json.dumps({"id": 1, "method": "Browser.getVersion"}))
+    reply = json.loads(await asyncio.wait_for(held.recv(), 30))
+    return held, reply["result"]["product"]
+
+
+async def _closed(held) -> int | None:
+    """The code a connection was closed with, once it is."""
+    await asyncio.wait_for(held.wait_closed(), 60)
+    return held.close_code

tst/integration/computer/conftest.py

@@ -1,10 +1,12 @@
 """The ubuntu-24.04 base computer, built from its folder once for the tests here and tagged for the local registry; the
 shipped desktop example's computer env, the cua-driver test bundle's, the chromium test bundle's, the website-browser
 test bundle's and agentenv's chromium deployed alone, each composed once for a module's tests;
-``state``, local stores for a test that runs a bundle; and ``docker``, which runs the docker CLI for them."""
+``state``, local stores for a test that runs a bundle; ``docker``, which runs the docker CLI for them; and
+``cdp_clients``, Playwright and Puppeteer installed to connect to a CDP address."""
 
 import logging
 import re
+import shutil
 import subprocess
 import uuid
 from pathlib import Path
@@ -29,6 +31,8 @@
 CHROMIUM = Path(__file__).parent / "chromium"
 WEBSITE_BROWSER = Path(__file__).parent / "website-browser"
 BROWSER_COMPUTER = Path(__file__).parent / "browser-computer"
+#: Playwright and Puppeteer, which connect to a cdp address from Node.
+CDP_CLIENTS = Path(__file__).parent / "cdp_clients"
 CHROMIUM_ALONE = "@local/agentenv-framework/browsers/chromium"
 
 
@@ -77,6 +81,22 @@ def browser_computer(tmp_path_factory):
     yield from _composed(tmp_path_factory.mktemp("browser-computer"), parse_bundle(BROWSER_COMPUTER), CHROMIUM_ALONE)
 
 
+@pytest.fixture(scope="module")
+def cdp_clients(tmp_path_factory) -> Path:
+    """The CDP clients installed in a folder of their own, or the test skipped without Node 20 or later and npm."""
+    node, npm = shutil.which("node"), shutil.which("npm")
+    version = subprocess.run([node, "--version"], capture_output=True, text=True).stdout if node else ""
+    if not npm or not (match := re.match(r"v(\d+)", version)) or int(match.group(1)) < 20:
+        pytest.skip("the CDP clients need Node 20 or later and npm")
+    folder = tmp_path_factory.mktemp("cdp-clients")
+    for name in ("package.json", "connect.mjs"):
+        shutil.copy(CDP_CLIENTS / name, folder / name)
+    installed = subprocess.run([npm, "install", "--no-audit", "--no-fund", "--loglevel=error"], cwd=folder,
+                               capture_output=True, text=True, timeout=600)
+    assert installed.returncode == 0, installed.stderr
+    return folder
+
+
 def _composed(root: Path, bundle, env_id: str | None = None):
     patch = pytest.MonkeyPatch()
     patch.chdir(root)