feat(gateway): WebSockets on a site's route, a site served at a front's own root, and the provider's headers kept from sites (FD-4683)

packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py

@@ -21,6 +21,7 @@
     DEFAULT_HEALTH_PATH,
     READY_TIMEOUT_S,
     SITE_READY_TIMEOUT_S,
+    WEBSITE_BROWSE_URL_PARAM,
     WEBSITE_EXTENSION_URI,
     WEBSITE_READY_OP,
     WebsiteReadiness,
@@ -36,6 +37,6 @@
     "UniverseClock", "ClockUnavailable", "SessionSigner", "reserved_path",
     "WEBSITE_EXTENSION_URI", "CLOCK_EXTENSION_URI", "UNIVERSE_NOW_HEADER", "DEFAULT_HEALTH_PATH", "SESSION_PATH",
     "SESSION_COOKIE", "SESSION_SECRET_ENV", "ROLE_KEY_ENV", "ROLE_KEY_HEADER",
-    "WEBSITE_READY_OP", "READY_TIMEOUT_S", "SITE_READY_TIMEOUT_S", "WebsiteReadiness",
+    "WEBSITE_READY_OP", "WEBSITE_BROWSE_URL_PARAM", "READY_TIMEOUT_S", "SITE_READY_TIMEOUT_S", "WebsiteReadiness",
     "website_readiness", "env_readiness",
 ]

packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py

@@ -4,7 +4,9 @@
 of an env, a composed card's children or a single site's own card, and answers for the env as a whole; asking one
 site may take ``READY_TIMEOUT_S`` by default. A site may say how long it may take to be ready (``timeout_s``,
 ``SITE_READY_TIMEOUT_S`` when it says nothing). A server whose card has no ``ready`` is asked at its health path,
-``DEFAULT_HEALTH_PATH`` unless it names one, which is deprecated."""
+``DEFAULT_HEALTH_PATH`` unless it names one, which is deprecated. A composer serving a site to a browser at a route of
+its own, as the env gateway serves ``/website/<name>/``, says where on the site's ``website/v1`` params, as
+``WEBSITE_BROWSE_URL_PARAM``: a path joined onto the composed card's address, or a full URL."""
 from __future__ import annotations
 
 import asyncio
@@ -17,6 +19,7 @@
 
 WEBSITE_EXTENSION_URI = "urn:agentenv:website/v1"
 WEBSITE_READY_OP = "ready"
+WEBSITE_BROWSE_URL_PARAM = "browse_url"
 READY_TIMEOUT_S = 10.0
 SITE_READY_TIMEOUT_S = 300.0
 DEFAULT_HEALTH_PATH = "/api/health"

pyproject.toml

@@ -159,6 +159,7 @@ dev = [
     "tomli-w==1.2.0",  # the Codex agent's config writer, which only its image installs; its tests import it
     "moto>=5.0.0",
     "psycopg2-binary>=2.9.0",  # the gateway server module's driver; outside its container only the tests import it
+    "websockets>=13",  # the gateway's WebSocket relay to a website; outside its container only the tests import it
     "pytest>=9.0",
     "pytest-asyncio>=0.23.0",
     "pytest-dependency>=0.6.0",

src/agent_env/env/gateway/browser_proxy.py

@@ -28,7 +28,7 @@
 
 # Relative import — gateway runs as a standalone container with no agent_env.*
 from .constants import EXT_BROWSER_URI, EXT_COMPUTER_APP_URI
-from .rest_proxy import HOP_BY_HOP, fetch_child_card, operations, upstream_headers, upstream_url
+from .rest_proxy import client_headers, fetch_child_card, operations, upstream_headers, upstream_url
 
 logger = logging.getLogger(__name__)
 
@@ -125,12 +125,14 @@ def filter(self, record: logging.LogRecord) -> bool:
 class BrowserRoutes:
     """The ``/browser/<key>/cdp/<token>/`` route over a gateway's ``/svc`` children, and the rewrite of a browser's
     ``cdp`` answer on their ``/svc`` route. ``log_event`` writes a trajectory event, ``attribution`` names a role's
-    caller for one, and ``role_of`` reads the role a connection's ``AgentEnv-Role`` header names."""
+    caller for one, and ``role_of`` reads the role a connection's ``AgentEnv-Role`` header names. The headers named in
+    ``withheld`` (lower case), the sandbox provider's credentials, reach no browser."""
 
     def __init__(self, rest_proxy_urls: dict[str, str], timeout: httpx.Timeout,
                  log_event: Callable[[dict], Awaitable[Any]], attribution: Callable[[str], dict],
-                 role_of: Callable[[Request | WebSocket], str]):
+                 role_of: Callable[[Request | WebSocket], str], *, withheld: frozenset[bytes] = frozenset()):
         self._urls = rest_proxy_urls
