mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(proxy): route Codex Live voice through a dedicated /v1/live transport (#2709)
## Description Codex Live traffic currently reaches an unrouted WebSocket path and receives HTTP 403 before the proxy can contact an upstream. Closes #2653 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature - [ ] Breaking change - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring ## Changes Made - Add a dedicated `/v1/live` WebSocket route family and transparent transport. - Preserve subscription auth routing, account headers, origin policy, subprotocols, and text/binary frame bytes. - Propagate WebSocket close metadata and cancel relay tasks deterministically on every exit. - Keep Live outside the Responses parser, compression, memory injection, and Responses beta-header path. - Keep generic HTTP paths on the existing catch-all and document the Live aliases plus the derived-path override. - Add real-app route, relay, and loopback integration proof. - Add coverage for authorization fallback, defensive receive events, and cancellation cleanup in the Live relay. ## Testing The focused Live handshake, preservation suites, Ruff, format, and diff checks pass. The base comparison, Codex Desktop owner round trip, and ChatGPT backend acceptance of the derived `/backend-api/codex/live` path remain untested. - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom --ignore-missing-imports`) - [x] New tests added for the reported failure - [x] Manual loopback testing performed ### Test Output ```text uv run pytest tests/test_codex_live.py -q: 6 passed, 7 warnings in 11.03s uv run pytest tests/test_provider_proxy_routes.py tests/test_proxy_codex_route_aliases.py tests/test_provider_codex_endpoints.py tests/test_openai_codex_routing.py -q: 53 passed in 90.54s (0:01:30) uv run ruff check headroom/providers/codex/live.py headroom/providers/proxy_routes.py headroom/proxy/handlers/openai.py headroom/proxy/ws_headers.py tests/test_codex_live.py: All checks passed uv run ruff format --check headroom/providers/codex/live.py headroom/providers/proxy_routes.py headroom/proxy/handlers/openai.py headroom/proxy/ws_headers.py tests/test_codex_live.py: 5 files already formatted uv run mypy headroom --ignore-missing-imports: Success: no issues found in 508 source files git diff --check: pass ``` ## Real Behavior Proof The local WebSocket integration floor uses real uvicorn, a real WebSocket client, and a real loopback WebSocket upstream. The base 403 comparison was not run. Head observes HTTP 101 on every Live alias and a byte-identical binary frame relay. - Environment: Windows, CPython 3.13, the Headroom proxy test environment. - Exact command / steps: run the focused Live test against the local uvicorn proxy and loopback WebSocket upstream, then run the preservation suite listed in `Test Output`. - Observed result: all four Live aliases return HTTP 101, negotiate `codex.live.v1`, preserve text and binary frames, and pass the preservation suite. - Not tested: Codex Desktop Live session; ChatGPT backend acceptance of `/backend-api/codex/live`; the base 403 comparison. ## Review Readiness - Live has a separate transport and does not enter Responses handling. - Existing Responses and generic passthrough suites remain preservation gates. - No `CHANGELOG.md` or install/crate changes are included. - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] Closes #2653 - [x] Real loopback handshake and binary-frame proof required - [x] No audio payload logging - [x] No unqualified end-to-end claim ## Screenshots Not applicable. ## Additional Notes The upstream Live path is derived from the repository’s Codex URL formula and remains explicitly unconfirmed until owner evidence is available.
This commit is contained in:
Rod Boev
GitHub
parent
3a27c4dacb
commit
232fb49c73
6 changed files with 894 additions and 14 deletions
|
|
@ -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.
|
||||
|
|
|
|||
253
headroom/providers/codex/live.py
Normal file
253
headroom/providers/codex/live.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
16
headroom/proxy/ws_headers.py
Normal file
16
headroom/proxy/ws_headers.py
Normal file
|
|
@ -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",
|
||||
}
|
||||
)
|
||||
581
tests/test_codex_live.py
Normal file
581
tests/test_codex_live.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Add a link
Reference in a new issue