diff --git a/docs/content/docs/proxy.mdx b/docs/content/docs/proxy.mdx index 68855f62c..5de6a334d 100644 --- a/docs/content/docs/proxy.mdx +++ b/docs/content/docs/proxy.mdx @@ -385,6 +385,22 @@ For Codex-compatible clients, the proxy also accepts these alias paths and route Matching WebSocket and subpath aliases are also supported for Codex flows. +### Codex Live voice WebSocket + +The proxy relays Codex Live voice frames without parsing or transforming them. +These paths use the same transparent transport: + +- `ws://localhost:8787/v1/live` +- `ws://localhost:8787/v1/codex/live` +- `ws://localhost:8787/backend-api/live` +- `ws://localhost:8787/backend-api/codex/live` + +Subscription authentication uses the derived ChatGPT backend path. API-key +authentication preserves the selected OpenAI-compatible base URL and inbound +path. The backend Live suffix defaults to `/live` and can be corrected with +`HEADROOM_CODEX_LIVE_WS_PATH` if the upstream contract changes. The exact +ChatGPT backend path is not confirmed by this proxy documentation. + ### `POST /v1internal:streamGenerateContent` Google Cloud Code Assist / Antigravity compatibility endpoint used by Pi-style `google-gemini-cli` and `google-antigravity` providers. diff --git a/headroom/providers/codex/live.py b/headroom/providers/codex/live.py new file mode 100644 index 000000000..f59baecb1 --- /dev/null +++ b/headroom/providers/codex/live.py @@ -0,0 +1,253 @@ +"""Transparent Codex Live WebSocket transport.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +from collections.abc import Mapping +from typing import Any, cast + +from fastapi import WebSocket + +from headroom.copilot_auth import apply_copilot_api_auth, build_copilot_upstream_url +from headroom.providers.codex.endpoints import codex_backend_ws_url +from headroom.providers.codex.runtime import resolve_codex_routing +from headroom.proxy.handlers.openai import _is_allowed_websocket_origin +from headroom.proxy.helpers import _strip_internal_headers, merge_extra_headers +from headroom.proxy.ws_headers import WS_HOP_BY_HOP_HEADERS + +logger = logging.getLogger("headroom.providers.codex.live") + +CODEX_LIVE_ROUTE_PATHS: tuple[str, ...] = ( + "/v1/live", + "/v1/codex/live", + "/backend-api/live", + "/backend-api/codex/live", +) +CODEX_LIVE_WS_PATH_ENV = "HEADROOM_CODEX_LIVE_WS_PATH" +DEFAULT_CODEX_LIVE_WS_PATH = "/live" + + +def codex_live_ws_path() -> str: + """Return the correctable, derived Codex backend Live path.""" + configured = os.environ.get(CODEX_LIVE_WS_PATH_ENV, "").strip() + path = configured or DEFAULT_CODEX_LIVE_WS_PATH + return path if path.startswith("/") else f"/{path}" + + +def codex_live_websocket_url( + *, + subscription: bool, + base_url: str, + path: str = "/v1/live", + query: str = "", +) -> str: + """Build the Live WebSocket URL from the selected auth mode.""" + if subscription: + url = codex_backend_ws_url(codex_live_ws_path()) + else: + ws_base = base_url.replace("https://", "wss://").replace("http://", "ws://") + url = build_copilot_upstream_url(ws_base, path) + return f"{url}?{query}" if query else url + + +def _forward_headers(headers: Mapping[str, str]) -> dict[str, str]: + forwarded = { + key: value for key, value in headers.items() if key.lower() not in WS_HOP_BY_HOP_HEADERS + } + return _strip_internal_headers(forwarded) + + +def _ensure_live_authorization(headers: Mapping[str, str]) -> dict[str, str]: + """Use the configured OpenAI key when a client omitted authorization.""" + if any(key.lower() == "authorization" for key in headers): + return dict(headers) + api_key = os.environ.get("OPENAI_API_KEY", "").strip() + if not api_key: + return dict(headers) + return {**headers, "Authorization": f"Bearer {api_key}"} + + +_NON_WIRE_CLOSE_CODES = {1004, 1005, 1006, 1015} + + +def _close_info(source: object, default_code: int) -> tuple[int, str]: + received = getattr(source, "rcvd", None) + if isinstance(source, Mapping): + code = source.get("code") + reason = source.get("reason") + else: + code = getattr(source, "close_code", None) + reason = getattr(source, "close_reason", None) + raw_code: object = code or getattr(received, "code", None) + try: + code = int(str(raw_code)) + except (TypeError, ValueError): + code = default_code + if code < 1000 or code > 4999 or code in _NON_WIRE_CLOSE_CODES: + code = default_code + + reason = reason or getattr(received, "reason", None) or "" + return code, str(reason)[:120] + + +async def handle_codex_live_websocket( + websocket: WebSocket, + proxy: object, + openai_base_url: str, + inbound_path: str, +) -> None: + """Relay Live text and binary frames without parsing or transforming them.""" + client_headers = dict(websocket.headers) + if not _is_allowed_websocket_origin(client_headers): + await websocket.close(code=1008, reason="origin not allowed") + return + + try: + import websockets + except ImportError: + logger.exception("Codex Live relay unavailable because websockets is not installed") + await websocket.accept() + await websocket.close( + code=1011, reason="websockets package not installed; pip install websockets" + ) + return + + forwarded_headers = _forward_headers(client_headers) + decision = resolve_codex_routing(forwarded_headers) + forwarded_headers = decision.headers + websocket_url = getattr(websocket, "url", None) + query = str(getattr(websocket_url, "query", "") or "") + upstream_url = codex_live_websocket_url( + subscription=decision.is_chatgpt_auth, + base_url=openai_base_url, + path=inbound_path, + query=query, + ) + forwarded_headers = await apply_copilot_api_auth(forwarded_headers, url=upstream_url) + config = getattr(proxy, "config", None) + forwarded_headers = merge_extra_headers( + forwarded_headers, + getattr(config, "openai_extra_headers", None), + ) + if not any(key.lower() == "authorization" for key in forwarded_headers): + if os.environ.get("OPENAI_API_KEY", "").strip(): + forwarded_headers = _ensure_live_authorization(forwarded_headers) + logger.debug("Codex Live injected Authorization from OPENAI_API_KEY") + else: + logger.warning("Codex Live has no Authorization header or OPENAI_API_KEY") + + raw_protocols = next( + (value for key, value in client_headers.items() if key.lower() == "sec-websocket-protocol"), + "", + ) + subprotocols = [value.strip() for value in raw_protocols.split(",") if value.strip()] + upstream = None + relay_tasks: set[asyncio.Task[None]] = set() + accepted = False + client_disconnected = False + close_code = 1000 + close_reason = "" + try: + try: + upstream = await websockets.connect( + upstream_url, + additional_headers=forwarded_headers, + # websockets types wire protocol tokens as a nominal wrapper. + subprotocols=cast(Any, subprotocols or None), + ssl=True if upstream_url.startswith("wss://") else None, + open_timeout=max(30, getattr(config, "connect_timeout_seconds", 10) * 3), + close_timeout=10, + ping_interval=20, + ping_timeout=None, + max_size=None, + ) + except Exception: + logger.exception("Codex Live upstream handshake failed url=%s", upstream_url) + await websocket.close(code=1011, reason="upstream connection failed") + return + + selected_subprotocol = getattr(upstream, "subprotocol", None) + if selected_subprotocol not in subprotocols: + selected_subprotocol = None + await websocket.accept(subprotocol=selected_subprotocol) + accepted = True + + async def client_to_upstream() -> None: + nonlocal client_disconnected, close_code, close_reason + while True: + message = await websocket.receive() + message_type = message.get("type") + if message_type == "websocket.disconnect": + client_disconnected = True + close_code, close_reason = _close_info(message, 1000) + with contextlib.suppress(Exception): + await upstream.close(code=close_code, reason=close_reason) + return + if message_type != "websocket.receive": + continue + 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 upstream_to_client() -> None: + nonlocal close_code, close_reason + try: + async for message in upstream: + if isinstance(message, bytes): + await websocket.send_bytes(message) + else: + await websocket.send_text(message) + except asyncio.CancelledError: + raise + except Exception as exc: + close_code, close_reason = _close_info(exc, 1011) + logger.warning( + "Codex Live upstream relay failed code=%s reason=%s error=%s", + close_code, + close_reason, + type(exc).__name__, + ) + else: + close_code, close_reason = _close_info(upstream, 1000) + + relay_tasks = { + asyncio.create_task(client_to_upstream()), + asyncio.create_task(upstream_to_client()), + } + done, _ = await asyncio.wait( + relay_tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + with contextlib.suppress(asyncio.CancelledError): + error = task.exception() + if error is not None: + close_code, close_reason = _close_info(error, 1011) + logger.warning( + "Codex Live relay task failed task=%s code=%s reason=%s error=%s", + task.get_name(), + close_code, + close_reason, + type(error).__name__, + ) + except asyncio.CancelledError: + raise + except Exception: + close_code, close_reason = 1011, "relay failed" + logger.exception("Codex Live relay failed url=%s", upstream_url) + finally: + pending = [task for task in relay_tasks if not task.done()] + for task in pending: + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + if upstream is not None: + with contextlib.suppress(Exception): + await upstream.close() + if accepted and not client_disconnected: + with contextlib.suppress(Exception): + await websocket.close(code=close_code, reason=close_reason) diff --git a/headroom/providers/proxy_routes.py b/headroom/providers/proxy_routes.py index 002c3661d..b7c6c595b 100644 --- a/headroom/providers/proxy_routes.py +++ b/headroom/providers/proxy_routes.py @@ -9,6 +9,10 @@ from typing import Any from fastapi import FastAPI, Request, WebSocket from headroom.providers.cloudcode import normalize_cloudcode_passthrough_path +from headroom.providers.codex.live import ( + CODEX_LIVE_ROUTE_PATHS, + handle_codex_live_websocket, +) from headroom.providers.codex.responses import handle_chatgpt_codex_responses_subpath from headroom.providers.model_metadata import ( MODEL_METADATA_LIST_ENDPOINT, @@ -175,6 +179,26 @@ def _register_openai_responses_routes(app: FastAPI, proxy: Any) -> None: _register_openai_responses_subpath_route(app, proxy, spec) +def _register_codex_live_routes(app: FastAPI, proxy: Any) -> None: + for path in CODEX_LIVE_ROUTE_PATHS: + + def register_websocket_route(route_path: str) -> None: + async def codex_live_websocket(websocket: WebSocket): + await handle_codex_live_websocket( + websocket, + proxy, + _api_target(proxy, "openai"), + route_path, + ) + + codex_live_websocket.__name__ = ( + route_path.strip("/").replace("/", "_") + "_live_websocket" + ) + app.websocket(route_path)(codex_live_websocket) + + register_websocket_route(path) + + def _register_openai_image_route(app: FastAPI, proxy: Any, endpoint: OpenAIImageEndpoint) -> None: async def openai_image_endpoint(request: Request): return await handle_openai_image_endpoint( @@ -434,6 +458,8 @@ def register_provider_routes(app: FastAPI, proxy: Any) -> None: _register_openai_image_routes(app, proxy) + _register_codex_live_routes(app, proxy) + _register_provider_passthrough_routes(app, proxy) @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "HEAD"]) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 22803d234..74bedd16a 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -30,6 +30,7 @@ from headroom.proxy.helpers import ( ) from headroom.proxy.loopback_guard import is_loopback_host from headroom.proxy.stage_timer import StageTimer, emit_stage_timings_log +from headroom.proxy.ws_headers import WS_HOP_BY_HOP_HEADERS from headroom.proxy.ws_session_registry import ( TerminationCause, WebSocketSessionRegistry, @@ -5718,20 +5719,7 @@ class OpenAIHandlerMixin: # These are WebSocket handshake mechanics that the `websockets` library # generates fresh for the upstream connection — forwarding them would conflict. # Everything else (auth, org, beta, user-agent, custom headers) is forwarded as-is. - _skip_headers = frozenset( - { - "host", # must match upstream, not local proxy - "connection", # hop-by-hop - "upgrade", # hop-by-hop - "sec-websocket-key", # per-connection cryptographic nonce - "sec-websocket-version", # protocol version (websockets lib sets this) - "sec-websocket-extensions", # per-connection negotiation - "sec-websocket-accept", # server-side only - "sec-websocket-protocol", # handled via subprotocols param below - "content-length", # hop-by-hop - "transfer-encoding", # hop-by-hop - } - ) + _skip_headers = WS_HOP_BY_HOP_HEADERS # PR-A5 (P5-49): also drop internal x-headroom-* from the upstream # WebSocket handshake. Inbound reads on `ws_headers` (memory user-id # below) keep working because we filter only when building diff --git a/headroom/proxy/ws_headers.py b/headroom/proxy/ws_headers.py new file mode 100644 index 000000000..dc6e1698b --- /dev/null +++ b/headroom/proxy/ws_headers.py @@ -0,0 +1,16 @@ +"""Shared WebSocket handshake header policy.""" + +WS_HOP_BY_HOP_HEADERS = frozenset( + { + "host", + "connection", + "upgrade", + "sec-websocket-key", + "sec-websocket-version", + "sec-websocket-extensions", + "sec-websocket-accept", + "sec-websocket-protocol", + "content-length", + "transfer-encoding", + } +) diff --git a/tests/test_codex_live.py b/tests/test_codex_live.py new file mode 100644 index 000000000..b5aa57fa3 --- /dev/null +++ b/tests/test_codex_live.py @@ -0,0 +1,581 @@ +from __future__ import annotations + +import asyncio +import base64 +import json +import logging +import os +import socket +import sys +import threading +import time +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch + +import pytest +import uvicorn +import websockets +from starlette.websockets import WebSocket + +from headroom.providers.codex.live import ( + CODEX_LIVE_ROUTE_PATHS, + DEFAULT_CODEX_LIVE_WS_PATH, + _close_info, + _ensure_live_authorization, + _forward_headers, + codex_live_websocket_url, + codex_live_ws_path, + handle_codex_live_websocket, +) +from headroom.providers.codex.runtime import resolve_codex_routing +from headroom.proxy.server import ProxyConfig, create_app + + +def _jwt(payload: dict[str, Any]) -> str: + encoded = base64.urlsafe_b64encode(json.dumps(payload).encode()).decode().rstrip("=") + return f"header.{encoded}.signature" + + +def test_live_aliases_are_registered_as_websocket_routes(monkeypatch) -> None: + monkeypatch.setenv("HEADROOM_REQUIRE_RUST_CORE", "false") + app = create_app( + ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + openai_api_url="https://api.openai.test", + ) + ) + + for path in CODEX_LIVE_ROUTE_PATHS: + matching = [route for route in app.routes if route.path == path] + assert any("live_websocket" in route.name for route in matching) + + +def test_live_auth_modes_and_derived_paths(monkeypatch) -> None: + subscription_headers = { + "authorization": "Bearer " + + _jwt( + { + "https://api.openai.com/auth": {"chatgpt_account_id": "acct-live"}, + } + ), + } + subscription = _forward_headers(subscription_headers) + decision = resolve_codex_routing(subscription) + assert ( + codex_live_websocket_url( + subscription=True, + base_url="https://api.openai.test", + path="/backend-api/live", + query="model=gpt-5.4", + ) + == "wss://chatgpt.com/backend-api/codex/live?model=gpt-5.4" + ) + assert subscription["authorization"] == subscription_headers["authorization"] + assert decision.headers["ChatGPT-Account-ID"] == "acct-live" + + monkeypatch.setenv("OPENAI_API_KEY", "env-live-key") + assert _ensure_live_authorization({}) == {"Authorization": "Bearer env-live-key"} + assert _ensure_live_authorization({"authorization": "Bearer client-key"}) == { + "authorization": "Bearer client-key" + } + monkeypatch.delenv("OPENAI_API_KEY") + existing_headers = {"X-Trace": "keep"} + assert _ensure_live_authorization(existing_headers) == existing_headers + + assert ( + codex_live_websocket_url( + subscription=False, + base_url="http://127.0.0.1:9000", + path="/backend-api/live", + query="mode=live", + ) + == "ws://127.0.0.1:9000/backend-api/live?mode=live" + ) + assert DEFAULT_CODEX_LIVE_WS_PATH == "/live" + monkeypatch.setenv("HEADROOM_CODEX_LIVE_WS_PATH", "/custom/live") + assert codex_live_ws_path() == "/custom/live" + assert ( + codex_live_websocket_url( + subscription=True, + base_url="https://api.openai.test", + query="mode=live", + ) + == "wss://chatgpt.com/backend-api/codex/custom/live?mode=live" + ) + + +def test_live_headers_strip_internal_and_handshake_headers_without_beta_injection() -> None: + headers = _forward_headers( + { + "Authorization": "Bearer live-token", + "ChatGPT-Account-ID": "acct-live", + "OpenAI-Beta": "client-beta", + "X-Headroom-User-ID": "private", + "Connection": "keep-alive", + "Sec-WebSocket-Key": "nonce", + "Sec-WebSocket-Protocol": "codex.live.v1", + "X-Trace": "keep", + }, + ) + lowered = {key.lower(): value for key, value in headers.items()} + assert lowered == { + "authorization": "Bearer live-token", + "chatgpt-account-id": "acct-live", + "openai-beta": "client-beta", + "x-trace": "keep", + } + + +def test_live_close_info_preserves_wire_codes_and_reasons() -> None: + assert _close_info({"code": 1008, "reason": "origin not allowed"}, 1011) == ( + 1008, + "origin not allowed", + ) + assert _close_info( + SimpleNamespace(rcvd=SimpleNamespace(code=1013, reason="upstream busy")), + 1011, + ) == (1013, "upstream busy") + assert _close_info({"code": 1006, "reason": "abnormal"}, 1011) == (1011, "abnormal") + + +class _LoopbackLiveUpstream: + def __init__(self) -> None: + self.server: Any = None + self.port = 0 + self.headers: dict[str, str] = {} + self.paths: list[str] = [] + self.received: list[str | bytes] = [] + + async def handler(self, connection: Any) -> None: + request = getattr(connection, "request", None) + request_headers = getattr(request, "headers", None) + if request_headers is None: + request_headers = getattr(connection, "request_headers", {}) + self.headers = dict(request_headers) + request_target = getattr(request, "path", None) or getattr(request, "target", None) + if request_target is not None: + self.paths.append(str(request_target)) + async for message in connection: + assert isinstance(message, (str, bytes)) + self.received.append(message) + await connection.send(message) + + async def start(self) -> None: + self.port = _free_port() + self.server = await websockets.serve( + self.handler, + "127.0.0.1", + self.port, + subprotocols=["codex.live.v1"], + ) + + async def stop(self) -> None: + if self.server is not None: + self.server.close() + await self.server.wait_closed() + + +class _ProxyThread: + def __init__(self, port: int, upstream_port: int) -> None: + app = create_app( + ProxyConfig( + host="127.0.0.1", + port=port, + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + openai_api_url=f"http://127.0.0.1:{upstream_port}", + ) + ) + self.server = uvicorn.Server( + uvicorn.Config( + app, + host="127.0.0.1", + port=port, + log_level="warning", + loop="asyncio", + lifespan="on", + ws="websockets", + ) + ) + self.thread = threading.Thread(target=self.server.run, daemon=True) + + def start(self) -> None: + self.thread.start() + deadline = time.perf_counter() + 15.0 + while time.perf_counter() < deadline: + if self.server.started: + return + time.sleep(0.05) + raise RuntimeError("uvicorn proxy failed to start") + + def stop(self) -> None: + self.server.should_exit = True + self.thread.join(timeout=15.0) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +@pytest.mark.asyncio +async def test_live_handler_propagates_close_metadata_and_cleans_tasks(monkeypatch, caplog) -> None: + class _Client: + headers = {"authorization": "Bearer sk-live"} + url = SimpleNamespace(query="turn=1") + + def __init__(self, message: dict[str, Any], *, fail_accept: bool = False) -> None: + self.message = message + self.fail_accept = fail_accept + self.accepted = False + self.closed: list[tuple[int | None, str | None]] = [] + self.cancelled = False + + async def accept(self, **kwargs: Any) -> None: + if self.fail_accept: + raise RuntimeError("accept failed") + self.accepted = True + + async def close(self, code=None, reason=None) -> None: + self.closed.append((code, reason)) + + async def receive(self) -> dict[str, Any]: + if self.message.get("type") == "wait": + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise + return self.message + + async def send_text(self, message: str) -> None: + del message + + async def send_bytes(self, message: bytes) -> None: + del message + + class _Upstream: + subprotocol = None + + def __init__(self, error: Exception | None = None) -> None: + self.error = error + self.close_calls: list[tuple[int | None, str | None]] = [] + self.sent: list[str | bytes] = [] + self.cancelled = False + + async def send(self, message: str | bytes) -> None: + self.sent.append(message) + + async def close(self, code=None, reason=None) -> None: + self.close_calls.append((code, reason)) + + def __aiter__(self): + async def _events(): + if self.error is not None: + raise self.error + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + self.cancelled = True + raise + if False: + yield b"" + + return _events() + + proxy = SimpleNamespace( + config=SimpleNamespace(openai_extra_headers=None, connect_timeout_seconds=10), + ) + + disconnect_client = _Client( + {"type": "websocket.disconnect", "code": 1008, "reason": "client stopped"} + ) + disconnect_upstream = _Upstream() + monkeypatch.setattr(websockets, "connect", lambda *args, **kwargs: _await(disconnect_upstream)) + await handle_codex_live_websocket( + disconnect_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert disconnect_upstream.close_calls[0] == (1008, "client stopped") + + subscription_client = _Client({"type": "websocket.disconnect", "code": 1000, "reason": "done"}) + subscription_client.headers = { + "authorization": "Bearer " + + _jwt( + { + "https://api.openai.com/auth": { + "chatgpt_account_id": "acct-live", + } + } + ) + } + subscription_upstream = _Upstream() + subscription_connect: dict[str, Any] = {} + + async def connect_subscription(*args: Any, **kwargs: Any) -> Any: + subscription_connect["url"] = args[0] + subscription_connect["headers"] = kwargs["additional_headers"] + return subscription_upstream + + monkeypatch.setattr(websockets, "connect", connect_subscription) + await handle_codex_live_websocket( + subscription_client, + proxy, + "https://api.openai.test", + "/backend-api/live", + ) + assert subscription_connect["url"] == "wss://chatgpt.com/backend-api/codex/live?turn=1" + assert subscription_connect["headers"]["ChatGPT-Account-ID"] == "acct-live" + + fallback_client = _Client({"type": "websocket.disconnect", "code": 1000, "reason": "done"}) + fallback_client.headers = {} + fallback_connect: dict[str, Any] = {} + fallback_upstream = _Upstream() + + async def connect_fallback(*args: Any, **kwargs: Any) -> Any: + fallback_connect["headers"] = kwargs["additional_headers"] + return fallback_upstream + + monkeypatch.setenv("OPENAI_API_KEY", "env-live-key") + monkeypatch.setattr(websockets, "connect", connect_fallback) + await handle_codex_live_websocket( + fallback_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert fallback_connect["headers"]["Authorization"] == "Bearer env-live-key" + + no_auth_client = _Client({"type": "websocket.disconnect", "code": 1000, "reason": "done"}) + no_auth_client.headers = {} + no_auth_connect: dict[str, Any] = {} + no_auth_upstream = _Upstream() + + async def connect_without_auth(*args: Any, **kwargs: Any) -> Any: + no_auth_connect["headers"] = kwargs["additional_headers"] + return no_auth_upstream + + monkeypatch.delenv("OPENAI_API_KEY") + caplog.set_level(logging.WARNING, logger="headroom.providers.codex.live") + monkeypatch.setattr(websockets, "connect", connect_without_auth) + await handle_codex_live_websocket( + no_auth_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert no_auth_connect["headers"] == {} + assert "Codex Live has no Authorization header or OPENAI_API_KEY" in caplog.text + + class _QueuedClient(WebSocket): + def __init__(self, messages: list[dict[str, Any]]) -> None: + scope = { + "type": "websocket", + "path": "/v1/live", + "raw_path": b"/v1/live", + "scheme": "ws", + "query_string": b"", + "headers": [], + "client": ("127.0.0.1", 1234), + "server": ("127.0.0.1", 8788), + "subprotocols": [], + } + super().__init__(scope, self._receive_asgi, self._send_asgi) + self._messages = iter([{"type": "websocket.connect"}, *messages]) + self.sent: list[dict[str, Any]] = [] + + async def _receive_asgi(self) -> dict[str, Any]: + return {"type": "websocket.connect"} + + async def _send_asgi(self, message: dict[str, Any]) -> None: + self.sent.append(message) + + async def receive(self) -> dict[str, Any]: + # Starlette rejects unknown connected-state messages before the + # handler can reach its defensive branch, so inject it at this seam. + try: + return next(self._messages) + except StopIteration: + return {"type": "websocket.disconnect", "code": 1000, "reason": "done"} + + queued_client = _QueuedClient( + [ + {"type": "websocket.other"}, + {"type": "websocket.receive"}, + {"type": "websocket.disconnect", "code": 1000, "reason": "done"}, + ] + ) + queued_upstream = _Upstream() + monkeypatch.setattr(websockets, "connect", lambda *args, **kwargs: _await(queued_upstream)) + await handle_codex_live_websocket( + queued_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert queued_upstream.sent == [] + assert queued_upstream.close_calls[0] == (1000, "done") + + cancellation_client = _Client({"type": "wait"}) + cancellation_upstream = _Upstream() + monkeypatch.setattr( + websockets, "connect", lambda *args, **kwargs: _await(cancellation_upstream) + ) + handler_task = asyncio.create_task( + handle_codex_live_websocket( + cancellation_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + ) + while not cancellation_client.accepted: + await asyncio.sleep(0) + handler_task.cancel() + with pytest.raises(asyncio.CancelledError): + await handler_task + assert cancellation_client.cancelled + assert cancellation_upstream.cancelled + assert cancellation_upstream.close_calls + + class _UpstreamFailure(Exception): + rcvd = SimpleNamespace(code=1013, reason="upstream busy") + + failure_client = _Client({"type": "wait"}) + failure_upstream = _Upstream(_UpstreamFailure("busy")) + monkeypatch.setattr(websockets, "connect", lambda *args, **kwargs: _await(failure_upstream)) + await handle_codex_live_websocket( + failure_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert failure_client.closed[-1] == (1013, "upstream busy") + + handshake_client = _Client({"type": "wait"}) + + async def fail_handshake(*args: Any, **kwargs: Any) -> Any: + del args, kwargs + raise RuntimeError("handshake failed") + + monkeypatch.setattr(websockets, "connect", fail_handshake) + await handle_codex_live_websocket( + handshake_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert handshake_client.closed[-1] == (1011, "upstream connection failed") + + class _ReceiveFailureClient(_Client): + async def receive(self) -> dict[str, Any]: + raise RuntimeError("client receive failed") + + task_failure_client = _ReceiveFailureClient({"type": "wait"}) + task_failure_upstream = _Upstream() + monkeypatch.setattr( + websockets, "connect", lambda *args, **kwargs: _await(task_failure_upstream) + ) + await handle_codex_live_websocket( + task_failure_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert task_failure_client.closed[-1] == (1011, "") + + accept_client = _Client({"type": "wait"}, fail_accept=True) + accept_upstream = _Upstream() + monkeypatch.setattr(websockets, "connect", lambda *args, **kwargs: _await(accept_upstream)) + await handle_codex_live_websocket( + accept_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert accept_upstream.close_calls + + missing_client = _Client({"type": "wait"}) + with patch.dict(sys.modules, {"websockets": None}): + await handle_codex_live_websocket( + missing_client, + proxy, + "https://api.openai.test", + "/v1/live", + ) + assert missing_client.accepted + assert missing_client.closed[-1] == ( + 1011, + "websockets package not installed; pip install websockets", + ) + + +async def _await(value: Any) -> Any: + return value + + +@pytest.mark.asyncio +async def test_live_rg4_real_uvicorn_and_websockets_binary_round_trip() -> None: + previous = os.environ.get("HEADROOM_REQUIRE_RUST_CORE") + os.environ["HEADROOM_REQUIRE_RUST_CORE"] = "false" + upstream = _LoopbackLiveUpstream() + await upstream.start() + proxy = _ProxyThread(_free_port(), upstream.port) + proxy.start() + payload = b"\x00\x01live-voice\xff" + expected: list[str | bytes] = [] + try: + with pytest.raises(websockets.exceptions.InvalidStatus): + async with websockets.connect( + f"ws://127.0.0.1:{proxy.server.config.port}/v1/live", + additional_headers={ + "Authorization": "Bearer sk-live", + "Origin": "https://remote.example", + }, + ): + raise AssertionError("disallowed origin unexpectedly connected") + assert upstream.received == [] + + for index, path in enumerate(CODEX_LIVE_ROUTE_PATHS): + text_payload = f"live-text-{index}" + binary_payload = payload + bytes([index]) + async with websockets.connect( + f"ws://127.0.0.1:{proxy.server.config.port}{path}?turn={index}", + additional_headers={ + "Authorization": "Bearer sk-live", + "OpenAI-Beta": "client-beta", + "X-Headroom-User-ID": "private", + }, + subprotocols=["codex.live.v1"], + ) as client: + response = getattr(client, "response", None) + assert response is not None + assert response.status_code == 101 + assert client.subprotocol == "codex.live.v1" + await client.send(text_payload) + assert await client.recv() == text_payload + await client.send(binary_payload) + assert await client.recv() == binary_payload + expected.extend([text_payload, binary_payload]) + + assert upstream.received == expected + assert upstream.paths == [ + f"{path}?turn={index}" for index, path in enumerate(CODEX_LIVE_ROUTE_PATHS) + ] + assert upstream.headers["authorization"] == "Bearer sk-live" + assert upstream.headers["openai-beta"] == "client-beta" + assert "x-headroom-user-id" not in upstream.headers + finally: + proxy.stop() + await upstream.stop() + if previous is None: + os.environ.pop("HEADROOM_REQUIRE_RUST_CORE", None) + else: + os.environ["HEADROOM_REQUIRE_RUST_CORE"] = previous