+        self._withheld = withheld
         self._timeout = timeout
         self._log_event = log_event
         self._attribution = attribution
@@ -168,7 +170,7 @@ async def discovery(self, request: Request) -> Response:
         base = public_base(request)
         if base is None:
             return _error(400, "the Host header isn't a host and port")
-        headers = [(name, value) for name, value in request.headers.raw if name.lower() not in HOP_BY_HOP]
+        headers = client_headers(request, self._withheld)
         try:
             async with httpx.AsyncClient(timeout=self._timeout) as client:
                 answer = await client.request(request.method, self._target(key, browser, request), headers=headers)
@@ -192,7 +194,7 @@ async def cdp_answer(self, request: Request, base_url: str, role: str) -> Respon
         try:
             async with httpx.AsyncClient(timeout=self._timeout) as client:
                 answer = await client.post(upstream_url(base_url, request), content=await request.body(),
-                                           headers=upstream_headers(request, role))
+                                           headers=upstream_headers(request, role, withheld=self._withheld))
         except httpx.HTTPError as e:
             return _error(502, f"the browser didn't answer: {type(e).__name__}")
         return _answered(answer, CDP_ANSWER_FIELDS, base, key)
@@ -210,8 +212,9 @@ async def relay(self, websocket: WebSocket) -> None:
             await self._log_event({**event, "opened_at": None, "closed_at": _now(), "close_code": REFUSED,
                                    "closed_by": "gateway", "status": browser.status_code})
             return
-        headers = [(name, value) for name, value in websocket.headers.items()
-                   if name.encode() not in HOP_BY_HOP and not name.startswith("sec-websocket-")]
+        headers = [(name.decode("latin-1"), value.decode("latin-1"))
+                   for name, value in client_headers(websocket, self._withheld)
+                   if not name.lower().startswith(b"sec-websocket-")]
         offered = [p.strip() for p in websocket.headers.get("sec-websocket-protocol", "").split(",") if p.strip()]
         try:
             upstream = await connect(

src/agent_env/env/gateway/constants.py

@@ -83,6 +83,7 @@ def data_plane_load_timeout_s(size_bytes: int | None) -> int:
 EXT_BROWSER_URI = "urn:agentenv:browser/v1"
 EXT_COMPUTER_APP_URI = "urn:agentenv:computer-app/v1"
 WEBSITE_READY_OP = "ready"
+WEBSITE_BROWSE_URL_PARAM = "browse_url"
 
 # Rows per GET /changes page: the default, and the most a client may ask for.
 CHANGES_PAGE_DEFAULT = 1000
@@ -145,6 +146,9 @@ def sources(names: list[str]) -> str:
 )
 
 AGENT_ENV_ROLE_HEADER = "AgentEnv-Role"
+# Sent by a computer's site proxy, which serves a website's route at the root of the site's own name.
+AGENT_ENV_SITE_ROOT_HEADER = "AgentEnv-Site-Root"
+GATEWAY_PROVIDER_HEADERS_ENV = "AGENTENV_PROVIDER_REQUEST_HEADERS"
 # MCP request `_meta` keys the gateway stamps on every child tools/call: the caller's role, and on the MCP path a
 # key naming the inbound session, so a child can keep per-caller state behind the gateway's one shared session.
 AGENT_ENV_ROLE_META_KEY = "agentenv.io/role"

src/agent_env/env/gateway/entrypoint.py

@@ -7,7 +7,7 @@
 import httpx
 
 from . import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode, InternalMCPServer
-from .constants import DEFAULT_MCP_SERVER_NAME
+from .constants import DEFAULT_MCP_SERVER_NAME, GATEWAY_PROVIDER_HEADERS_ENV
 from .gateway import Gateway
 
 logger = logging.getLogger(__name__)
@@ -152,5 +152,6 @@ def parse_name_url_map(env_str: str) -> dict[str, str]:
         rest_proxy_urls=rest_proxy_urls,
         gateway_mode=gateway_mode,
         service_db_url=service_db_url,
+        provider_headers=os.environ.get(GATEWAY_PROVIDER_HEADERS_ENV, "").split(","),
     )
     gateway.run()

src/agent_env/env/gateway/gateway.py

@@ -10,6 +10,7 @@
 import posixpath
 import socket
 import subprocess
+from collections.abc import Iterable
 from contextlib import asynccontextmanager
 from datetime import datetime, timezone
 from pathlib import Path
