From 8ed27e6d7aeb04d4df7df1fe322cd916ff55baca Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 09:47:42 -0800 Subject: [PATCH 1/3] feat(gateway): WebSockets on a site's route, a front that serves a site at its own root, and the provider's headers kept from sites The gateway side of a computer reaching the run's websites by their names, lifted from #448: - A WebSocket to /website// is relayed to the site, under the HTTP route's rules: refused at the paths the site's card reserves, with the role and forwarded address the gateway writes. - AgentEnv-Site-Root: a front that serves a site's route at a root of its own, as a computer's site proxy does, gets no X-Forwarded-Prefix. - The headers the sandbox provider in front of the gateway is sent, its credentials, are taken off every request the gateway forwards (AGENTENV_PROVIDER_REQUEST_HEADERS). - website/v1 browse_url: where a composer serves a site to a browser. Co-Authored-By: William Loo Co-Authored-By: Claude Opus 5.5 --- .../src/agentenv_protocol/website/__init__.py | 3 +- .../agentenv_protocol/website/readiness.py | 5 +- pyproject.toml | 1 + src/agent_env/env/gateway/constants.py | 4 + src/agent_env/env/gateway/entrypoint.py | 3 +- src/agent_env/env/gateway/gateway.py | 68 +++- src/agent_env/env/gateway/rest_proxy.py | 168 +++++++++- .../env_providers/env_gateway_provider.py | 9 +- .../sandbox_providers/sandbox_provider.py | 9 + tst/unit/env/gateway/gateway_card_test.py | 22 ++ .../env/gateway/rest_proxy_stream_test.py | 14 + .../gateway/website_reserved_paths_test.py | 20 ++ .../env/gateway/website_websocket_test.py | 295 ++++++++++++++++++ .../env_gateway_provider_test.py | 61 +++- .../test_reachability_contributions.py | 18 ++ uv.lock | 2 + 16 files changed, 668 insertions(+), 34 deletions(-) create mode 100644 tst/unit/env/gateway/website_websocket_test.py diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py b/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py index c6bcb29fa..48d678519 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/__init__.py +++ b/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", ] diff --git a/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py b/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py index 756938d36..91764f46c 100644 --- a/packages/agentenv-protocol/src/agentenv_protocol/website/readiness.py +++ b/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//``, 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" diff --git a/pyproject.toml b/pyproject.toml index 3751af021..dc4d207fc 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/agent_env/env/gateway/constants.py b/src/agent_env/env/gateway/constants.py index eb37d72bd..a04bd51f2 100644 --- a/src/agent_env/env/gateway/constants.py +++ b/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" diff --git a/src/agent_env/env/gateway/entrypoint.py b/src/agent_env/env/gateway/entrypoint.py index b0b936ae9..65cf7e3c3 100644 --- a/src/agent_env/env/gateway/entrypoint.py +++ b/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() diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index 14784056b..fb3248265 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/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,7 +143,10 @@ 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 @@ -153,6 +162,7 @@ def __init__( log.addFilter(TokenRedaction()) self.gateway_mode = gateway_mode self._service_db_url = service_db_url + self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: @@ -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//``, 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) diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index 33dc34570..e9fca4336 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -13,9 +13,15 @@ inside its JSON-RPC ``params._meta``; - the upstream learns the address the client reached it at: ``X-Forwarded-Host`` and ``X-Forwarded-Proto`` (a front's own, when one before the gateway sent them, else the request's ``Host`` and scheme) and ``X-Forwarded-Prefix``, the - route prefix (``/svc/{service}`` or ``/website/{service}``) the gateway strips, which a client never sets; + route prefix (``/svc/{service}`` or ``/website/{service}``) the gateway strips, which a client never sets: a client + that serves the route at a root of its own (a computer's site proxy, at ``https://notes.test/``) says so with + ``AgentEnv-Site-Root``, and the upstream is then told no prefix; - the upstream response and the client it came through are closed when the response ends, iterated or not. +A WebSocket on ``/website/{service}/{path}`` is relayed under the same rules (:func:`relay_websocket`): the gateway +opens one of its own to the service, with the headers an HTTP request would carry upstream, and relays each message +and the close, either side's ending ending both. + A backing server's env card names its endpoints relative to itself; :func:`rewrite_child_card` roots them under ``/svc/{key}`` so they resolve against the env's address, keeping the gateway's own routes where the gateway offers the same operation. A website's card names its routes on ``urn:agentenv:website/v1``; :func:`website_routes` reads @@ -23,6 +29,8 @@ """ from __future__ import annotations +import asyncio +import contextlib import copy import json import logging @@ -33,15 +41,19 @@ from urllib.parse import unquote import httpx -from starlette.requests import Request +from starlette.requests import HTTPConnection, Request from starlette.responses import Response, StreamingResponse from starlette.types import Receive, Scope, Send +from starlette.websockets import WebSocket, WebSocketState +from websockets.asyncio.client import ClientConnection, connect as connect_websocket +from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus # Relative import — gateway runs as a standalone container with no agent_env.* from .constants import ( AGENT_ENV_ROLE_HEADER, AGENT_ENV_ROLE_META_KEY, AGENT_ENV_SESSION_META_KEY, + AGENT_ENV_SITE_ROOT_HEADER, CARD_FETCH_TIMEOUT_S, EXT_WEBSITE_URI, WELL_KNOWN_PATH, @@ -59,12 +71,24 @@ ROLE_HEADER_BYTES = AGENT_ENV_ROLE_HEADER.lower().encode() FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX = "x-forwarded-host", "x-forwarded-proto", "x-forwarded-prefix" -_FORWARDED = frozenset(name.encode() for name in (FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX)) +_FORWARDED = frozenset(name.encode() for name in (FORWARDED_HOST, FORWARDED_PROTO, FORWARDED_PREFIX, + AGENT_ENV_SITE_ROOT_HEADER.lower())) +# A WebSocket's handshake fields, which the gateway's own connection upstream writes for itself. +_WEBSOCKET_HANDSHAKE = frozenset({"sec-websocket-key", "sec-websocket-version", "sec-websocket-extensions", + "sec-websocket-protocol"}) +# What the gateway's own WebSocket handshake answer writes for itself, so a site's copies aren't relayed. +_WEBSOCKET_ANSWER_DROP = HOP_BY_HOP | {b"date", b"server", b"content-length", b"sec-websocket-accept", + b"sec-websocket-protocol", b"sec-websocket-extensions"} +_WEBSOCKET_REFUSAL_DROP = HOP_BY_HOP | {b"date", b"server", b"content-length"} +_HTTP_SCHEME = {"ws": "http", "wss": "https"} +WEBSOCKET_MESSAGE_MAX_BYTES = 16 * 1024 * 1024 async def forward(request: Request, base_url: str, *, role: str, timeout: httpx.Timeout, mcp_server: bool, - on_done: Callable[[int], Awaitable[None]] | None = None) -> Response: - """Forward `request` to the service at `base_url` as `role` and stream the response through. `mcp_server` says + on_done: Callable[[int], Awaitable[None]] | None = None, + withheld: frozenset[bytes] = frozenset()) -> Response: + """Forward `request` to the service at `base_url` as `role`, without the headers named in `withheld`, and stream + the response through. `mcp_server` says the service is a backing MCP server, whose own endpoint takes the role in its JSON-RPC `_meta` as well as in the header. A 400 names a request-target the proxy cannot route; a 502 names an upstream that never answered. `on_done` is awaited with the status once the response has ended, after its last byte when streamed.""" @@ -73,7 +97,7 @@ async def forward(request: Request, base_url: str, *, role: str, timeout: httpx. except ValueError as e: # an escaped slash inside the route prefix; refuse rather than forward a wrong path return await _done(_json_response({"error": str(e)}, 400), on_done) content = await request.body() # /svc and /website request bodies are small (JSON-RPC, forms) - headers = upstream_headers(request, role) + headers = upstream_headers(request, role, withheld=withheld) if mcp_server and request.method == "POST" and is_child_mcp_path(request.path_params.get("path", "")): # The gateway is the trust boundary for identity: a child's proxied MCP endpoint stays reachable, but the # role in its JSON-RPC `_meta` is the gateway's to write, like the header above. @@ -96,7 +120,7 @@ async def _done(response: Response, on_done: Callable[[int], Awaitable[None]] | return response -def upstream_url(base_url: str, request: Request) -> str: +def upstream_url(base_url: str, request: HTTPConnection) -> str: # Forward the path as received: uvicorn decodes scope["path"] for routing, so Starlette's {path} param # would send %2F, %3F and %23 to the upstream as separators. latin-1 keeps one char per wire byte. url = base_url @@ -110,14 +134,15 @@ def upstream_url(base_url: str, request: Request) -> str: return url -def upstream_headers(request: Request, role: str) -> httpx.Headers: +def upstream_headers(request: HTTPConnection, role: str, *, withheld: frozenset[bytes] = frozenset()) -> httpx.Headers: """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity. The role header is the gateway's to set: the client's copies are dropped and the role the gateway attributed the request to goes in their place, so the upstream never reads a value the gateway did not write. So are the - forwarded-address headers (`forwarded_address`).""" + forwarded-address headers (`forwarded_address`). The headers named in `withheld` (lower case), the credentials of + the provider in front of the gateway, are dropped, so no service behind it reads them.""" headers = httpx.Headers([(k, v) for k, v in request.headers.raw if k.lower() not in HOP_BY_HOP and k.lower() != ROLE_HEADER_BYTES - and k.lower() not in _FORWARDED]) + and k.lower() not in _FORWARDED and k.lower() not in withheld]) headers[AGENT_ENV_ROLE_HEADER] = role headers.update(forwarded_address(request)) # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take. @@ -125,15 +150,20 @@ def upstream_headers(request: Request, role: str) -> httpx.Headers: return headers -def forwarded_address(request: Request) -> dict[str, str]: +def forwarded_address(request: HTTPConnection) -> dict[str, str]: """Where the client reached the proxied service: the host and scheme it addressed (as a front before the gateway names them in its own ``X-Forwarded-Host``/``X-Forwarded-Proto``, else the request's ``Host`` and scheme; the - client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent. - A header with nothing to say is left out.""" + client could set ``Host`` anyway, so believing its front names nothing it couldn't), and the route prefix as sent, + or none for a client that says with ``AgentEnv-Site-Root`` it serves the route at a root of its own: a computer's + site proxy, whose browser would otherwise get the route's prefix twice. Like the host, that only changes the links + the client's own answer names; what the route reaches is the path's alone. A header with nothing to say is left + out.""" host = _first(request.headers.get(FORWARDED_HOST, "")) or request.headers.get("host", "") proto = _first(request.headers.get(FORWARDED_PROTO, "")).lower() - address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else request.url.scheme, - FORWARDED_PREFIX: route_prefix(request)} + at_root = bool(request.headers.get(AGENT_ENV_SITE_ROOT_HEADER, "").strip()) + scheme = _HTTP_SCHEME.get(request.url.scheme, request.url.scheme) + address = {FORWARDED_HOST: host, FORWARDED_PROTO: proto if proto in ("http", "https") else scheme, + FORWARDED_PREFIX: "" if at_root else route_prefix(request)} return {name: value for name, value in address.items() if value} @@ -141,7 +171,7 @@ def _first(value: str) -> str: return value.split(",")[0].strip() -def route_prefix(request: Request) -> str: +def route_prefix(request: HTTPConnection) -> str: """The path the client reached the proxied service under, as it sent it: the request-target up to the `{path}` remainder, without its trailing slash, e.g. ``/website/notes``, or under a root path ``/env/website/notes``.""" raw_path = request.scope.get("raw_path") or request.scope["path"].encode() @@ -200,6 +230,112 @@ def raw_remainder(raw_path: bytes, decoded_path: str, decoded_rest: str) -> byte return rest +async def relay_websocket(websocket: WebSocket, base_url: str, *, role: str, + open_timeout: float = 10.0, withheld: frozenset[bytes] = frozenset()) -> None: + """Relay `websocket` to the service at `base_url` as `role`: a WebSocket of the gateway's own is opened to the same + path, with the headers ``upstream_headers`` sends, asking for the subprotocols the client offered; `websocket` is + accepted with the one the service chose and the headers its answer carried (a ``Set-Cookie`` among them, each + field kept apart), but those the gateway's own answer writes. Each text and binary message is relayed as it came, and either side's + close, with its code, closes the other; a message from the service larger than ``WEBSOCKET_MESSAGE_MAX_BYTES``, the + size uvicorn takes from a client, closes both with 1009 rather than be held whole. A service that refuses the + handshake has its answer relayed as the refusal, its headers as an accepted one's are; one that never answers is a + 502. The headers named in `withheld` don't go upstream, as with ``upstream_headers``.""" + try: + target_url = upstream_url(base_url, websocket) + except ValueError as e: + await refuse_websocket(websocket, _json_response({"error": str(e)}, 400)) + return + headers = [(key, value) for key, value in upstream_headers(websocket, role, withheld=withheld).multi_items() + if key.lower() not in _WEBSOCKET_HANDSHAKE] + offered = [protocol.strip() for value in websocket.headers.getlist("sec-websocket-protocol") + for protocol in value.split(",") if protocol.strip()] + try: + upstream = await connect_websocket("ws" + target_url.removeprefix("http"), additional_headers=headers, + subprotocols=offered or None, user_agent_header=None, + open_timeout=open_timeout, max_size=WEBSOCKET_MESSAGE_MAX_BYTES) + except InvalidStatus as e: + refusal = Response(content=bytes(e.response.body or b""), status_code=e.response.status_code) + refusal.raw_headers = [*(item for item in refusal.raw_headers if item[0] == b"content-length"), + *_answer_headers(e.response.headers, drop=_WEBSOCKET_REFUSAL_DROP)] + await refuse_websocket(websocket, refusal) + return + except (InvalidHandshake, OSError, TimeoutError) as e: + logger.error(f"WebSocket proxy to {target_url} failed: {e!r}") + await refuse_websocket(websocket, _json_response({"error": f"Proxy request failed: {e!r}"}, 502)) + return + async with upstream: + await websocket.accept(subprotocol=upstream.subprotocol, + headers=_answer_headers(upstream.response.headers, drop=_WEBSOCKET_ANSWER_DROP)) + await _relay_messages(websocket, upstream) + + +def _answer_headers(headers: Any, *, drop: frozenset[bytes]) -> list[tuple[bytes, bytes]]: + """A site's handshake answer's headers as the gateway relays them, each field apart (a ``Set-Cookie`` per cookie), + but those named in `drop`, which the gateway's own answer writes.""" + return [(key.lower(), value) for key, value in + ((key.encode("ascii", "surrogateescape"), value.encode("ascii", "surrogateescape")) + for key, value in headers.raw_items()) if key.lower() not in drop] + + +async def refuse_websocket(websocket: WebSocket, response: Response) -> None: + """Refuse `websocket`'s handshake with `response`, or, on a server that can't send one, a 403.""" + if "websocket.http.response" in websocket.scope.get("extensions", {}): + await websocket.send_denial_response(response) + else: + await websocket.close(1008) + + +async def _relay_messages(websocket: WebSocket, upstream: ClientConnection) -> None: + """Both directions of an accepted WebSocket until either ends, then both closed: the client's close goes upstream + with its code, and the service's to the client.""" + + async def from_client() -> None: + while True: + message = await websocket.receive() + if message["type"] == "websocket.disconnect": + await upstream.close(_sendable(message.get("code"), abnormal=1001), message.get("reason") or "") + return + if message.get("text") is not None: + await upstream.send(message["text"]) + elif message.get("bytes") is not None: + await upstream.send(message["bytes"]) + + async def from_upstream() -> None: + async for data in upstream: + if isinstance(data, str): + await websocket.send_text(data) + else: + await websocket.send_bytes(data) + + tasks = [asyncio.create_task(from_client()), asyncio.create_task(from_upstream())] + try: + await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await upstream.close() + if WebSocketState.DISCONNECTED not in (websocket.client_state, websocket.application_state): + sent = upstream.protocol.close_sent + code = sent.code if sent is not None and sent.code == 1009 else upstream.close_code + with contextlib.suppress(RuntimeError, OSError, ConnectionClosed): + await websocket.close(_sendable(code, abnormal=1011), upstream.close_reason or "") + + +def _sendable(code: int | None, *, abnormal: int) -> int: + """A close code one side reported, as one the other may be sent: none (1005) as a normal close, and a closure + with no close frame (1006, 1015) or a code no endpoint may send as `abnormal`. + + >>> [_sendable(c, abnormal=1011) for c in (None, 1005, 1000, 1006, 4001, 1004)] + [1000, 1000, 1000, 1011, 4001, 1011] + """ + if code in (None, 1005): + return 1000 + if 1000 <= code <= 1003 or 1007 <= code <= 1014 or 3000 <= code <= 4999: + return code + return abnormal + + class UpstreamResponse(StreamingResponse): """Streams an httpx response through untouched and closes it when the response ends, iterated or not. diff --git a/src/agent_env/providers/env_providers/env_gateway_provider.py b/src/agent_env/providers/env_providers/env_gateway_provider.py index c6da7c94b..b8df4c430 100644 --- a/src/agent_env/providers/env_providers/env_gateway_provider.py +++ b/src/agent_env/providers/env_providers/env_gateway_provider.py @@ -31,6 +31,7 @@ from agent_env.env.gateway import AGENT_ENV_GATEWAY_MCP_PORT, GatewayMode from agent_env.env.gateway.constants import ( DEFAULT_MCP_SERVER_NAME, + GATEWAY_PROVIDER_HEADERS_ENV, STATE_TOOL_NAME_CLASHES_KEY, WELL_KNOWN_PATH, random_mcp_server_name, @@ -68,6 +69,7 @@ SANDBOX_MODE_VM, Accepts, SandboxProvider, + sandbox_request_header_names, ) from agent_env.store.image_store.oci_registry_credentials import image_reference_problem @@ -498,6 +500,8 @@ def create_docker_compose( gateway_env.append(f" - REST_PROXY_URLS={','.join(rest_proxy_urls)}") gateway_env.append(f" - GATEWAY_MODE={gateway_mode.value}") gateway_env.append(f" - SERVICE_DB_URL={state_provider.base_url(state_instance)}") + if provider_headers := sandbox_request_header_names(): + gateway_env.append(f" - {GATEWAY_PROVIDER_HEADERS_ENV}={','.join(provider_headers)}") publish = host_port or (lambda port: port) exposed_gateway_ports = [] if expose_gateway_port: @@ -912,7 +916,9 @@ async def _stand_up(_spec): card = await self._wait_for_tunnel(gateway_url, timeout=120) if card is None: raise RuntimeError("Tunnel did not become ready in time") - await _wait_for_websites(gateway_url, website_configs or []) + if website_configs: + await _wait_for_websites(gateway_url, website_configs) + card = await self._wait_for_tunnel(gateway_url, timeout=120) or card card_read_at_utc = datetime.now(timezone.utc).isoformat() await self._refuse_tool_name_clashes(gateway_url) mcp_url = _mcp_url(gateway_url, card) @@ -1046,6 +1052,7 @@ async def _deploy_via_containers( "GATEWAY_MODE": gateway_mode.value, "SERVICE_DB_URL": self._state_provider.base_url(state_instance), "MCP_SERVER_NAME": mcp_server_name or DEFAULT_MCP_SERVER_NAME, + **({GATEWAY_PROVIDER_HEADERS_ENV: ",".join(names)} if (names := sandbox_request_header_names()) else {}), } gateway_sb = await sandbox_provider.create_container( image_name=gateway_env.docker_image_artifact.image_name, diff --git a/src/agent_env/providers/sandbox_providers/sandbox_provider.py b/src/agent_env/providers/sandbox_providers/sandbox_provider.py index b5c228ab4..3e0f47803 100644 --- a/src/agent_env/providers/sandbox_providers/sandbox_provider.py +++ b/src/agent_env/providers/sandbox_providers/sandbox_provider.py @@ -629,6 +629,15 @@ def sandbox_request_headers_for_url(url: str) -> dict[str, str]: return headers +def sandbox_request_header_names() -> tuple[str, ...]: + """The names of every header any registered provider attaches to requests for its URLs + (``sandbox_request_headers_for_url``), lower-cased and sorted: what a gateway takes off the requests it forwards, + so no service behind it reads a provider's credentials. Names only, so no credential is resolved.""" + names = {header.lower() for cls in registered_sandbox_provider_classes() + for entries in (cls.REQUEST_HEADERS or {}).values() for header in entries} + return tuple(sorted(names)) + + def _resolve_container_env(provider_cls: type[SandboxProvider]) -> dict[str, str]: """``provider_cls.CONTAINER_ENV`` resolved to values. diff --git a/tst/unit/env/gateway/gateway_card_test.py b/tst/unit/env/gateway/gateway_card_test.py index 6b71a6b7c..48979ddfa 100644 --- a/tst/unit/env/gateway/gateway_card_test.py +++ b/tst/unit/env/gateway/gateway_card_test.py @@ -267,6 +267,28 @@ async def fake_fetch(client, key, base_url): assert {c["url"] for c in card["children_environments"]} == {"/svc/mcp-items/agentenv", "/svc/webitems/agentenv"} +def test_a_website_childs_card_names_the_route_the_gateway_serves_it_to_a_browser_at(): + """The website/v1 param naming it mirrors the protocol's, as the gateway image has no agentenv_protocol.""" + from agentenv_protocol.website import WEBSITE_BROWSE_URL_PARAM + + gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": "http://notes:8000", "bare": "http://bare:8000"}, + rest_proxy_urls={"notes": "http://notes:8000", "bare": "http://bare:8000", "mcp-notes": "http://m"}) + site = {"name": "notes", "url": "/agentenv", "capabilities": {"extensions": [ + {"uri": "urn:agentenv:website/v1", "params": {"domains": ["notes.test"]}}]}} + + notes = gw._rewrite_child_card("notes", site) + mcp = gw._rewrite_child_card("mcp-notes", site) + bare = gw._rewrite_child_card("bare", {"name": "bare", "url": "/agentenv", "capabilities": {"extensions": []}}) + + assert gateway_constants.WEBSITE_BROWSE_URL_PARAM == WEBSITE_BROWSE_URL_PARAM + assert protocol_client.extension_params(notes, "urn:agentenv:website/v1") == { + "domains": ["notes.test"], WEBSITE_BROWSE_URL_PARAM: "/website/notes/"} + assert protocol_client.extension_params(mcp, "urn:agentenv:website/v1") == {"domains": ["notes.test"]} + assert bare["capabilities"]["extensions"] == [] + assert site["capabilities"]["extensions"][0]["params"] == {"domains": ["notes.test"]} + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["data/reset", "data/add", "data/get"]) async def test_env_data_plane_requires_single_child(method): diff --git a/tst/unit/env/gateway/rest_proxy_stream_test.py b/tst/unit/env/gateway/rest_proxy_stream_test.py index 779be614f..33378c399 100644 --- a/tst/unit/env/gateway/rest_proxy_stream_test.py +++ b/tst/unit/env/gateway/rest_proxy_stream_test.py @@ -506,6 +506,20 @@ async def test_the_upstream_learns_the_host_scheme_and_route_prefix_the_client_r ] +@pytest.mark.asyncio +async def test_a_client_serving_the_route_at_its_own_root_is_told_no_prefix_and_its_header_stays_with_the_gateway( + monkeypatch): + up = _Upstream([b"ok"]) + up.install(monkeypatch) + gw = _gateway(website_urls={"web": "http://web"}) + await _call(gw, "/website/web/api/x", headers={"host": "gw.example", "x-forwarded-host": "notes.test", + "x-forwarded-proto": "https", "agentenv-site-root": "1"}) + await _call(gw, "/website/web/api/x", headers={"host": "gw.example", "agentenv-site-root": ""}) + assert [_forwarded(r) for r in up.forwarded] == [(["notes.test"], ["https"], []), + (["gw.example"], ["http"], ["/website/web"])] + assert all("agentenv-site-root" not in r.headers for r in up.forwarded) + + @pytest.mark.asyncio async def test_a_front_s_own_forwarded_host_and_scheme_win_and_reach_the_upstream_once(monkeypatch): up = _Upstream([b"ok"]) diff --git a/tst/unit/env/gateway/website_reserved_paths_test.py b/tst/unit/env/gateway/website_reserved_paths_test.py index 372b1a2fe..306a31d76 100644 --- a/tst/unit/env/gateway/website_reserved_paths_test.py +++ b/tst/unit/env/gateway/website_reserved_paths_test.py @@ -132,3 +132,23 @@ class Site(WebsiteEnvironment): reserved = site_reserved_paths(card) assert find_extension_method(card, WEBSITE_EXTENSION_URI, WEBSITE_READY_OP)["endpoint"] in reserved assert reserved_site_path("/agentenv/ext/website/ready", reserved) and not reserved_site_path("/", reserved) + + +@pytest.mark.asyncio +async def test_a_cors_preflight_on_a_sites_route_reaches_the_site(monkeypatch): + reached = [] + + def site(request): + reached.append((request.method, request.url.path)) + if request.url.path == WELL_KNOWN_PATH: + return httpx.Response(404) + return httpx.Response(200, headers={"Access-Control-Allow-Origin": "https://app.test"}, content=_streamed()) + + transport = httpx.MockTransport(site) + monkeypatch.setattr(rest_proxy_module.httpx, "AsyncClient", lambda **kw: _RealAsyncClient(transport=transport, **kw)) + gw = _gateway(_Changelog(), website_urls={"api": "http://api-website-backend:8000"}) + + preflight = {"origin": "https://app.test", "access-control-request-method": "POST"} + assert (await _http(gw, "OPTIONS", "/website/api/items", headers=preflight))[0] == 200 + assert (await _http(gw, "OPTIONS", "/website/api", headers=preflight))[0] == 200 + assert ("OPTIONS", "/items") in reached and ("OPTIONS", "/") in reached diff --git a/tst/unit/env/gateway/website_websocket_test.py b/tst/unit/env/gateway/website_websocket_test.py new file mode 100644 index 000000000..002e7791b --- /dev/null +++ b/tst/unit/env/gateway/website_websocket_test.py @@ -0,0 +1,295 @@ +"""A WebSocket on a site's route reaches the site's server through the real gateway, under the HTTP route's rules: its +data plane and card are refused, the role and forwarded address go upstream as the gateway writes them, and each +message and either side's close is relayed.""" + +from __future__ import annotations + +import asyncio +import contextlib +import socket + +import httpx +import pytest +import pytest_asyncio +import uvicorn +from starlette.applications import Starlette +from starlette.responses import JSONResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket, WebSocketDisconnect +from websockets.asyncio.client import connect +from websockets.exceptions import ConnectionClosed, InvalidStatus + +from agent_env.env.gateway.gateway import Gateway + +pytestmark = pytest.mark.usefixtures("socket_enabled") + + +class EchoSite: + """A site whose ``/live`` echoes each message, closes with 4001 on ``bye``, and records each handshake's headers + and the code its client closed with.""" + + def __init__(self): + self.handshakes, self.closed = [], [] + self.app = Starlette(routes=[WebSocketRoute("/live", self.live), WebSocketRoute("/session", self.session), + WebSocketRoute("/{path:path}", self.live)]) + + async def session(self, websocket: WebSocket) -> None: + await websocket.accept(headers=[(b"set-cookie", b"s=2; Path=/"), (b"set-cookie", b"t=3; Path=/"), + (b"x-site", b"yes")]) + await websocket.close(1000) + + async def live(self, websocket: WebSocket) -> None: + self.handshakes.append((websocket.url.path, dict(websocket.headers))) + offered = websocket.headers.get("sec-websocket-protocol") + await websocket.accept(subprotocol="chat.v2" if offered and "chat.v2" in offered else None) + try: + while True: + message = await websocket.receive() + if message["type"] == "websocket.disconnect": + self.closed.append(message.get("code")) + return + if message.get("text") == "bye": + await websocket.close(4001, "done") + return + if message.get("text") is not None: + await websocket.send_text(message["text"].upper()) + else: + await websocket.send_bytes(message["bytes"][::-1]) + except WebSocketDisconnect as e: + self.closed.append(e.code) + + +async def serve(app) -> tuple[uvicorn.Server, int]: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("127.0.0.1", 0)) + config = uvicorn.Config(app, lifespan="off", log_level="critical", access_log=False) + config.load() + server = uvicorn.Server(config) + server.lifespan = config.lifespan_class(config) + await server.startup(sockets=[sock]) + return server, sock.getsockname()[1] + + +async def _until(condition, timeout_s: float = 5.0) -> None: + deadline = asyncio.get_running_loop().time() + timeout_s + while not condition(): + assert asyncio.get_running_loop().time() < deadline + await asyncio.sleep(0.02) + + +@pytest_asyncio.fixture +async def stack(): + site = EchoSite() + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + rest_proxy_urls={"notes": f"http://127.0.0.1:{site_port}"}, + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + yield site, f"ws://127.0.0.1:{gateway_port}/website/notes" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + +@pytest.mark.asyncio +async def test_messages_and_the_sites_close_are_relayed_with_the_gateways_headers(stack): + site, route = stack + async with connect(f"{route}/live?room=1", subprotocols=["chat.v1", "chat.v2"], + additional_headers=[("AgentEnv-Role", "grader"), ("X-Forwarded-Prefix", "/elsewhere"), + ("Cookie", "s=1")]) as websocket: + await websocket.send("hello") + text = await asyncio.wait_for(websocket.recv(), 5) + await websocket.send(b"\x01\x02\x03") + data = await asyncio.wait_for(websocket.recv(), 5) + await websocket.send("bye") + with pytest.raises(ConnectionClosed): + await asyncio.wait_for(websocket.recv(), 5) + subprotocol, close = websocket.subprotocol, websocket.close_code + + ((path, headers),) = site.handshakes + assert (text, data, subprotocol, close) == ("HELLO", b"\x03\x02\x01", "chat.v2", 4001) + assert path == "/live" and headers["cookie"] == "s=1" + assert headers["agentenv-role"] == "grader" + assert (headers["x-forwarded-proto"], headers["x-forwarded-prefix"]) == ("http", "/website/notes") + + +@pytest.mark.asyncio +async def test_the_sites_handshake_answer_reaches_the_client_with_each_cookie_apart(stack): + _, route = stack + async with connect(f"{route}/session") as websocket: + headers = websocket.response.headers + + assert headers.get_all("set-cookie") == ["s=2; Path=/", "t=3; Path=/"] and headers["x-site"] == "yes" + assert len(headers.get_all("sec-websocket-accept")) == 1 and len(headers.get_all("server")) <= 1 + + +@pytest.mark.asyncio +async def test_the_clients_close_reaches_the_site(stack): + site, route = stack + async with connect(f"{route}/live") as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + await websocket.close(4002, "leaving") + + await _until(lambda: site.closed) + assert site.closed == [4002] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/.well-known/agent-env.json", "/agentenv", "/%61gentenv", "/x/..%2Fagentenv"]) +async def test_a_websocket_to_the_sites_data_plane_or_card_is_refused(stack, path): + site, route = stack + with pytest.raises(InvalidStatus) as refused: + async with connect(f"{route}{path}"): + pass + + assert refused.value.response.status_code == 404 and site.handshakes == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/grade", "/grade/", "/x/..%2Fgrade", "/grading/now", "/agentenv/ext/website/ready"]) +async def test_a_websocket_to_a_path_the_sites_card_advertises_is_refused(path): + site = EchoSite() + card = {"name": "notes", "url": "/agentenv", "capabilities": {"extensions": [{ + "uri": "urn:agentenv:website/v1", "params": {"methods": { + "ready": {"method": "GET", "endpoint": "/agentenv/ext/website/ready"}, + "grade": {"method": "GET", "endpoint": "/grade"}}}}, + {"uri": "urn:example:grading/v1", "params": {"endpoint": "/grading"}}]}} + site.app.router.routes.insert(0, Route("/.well-known/agent-env.json", lambda request: JSONResponse(card))) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + route = f"ws://127.0.0.1:{gateway_port}/website/notes" + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"{route}{path}"): + pass + async with connect(f"{route}/live") as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 404 and [p for p, _ in site.handshakes] == ["/live"] + + +@pytest.mark.asyncio +async def test_a_message_from_the_site_over_the_limit_closes_both_sides_with_1009(stack, monkeypatch): + from agent_env.env.gateway import rest_proxy + + site, route = stack + monkeypatch.setattr(rest_proxy, "WEBSOCKET_MESSAGE_MAX_BYTES", 1024) + + async def big(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_bytes(b"x" * 2048) + with contextlib.suppress(WebSocketDisconnect): + while (await websocket.receive())["type"] != "websocket.disconnect": + pass + + site.app.router.routes.insert(0, WebSocketRoute("/big", big)) + async with connect(f"{route}/big") as websocket: + with pytest.raises(ConnectionClosed) as closed: + await asyncio.wait_for(websocket.recv(), 5) + + assert closed.value.rcvd.code == 1009 + + +@pytest.mark.asyncio +async def test_a_websocket_is_refused_while_the_sites_card_cant_be_read(): + site = EchoSite() + site.app.router.routes.insert(0, Route("/.well-known/agent-env.json", lambda request: JSONResponse({}, 503))) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 503 and site.handshakes == [] + + +@pytest.mark.asyncio +async def test_the_providers_headers_never_reach_the_site_over_http_or_a_websocket(): + site, seen = EchoSite(), [] + + async def echo_headers(request): + seen.append(dict(request.headers)) + return JSONResponse(dict(request.headers)) + + site.app.router.routes.insert(0, Route("/whoami", echo_headers)) + site_server, site_port = await serve(site.app) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}, + provider_headers=["CF-Access-Client-Id", "cf-access-client-secret", " "]) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + sent = {"CF-Access-Client-Id": "id", "CF-Access-Client-Secret": "secret", "X-Site": "kept"} + try: + async with httpx.AsyncClient() as client: + answer = await client.get(f"http://127.0.0.1:{gateway_port}/website/notes/whoami", headers=sent) + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live", additional_headers=sent) as websocket: + await websocket.send("ping") + assert await asyncio.wait_for(websocket.recv(), 5) == "PING" + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + [(_, handshake)] = site.handshakes + for headers in (seen[0], handshake, answer.json()): + assert "cf-access-client-id" not in headers and "cf-access-client-secret" not in headers + assert headers["x-site"] == "kept" + + +@pytest.mark.asyncio +async def test_a_refused_handshake_keeps_the_sites_headers_each_cookie_apart(): + async def refuse(websocket: WebSocket) -> None: + answer = JSONResponse({"error": "signed out"}, status_code=401) + answer.raw_headers += [(b"set-cookie", b"s=; Max-Age=0; Path=/"), (b"set-cookie", b"t=; Max-Age=0; Path=/"), + (b"x-site", b"yes")] + await websocket.send_denial_response(answer) + + site_server, site_port = await serve(Starlette(routes=[WebSocketRoute("/live", refuse)])) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._asgi_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + answer = refused.value.response + assert answer.status_code == 401 and answer.headers.get_all("set-cookie") == [ + "s=; Max-Age=0; Path=/", "t=; Max-Age=0; Path=/"] + assert answer.headers["x-site"] == "yes" and answer.headers["content-type"] == "application/json" + assert answer.headers.get_all("content-length") == [str(len(answer.body))] + + +@pytest.mark.asyncio +async def test_a_site_that_refuses_the_handshake_has_its_answer_relayed(): + site_server, site_port = await serve(Starlette(routes=[])) + gateway = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], + website_urls={"notes": f"http://127.0.0.1:{site_port}"}) + gateway_server, gateway_port = await serve(gateway._mcp.streamable_http_app()) + try: + with pytest.raises(InvalidStatus) as refused: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/notes/missing"): + pass + with pytest.raises(InvalidStatus) as unknown: + async with connect(f"ws://127.0.0.1:{gateway_port}/website/other/live"): + pass + finally: + await gateway_server.shutdown() + await site_server.shutdown() + + assert refused.value.response.status_code == 403 and unknown.value.response.status_code == 404 diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 58aad5c49..8ed04f46c 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -305,8 +305,10 @@ async def test_a_vm_gateway_refuses_an_image_its_vm_cant_get_before_creating_the @pytest.mark.asyncio -async def test_container_mode_constructs_internal_mcp_servers_url_format(): - """Container mode must pass INTERNAL_MCP_SERVERS as 'name=url,name=url' to the gateway.""" +async def test_container_mode_constructs_internal_mcp_servers_url_format(monkeypatch): + """Container mode must pass INTERNAL_MCP_SERVERS as 'name=url,name=url' to the gateway, and the names of the + providers' request headers, which it takes off what it forwards.""" + monkeypatch.setattr(env_gateway_provider, "sandbox_request_header_names", lambda: ("cf-access-client-id", "x-a")) provider = ModalSandboxProvider() gp = EnvironmentGatewayProvider() gp._state_provider = LocalPostgresStateProvider() # normally set by create_gateway @@ -395,6 +397,7 @@ async def _spy_acquire(self, ctx): assert re.search(r"^postgresql://agentenv:agentenv@\[fdaa::fake:[0-9a-f]+\]:5432/agentenv", env["SERVICE_DB_URL"]), env["SERVICE_DB_URL"] assert env["MCP_SERVER_NAME"] == "crm" assert result.mcp_server_name == "crm" + assert env["AGENTENV_PROVIDER_REQUEST_HEADERS"] == "cf-access-client-id,x-a" # Each MCP got DATABASE_URL with its scoped search_path mcp_calls = [c for c in create_calls if c["image_name"].startswith("mcp-")] @@ -1191,6 +1194,26 @@ def test_redact_compose_secrets_masks_only_secret_env_lines(): assert "s3cr3t-token-value" in compose +@pytest.mark.parametrize("names", [("cf-access-client-id", "cf-access-client-secret"), ()], ids=["named", "none"]) +def test_compose_gives_the_gateway_the_names_of_the_providers_request_headers_alone(names, monkeypatch): + from agent_env.env.envs.service_db import ServiceDBConfig + + monkeypatch.setattr(env_gateway_provider, "sandbox_request_header_names", lambda: names) + compose = EnvironmentGatewayProvider().create_docker_compose( + mcp_servers=[MCPServerConfig(image="mcp-slack", environment_name="slack")], + gateway_image="agent-gateway", + website_configs=[WebsiteConfig(backend_image="notes-be", frontend_image=None, environment_name="notes")], + state_provider=LocalPostgresStateProvider(service_db_config=ServiceDBConfig()), + state_instance=LocalPostgresStateProvider.default_instance(), + ) + + services = yaml.safe_load(compose)["services"] + named = {name: [entry for entry in service.get("environment") or [] if "PROVIDER_REQUEST_HEADERS" in entry] + for name, service in services.items()} + expected = {GATEWAY_SERVICE_NAME: [f"AGENTENV_PROVIDER_REQUEST_HEADERS={','.join(names)}"]} if names else {} + assert {name: entries for name, entries in named.items() if entries} == expected + + @pytest.mark.asyncio async def test_acquire_state_for_deploy_local_returns_none(): """Local Postgres cannot be pre-provisioned (its container's host doesn't exist yet), so the @@ -2018,6 +2041,40 @@ async def test_deploy_returns_the_gateway_record_for_the_env(caplog): state_provider.teardown.assert_awaited_once_with(external) +@pytest.mark.asyncio +async def test_a_deploy_with_websites_records_the_card_read_once_they_are_ready(): + """The gateway composes its card from its children's as it reads them, so the card read before the sites were up + may lack them; the one recorded is read once they're ready, as what reads the sites from it expects.""" + from agent_env.providers.env_state import DatabaseStateProvider + + gp, external = EnvironmentGatewayProvider(), MagicMock(state_type="remote", instance_id="st-ext") + state_provider = MagicMock(spec=DatabaseStateProvider, prepare=AsyncMock(), teardown=AsyncMock()) + vm = MagicMock(sandbox_id="vm-1", type="modal_vm", tunnel_urls={18765: "https://vm.example"}, host_port=lambda port: port, host_ips=(), + load_docker_images=AsyncMock(), exec_script=AsyncMock(), exec_with_output=AsyncMock(return_value=(0, "gateway", ""))) + env = WebsiteEnv(id="web-shop", version=1, backend_docker_image_artifact=_loadable("shop-be"), + frontend_docker_image_artifact=_loadable("shop-fe"), environment_name="shop") + before, after = {"name": "env1234", "children_environments": []}, {"name": "env1234", "children_environments": [{"name": "shop"}]} + + with patch("agent_env.providers.env_state.acquire_state_for_deploy", AsyncMock(return_value=external)), \ + patch("agent_env.providers.env_state.build_state_provider", return_value=state_provider), \ + patch("agent_env.env.env.Env.get", MagicMock(return_value=MagicMock(docker_image_artifact=_loadable("gw")))), \ + patch("agent_env.config.get_config", return_value=MagicMock(default_gateway_env_id="gw-id")), \ + patch.object(gp, "create_docker_compose", return_value="services: {}"), \ + patch.object(gp, "_wait_for_gateway", AsyncMock(return_value=True)), \ + patch.object(gp, "_get_container_id", AsyncMock(return_value=None)), \ + patch.object(gp, "_wait_for_tunnel", AsyncMock(side_effect=[before, after])), \ + patch.object(gp, "_refuse_tool_name_clashes", AsyncMock()), \ + patch.object(env_gateway_provider, "_wait_for_websites", AsyncMock()) as websites_waited, \ + patch.object(env_gateway_provider, "_tool_names", AsyncMock(return_value=set())), \ + patch("agent_env.providers.env_providers.env_gateway_provider.asyncio.sleep", AsyncMock()): + record = await gp.deploy(env, _vm_provider(create_vm=AsyncMock(return_value=vm)), ttl_seconds=60, env_state_type="remote") + + websites_waited.assert_awaited_once() + assert record.environment_card == after + assert record.website_frontend_urls == {"shop": "https://vm.example/website/shop/"} + await gp.close() + + @pytest.mark.asyncio async def test_deploy_records_the_gateways_sidecar_urls_and_its_containers(): gp, sandboxes = EnvironmentGatewayProvider(), {i: _FakeSandbox(i, 18765, f"https://{i}") for i in ("gw", "srv", "db", "pg", "dm")} diff --git a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py index 9e1c10ad0..2f79f44a5 100644 --- a/tst/unit/providers/sandbox_providers/test_reachability_contributions.py +++ b/tst/unit/providers/sandbox_providers/test_reachability_contributions.py @@ -21,6 +21,7 @@ SandboxProvider, _resolve_container_env, all_sandbox_container_env, + sandbox_request_header_names, sandbox_request_headers_for_url, ) @@ -176,6 +177,23 @@ def test_headers_union_across_providers_and_a_raising_sibling_is_skipped(monkeyp assert "X-Typo" not in headers +def test_the_header_names_every_provider_attaches_are_listed_without_resolving_a_credential(monkeypatch, no_cf_creds, + tmp_path): + (tmp_path / "config.toml").write_text(textwrap.dedent( + """ + [sandbox.providers.static] + impl = "tst.unit.providers.sandbox_providers.test_reachability_contributions:_StaticHeaderProvider" + [sandbox.providers.extra] + impl = "tst.unit.providers.sandbox_providers.test_reachability_contributions:_ExtraHeaderProvider" + """ + ) + _CF_TOML) + monkeypatch.setenv("AGENT_ENV_CONFIG", str(tmp_path / "config.toml")) + reset_config() + + assert sandbox_request_header_names() == ("cf-access-client-id", "cf-access-client-secret", "x-extra", + "x-extra-static", "x-static") + + def test_static_headers_emit_without_any_resolution(monkeypatch, tmp_path): import agent_env.config as config_mod diff --git a/uv.lock b/uv.lock index 13add648c..9f3fb1256 100644 --- a/uv.lock +++ b/uv.lock @@ -131,6 +131,7 @@ dev = [ { name = "transformers" }, { name = "urllib3" }, { name = "vercel-sandbox" }, + { name = "websockets" }, ] gcp = [ { name = "google-api-core" }, @@ -266,6 +267,7 @@ requires-dist = [ { name = "urllib3", marker = "extra == 'boat'", specifier = ">=2.1.0,<3" }, { name = "uvicorn", specifier = ">=0.47" }, { name = "vercel-sandbox", marker = "extra == 'vercel'", specifier = ">=0.7.0,<0.8" }, + { name = "websockets", marker = "extra == 'dev'", specifier = ">=13" }, ] provides-extras = ["openenv", "harness", "trl", "nemo-gym", "verl", "skyrl", "miles", "slime", "aws", "gcp", "sail", "vercel", "tensorlake", "tenki", "boxd", "daytona", "boat", "coreweave", "dev", "smol"] From 4872309430f5928dfdfba9af398cc29aaa6b30e4 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 09:58:47 -0800 Subject: [PATCH 2/3] fix(gateway): the provider's headers reach no browser on the CDP routes either The website browser's CDP answer, its discovery and its WebSocket relay forwarded a request's headers on routes of their own, before the withheld filter; BrowserRoutes now takes the withheld names and drops them on all three. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/browser_proxy.py | 14 +++++---- src/agent_env/env/gateway/gateway.py | 4 +-- tst/unit/env/gateway/browser_route_test.py | 35 ++++++++++++++++++---- 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/src/agent_env/env/gateway/browser_proxy.py b/src/agent_env/env/gateway/browser_proxy.py index 4f5a2cc90..0b7735047 100644 --- a/src/agent_env/env/gateway/browser_proxy.py +++ b/src/agent_env/env/gateway/browser_proxy.py @@ -125,12 +125,14 @@ def filter(self, record: logging.LogRecord) -> bool: class BrowserRoutes: """The ``/browser//cdp//`` 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,8 @@ 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 = [(name, value) for name, value in request.headers.raw + if name.lower() not in HOP_BY_HOP and name.lower() not in 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 +195,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) @@ -211,7 +214,8 @@ async def relay(self, websocket: WebSocket) -> None: "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-")] + if name.encode() not in HOP_BY_HOP and name.encode() not in self._withheld + and not name.startswith("sec-websocket-")] offered = [p.strip() for p in websocket.headers.get("sec-websocket-protocol", "").split(",") if p.strip()] try: upstream = await connect( diff --git a/src/agent_env/env/gateway/gateway.py b/src/agent_env/env/gateway/gateway.py index fb3248265..79f6d1911 100644 --- a/src/agent_env/env/gateway/gateway.py +++ b/src/agent_env/env/gateway/gateway.py @@ -154,15 +154,15 @@ def __init__( 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): log.addFilter(TokenRedaction()) self.gateway_mode = gateway_mode self._service_db_url = service_db_url - self._withheld = frozenset(name.strip().lower().encode() for name in provider_headers if name.strip()) self._website_routes: dict[str, list[WebsiteRoute] | None] = {} self._db_conn: psycopg2.extensions.connection | None = None if service_db_url: diff --git a/tst/unit/env/gateway/browser_route_test.py b/tst/unit/env/gateway/browser_route_test.py index 9c213f9c6..eb793a945 100644 --- a/tst/unit/env/gateway/browser_route_test.py +++ b/tst/unit/env/gateway/browser_route_test.py @@ -54,14 +54,21 @@ def _computer_card(*apps: dict) -> dict: "children_environments": list(apps)} -def _browser_computer(card: dict) -> Starlette: +def _browser_computer(card: dict, seen: list[dict] | None = None) -> Starlette: + """A computer serving ``card`` and a browser's CDP routes; ``seen`` gets the headers of each request to a CDP route.""" + def saw(connection) -> None: + if seen is not None: + seen.append(dict(connection.headers)) + async def served_card(_: Request) -> JSONResponse: return JSONResponse(card) - async def cdp(_: Request) -> JSONResponse: + async def cdp(request: Request) -> JSONResponse: + saw(request) return JSONResponse({"http_url": MINTED, "ws_url": "ws://browser-host:7000/svc/chromium/cdp/t0k3n/devtools/browser"}) async def discovery(request: Request) -> JSONResponse: + saw(request) if request.path_params["token"] != TOKEN: return JSONResponse({"error": "agentenv didn't mint this CDP address"}, status_code=403) if request.path_params["rest"].rstrip("/") == "json/list": @@ -72,6 +79,7 @@ async def discovery(request: Request) -> JSONResponse: f"ws://browser-host:7000/svc/chromium/cdp/{TOKEN}/devtools/browser"}) async def browser_socket(websocket: WebSocket) -> None: + saw(websocket) if websocket.path_params["token"] != TOKEN: await websocket.close(code=1008) return @@ -115,14 +123,15 @@ async def _serve(app) -> tuple[uvicorn.Server, str]: @asynccontextmanager -async def _gateway(tmp_path, monkeypatch, card: dict | None = None): +async def _gateway(tmp_path, monkeypatch, card: dict | None = None, seen: list[dict] | None = None, + provider_headers: tuple[str, ...] = ()): """The real gateway app on loopback, its /svc child ``web`` a browser computer serving ``card``; yields the gateway's base URL and a reader of its browser_connection events, which waits for ``at_least``: each is written as its connection ends, just after the client sees the close.""" monkeypatch.setattr(gw_mod, "GATEWAY_TRAJECTORY_FILE", str(tmp_path / "trajectory.jsonl")) - upstream, upstream_url = await _serve(_browser_computer(_computer_card(_app_card()) if card is None else card)) + upstream, upstream_url = await _serve(_browser_computer(_computer_card(_app_card()) if card is None else card, seen)) gw = Gateway(host="127.0.0.1", port=0, server_name="t", internal_mcp_servers=[], - rest_proxy_urls={"web": upstream_url}) + rest_proxy_urls={"web": upstream_url}, provider_headers=provider_headers) gateway, gateway_url = await _serve(gw._asgi_app()) async def events(at_least: int) -> list[dict]: for _ in range(50): @@ -183,6 +192,22 @@ async def test_a_websocket_is_relayed_both_ways_and_each_connection_is_an_event_ assert all(TOKEN not in json.dumps(e) and "forged" not in json.dumps(e) for e in recorded) +@pytest.mark.asyncio +async def test_the_providers_headers_reach_the_browser_on_no_cdp_route(tmp_path, monkeypatch): + seen: list[dict] = [] + sent = {"CF-Access-Client-Id": "id", "X-Kept": "yes"} + async with (_gateway(tmp_path, monkeypatch, seen=seen, provider_headers=("CF-Access-Client-Id",)) as (url, _), + httpx.AsyncClient() as client): + await client.post(f"{url}/svc/web{CDP_PATH}", json={}, headers=sent) + discovered = (await client.get(f"{url}/browser/web/cdp/{TOKEN}/json/version", headers=sent)).json() + async with connect(discovered["webSocketDebuggerUrl"], additional_headers=sent) as conn: + await conn.send("hi") + await conn.recv() + + assert len(seen) == 3 + assert [("cf-access-client-id" in headers, headers.get("x-kept")) for headers in seen] == [(False, "yes")] * 3 + + @pytest.mark.asyncio async def test_a_relay_that_fails_in_between_closes_both_sides_and_records_the_failures_type(tmp_path, monkeypatch): real = browser_proxy.connect From 7b0b157675286820aeb35c4fd63547af74472793 Mon Sep 17 00:00:00 2001 From: Edgar Arakelyan Date: Thu, 5 Nov 2026 11:06:28 -0800 Subject: [PATCH 3/3] refactor(gateway): one rule for the headers any route passes on client_headers is what the gateway may pass on of a connection's headers, minus the hop-by-hop set and the provider's credentials; upstream_headers builds on it, and the CDP routes' discovery and WebSocket relay call it instead of filtering their own copies, so a header withheld once is withheld on every route. Co-Authored-By: Claude Opus 5.5 --- src/agent_env/env/gateway/browser_proxy.py | 11 +++++------ src/agent_env/env/gateway/rest_proxy.py | 22 ++++++++++++++-------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/agent_env/env/gateway/browser_proxy.py b/src/agent_env/env/gateway/browser_proxy.py index 0b7735047..2d08ddd03 100644 --- a/src/agent_env/env/gateway/browser_proxy.py +++ b/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__) @@ -170,8 +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 and name.lower() not in self._withheld] + 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) @@ -213,9 +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 name.encode() not in self._withheld - 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( diff --git a/src/agent_env/env/gateway/rest_proxy.py b/src/agent_env/env/gateway/rest_proxy.py index e9fca4336..6dbe137f9 100644 --- a/src/agent_env/env/gateway/rest_proxy.py +++ b/src/agent_env/env/gateway/rest_proxy.py @@ -134,15 +134,21 @@ def upstream_url(base_url: str, request: HTTPConnection) -> str: return url +def client_headers(connection: HTTPConnection, withheld: frozenset[bytes] = frozenset()) -> list[tuple[bytes, bytes]]: + """What any route of the gateway may pass on of a connection's headers: one entry per raw header, so repeated + fields (cookie, x-forwarded-for) keep their multiplicity, minus the hop-by-hop set and the headers named in + `withheld` (lower case), the credentials of the provider in front of the gateway, so no service behind it reads + them.""" + return [(k, v) for k, v in connection.headers.raw if k.lower() not in HOP_BY_HOP and k.lower() not in withheld] + + def upstream_headers(request: HTTPConnection, role: str, *, withheld: frozenset[bytes] = frozenset()) -> httpx.Headers: - """One entry per raw header, so repeated fields (cookie, x-forwarded-for) keep their multiplicity. The role - header is the gateway's to set: the client's copies are dropped and the role the gateway attributed the - request to goes in their place, so the upstream never reads a value the gateway did not write. So are the - forwarded-address headers (`forwarded_address`). The headers named in `withheld` (lower case), the credentials of - the provider in front of the gateway, are dropped, so no service behind it reads them.""" - headers = httpx.Headers([(k, v) for k, v in request.headers.raw - if k.lower() not in HOP_BY_HOP and k.lower() != ROLE_HEADER_BYTES - and k.lower() not in _FORWARDED and k.lower() not in withheld]) + """The client's headers (``client_headers``) as a proxied request sends them. The role header is the gateway's to + set: the client's copies are dropped and the role the gateway attributed the request to goes in their place, so + the upstream never reads a value the gateway did not write. So are the forwarded-address headers + (`forwarded_address`).""" + headers = httpx.Headers([(k, v) for k, v in client_headers(request, withheld) + if k.lower() != ROLE_HEADER_BYTES and k.lower() not in _FORWARDED]) headers[AGENT_ENV_ROLE_HEADER] = role headers.update(forwarded_address(request)) # Bytes are forwarded undecoded, so only ask upstream for encodings the caller can take.