@@ -30,9 +31,10 @@
 from mcp.shared.exceptions import McpError
 from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent, Tool as MCPTool
 from pydantic import create_model
-from starlette.requests import Request
+from starlette.requests import HTTPConnection, Request
 from starlette.responses import Response, StreamingResponse
 from starlette.routing import WebSocketRoute
+from starlette.websockets import WebSocket
 from anyio import ClosedResourceError
 
 if TYPE_CHECKING:
@@ -50,6 +52,7 @@
     CHILD_DISCOVERY_TIMEOUT_S,
     DATA_PLANE_LOAD_TIMEOUT_MAX_S,
     DEFAULT_ROLE,
+    EXT_WEBSITE_URI,
     GATEWAY_COMPOSE_HOST,
     GATEWAY_EXTENSIONS,
     GATEWAY_TRAJECTORY_FILE,
@@ -63,6 +66,7 @@
     STEP_SOURCE_POST_TURN_HOOK,
     TOOL_DISABLE_ACTION,
     TOOL_ENABLE_ACTION,
+    WEBSITE_BROWSE_URL_PARAM,
     WELL_KNOWN_PATH,
     WILDCARD,
     GatewayMode,
@@ -81,6 +85,8 @@
     fetch_child_card,
     forward,
     operations,
+    refuse_websocket,
+    relay_websocket,
     rewrite_child_card,
     route_name,
     upstream_url,
@@ -137,16 +143,20 @@ def __init__(
         rest_proxy_urls: dict[str, str] | None = None,
         gateway_mode: GatewayMode = GatewayMode.PERFORMANCE,
         service_db_url: str | None = None,
+        provider_headers: Iterable[str] = (),
     ):
+        """`provider_headers` names the headers the sandbox provider in front of the gateway is sent, its credentials,
+        which the gateway takes off every request it forwards to a service behind it."""
         self.host = host
         self.port = port
         self.server_name = server_name
         self.internal_mcp_servers = internal_mcp_servers
         self.website_urls = website_urls or {}
         self.rest_proxy_urls = rest_proxy_urls or {}
         self._site_reserved: dict[str, frozenset[str]] = {}
+        self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip())
         self._browsers = BrowserRoutes(self.rest_proxy_urls, self.REST_PROXY_TIMEOUT, self._log_event,
-                                       self._attribution, self._request_role)
+                                       self._attribution, self._request_role, withheld=self._withheld)
         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):
@@ -220,18 +230,23 @@ async def browser_discovery(request: Request) -> Response:
         if self.website_urls:
             @self._mcp.custom_route(
                 "/website/{service_name}/{path:path}",
-                methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+                methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
             )
             async def proxy_to_frontend(request: Request) -> Response:
                 return await self._proxy_website_request(request)
 
             @self._mcp.custom_route(
                 "/website/{service_name}",
-                methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
+                methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS"],
             )
             async def proxy_to_frontend_root(request: Request) -> Response:
                 return await self._proxy_website_request(request)
 
+            self._mcp._custom_starlette_routes += [
+                WebSocketRoute("/website/{service_name}/{path:path}", self._proxy_website_socket),
+                WebSocketRoute("/website/{service_name}", self._proxy_website_socket),
+            ]
+
         @self._mcp.custom_route(f"/tools/{TOOL_DISABLE_ACTION}", methods=["POST"])
         async def tools_disable(request: Request) -> Response:
             return await self._handle_role_mutation(request, action=TOOL_DISABLE_ACTION)
@@ -640,7 +655,8 @@ async def _proxy_rest_request(self, request: Request, url_map: dict[str, str] |
         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)
+            return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server,
+                                 withheld=self._withheld)
         changelog_before = self._query_changelog_id()
         attribution = self._attribution(role)  # at the start: a fire may be gone by the time the response ends
 
@@ -665,7 +681,7 @@ async def log_write(status: int) -> None:
             })
 
         return await forward(request, base_url, role=role, timeout=self.REST_PROXY_TIMEOUT, mcp_server=mcp_server,
-                             on_done=log_write)
+                             on_done=log_write, withheld=self._withheld)
 
     async def _proxy_website_request(self, request: Request) -> Response:
         """Proxy a website request; with a website or state trigger registered, feed it to the engine and hold barriers.
@@ -692,6 +708,25 @@ async def _proxy_website_request(self, request: Request) -> Response:
             await self._await_barriers(barrier_tasks, barrier_timeout_s)
         return response
 
+    async def _proxy_website_socket(self, websocket: WebSocket) -> None:
+        """Relay a WebSocket on a site's route to its server, as its HTTP requests are: its data plane and card, and
+        every path its card advertises for agentenv (``site_reserved_paths``), are refused, and it goes upstream with
+        the role, role key and forwarded address an HTTP request would carry."""
+        reserved = await self._site_reserved_paths(websocket.path_params["service_name"])
+        if reserved is None:
+            await refuse_websocket(websocket, self._json_response({"error": "website card unavailable"}, 503))
+            return
+        if reserved_site_request(websocket, reserved):
+            await refuse_websocket(websocket, self._json_response({"error": "not found"}, 404))
+            return
+        service_name = websocket.path_params["service_name"]
+        base_url = self.website_urls.get(service_name)
+        if base_url is None:
+            await refuse_websocket(websocket, self._json_response({"error": f"Unknown service: {service_name}"}, 404))
+            return
+        await relay_websocket(websocket, base_url, role=self._request_role(websocket),
+                              open_timeout=self.REST_PROXY_CONNECT_TIMEOUT_S, withheld=self._withheld)
+
     async def _fire_website_triggers(self, request: Request, app: str,
                                      status: int) -> tuple[list[asyncio.Task], float | None]:
         """Feed an answered website request to the trigger engine (fail-open)."""
@@ -706,7 +741,7 @@ async def _fire_website_triggers(self, request: Request, app: str,
             return [], None
 
     @staticmethod
-    def _request_role(request: Request) -> str:
+    def _request_role(request: HTTPConnection) -> str:
         """The role the gateway attributes a REST request to: the first AgentEnv-Role header, else the default role."""
         return request.headers.get(AGENT_ENV_ROLE_HEADER, "").strip() or DEFAULT_ROLE
 
@@ -771,7 +806,15 @@ def _backing_keys(self) -> list[str]:
         return list(self.rest_proxy_urls)
 
     def _rewrite_child_card(self, key: str, card: dict) -> dict:
-        return rewrite_child_card(key, card, self._gateway_extensions())
+        """One backing server's card rooted at the gateway (``rewrite_child_card``); a website's ``website/v1`` params
+        also name the route the gateway serves it to a browser at, ``/website/<key>/``, as ``browse_url``."""
+        card = rewrite_child_card(key, card, self._gateway_extensions())
+        if key in self.website_urls:
+            for ext in (card.get("capabilities") or {}).get("extensions") or []:
+                if isinstance(ext, dict) and ext.get("uri") == EXT_WEBSITE_URI:
+                    params = ext.get("params") if isinstance(ext.get("params"), dict) else {}
+                    ext["params"] = {**params, WEBSITE_BROWSE_URL_PARAM: f"/website/{key}/"}
+        return card
 
     async def _fetch_child_card(self, client: httpx.AsyncClient, key: str, base_url: str) -> dict | None:
         """Fetch + rewrite one backing server's card; None (omit) on 404 / error / timeout / malformed JSON. A
@@ -1638,7 +1681,7 @@ def site_reserved_paths(card: dict) -> frozenset[str]:
     return frozenset(_normalized(p) for p in paths if isinstance(p, str) and p.startswith("/") and _normalized(p) != "/")
 
 
-def site_request_path(request: Request) -> str | None:
+def site_request_path(request: HTTPConnection) -> str | None:
     """The path a site's server receives for a request on its route: the raw path forwarded, its dot segments removed
     by httpx and then decoded. It differs from the path as routed here when an encoded slash sits beside a dot segment
     (``public%2Fx/../api`` reaches ``/api``). None for a request-target the proxy can't build a URL for: ``forward()``
@@ -1649,10 +1692,11 @@ def site_request_path(request: Request) -> str | None:
         return None
 
 
-def reserved_site_request(request: Request, reserved: frozenset[str] = frozenset()) -> bool:
-    """Whether a request on a site's route reaches its data plane or env card, or one of the paths ``reserved`` names
-    (``site_reserved_paths``): by the path as routed here, or by the one its server receives (``site_request_path``).
-    A request-target the proxy can't build a URL for isn't reserved here: ``forward()`` refuses it."""
+def reserved_site_request(request: HTTPConnection, reserved: frozenset[str] = frozenset()) -> bool:
+    """Whether a request or WebSocket on a site's route reaches its data plane or env card, or one of the paths
+    ``reserved`` names (``site_reserved_paths``): by the path as routed here, or by the one its server receives
+    (``site_request_path``). A request-target the proxy can't build a URL for isn't reserved here: ``forward()``
+    refuses it."""
     if reserved_site_path("/" + request.path_params.get("path", ""), reserved):
         return True
     sent = site_request_path(request)