diff --git a/headroom/proxy/auth_mode.py b/headroom/proxy/auth_mode.py index 440e4db5a..220ac4397 100644 --- a/headroom/proxy/auth_mode.py +++ b/headroom/proxy/auth_mode.py @@ -253,10 +253,37 @@ def classify_client(headers: Mapping[str, Any] | Any, *, default: str | None = N return default +# OpenAI's Responses API endpoint. In practice this is Codex's endpoint, but a +# proxy can't assume every caller here is Codex — hence +# :func:`should_stamp_codex_client` only stamps callers that don't already +# classify. +CODEX_RESPONSES_PATH = "/v1/responses" + + +def should_stamp_codex_client(path: str, headers: Mapping[str, Any] | Any) -> bool: + """Whether to stamp ``X-Client: codex`` on a request to the proxy. + + Stamping ``X-Client: codex`` on the Responses endpoint makes the backend + take the codex fail-open branch on a compression timeout — Codex treats the + proxy's 413/1009 refusal as a hard connection failure. This is needed + because Codex Desktop's User-Agent (``Codex Desktop/...``) isn't in + :data:`CLIENT_UA_MAP` and would otherwise be refused. + + Returns ``True`` only for an unidentified caller (no ``X-Client`` and no + recognized User-Agent) on the Responses endpoint. A caller that already + classifies is left untouched. + """ + if path != CODEX_RESPONSES_PATH and not path.startswith(CODEX_RESPONSES_PATH + "/"): + return False + return classify_client(headers) is None + + __all__ = [ "AuthMode", "CLIENT_UA_MAP", + "CODEX_RESPONSES_PATH", "SUBSCRIPTION_UA_PREFIXES", "classify_auth_mode", "classify_client", + "should_stamp_codex_client", ] diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index e29274425..067bdcb23 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -45,7 +45,11 @@ import httpx from headroom.agent_savings import proxy_pipeline_kwargs from headroom.copilot_auth import apply_copilot_api_auth, build_copilot_upstream_url from headroom.pipeline import PipelineStage, summarize_routing_markers -from headroom.proxy.auth_mode import classify_auth_mode, classify_client +from headroom.proxy.auth_mode import ( + classify_auth_mode, + classify_client, + should_stamp_codex_client, +) from headroom.proxy.compression_decision import CompressionDecision from headroom.proxy.cost import _summarize_transforms, header_safe_transforms from headroom.proxy.outcome import RequestOutcome @@ -3498,17 +3502,22 @@ class OpenAIHandlerMixin: # Forward client headers to upstream, adding required OpenAI-Beta header ws_headers = dict(websocket.headers) + _ws_url_obj = getattr(websocket, "url", None) + _ws_url = str(_ws_url_obj) if _ws_url_obj is not None else "" + _ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else "" + if not _ws_path: + _ws_path = "/v1/responses" + # WS sessions bypass the HTTP middleware that stamps X-Client: codex on + # the Responses endpoint, so apply the same path-based stamp here before + # classify_client runs (parallels server.py / should_stamp_codex_client). + if should_stamp_codex_client(_ws_path, ws_headers): + ws_headers["x-client"] = "codex" # Identify the WS harness before downstream auth/header rewrites. # Captured in closure so per-turn RequestOutcome can stamp it. client = classify_client(ws_headers) # WS sessions bypass the HTTP middleware, so bind the project here; # per-turn outcome emission inside this task inherits the context. set_current_project(classify_project(ws_headers)) - _ws_url_obj = getattr(websocket, "url", None) - _ws_url = str(_ws_url_obj) if _ws_url_obj is not None else "" - _ws_path = getattr(_ws_url_obj, "path", "") if _ws_url_obj is not None else "" - if not _ws_path: - _ws_path = "/v1/responses" metrics_for_inbound_ws = getattr(self, "metrics", None) if metrics_for_inbound_ws is not None and hasattr( metrics_for_inbound_ws, "record_inbound_request" diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index d4c147fc7..8b732fbaf 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -103,6 +103,7 @@ from headroom.providers.registry import ( format_backend_status, resolve_api_targets, ) +from headroom.proxy.auth_mode import should_stamp_codex_client # ============================================================================= # Extracted modules (re-exported for backward compatibility) @@ -2152,6 +2153,15 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: query = request.url.query headers = dict(request.headers.items()) set_current_project(classify_project(headers) or prefix_project) + # Path-based Codex identification: stamp X-Client: codex on the + # Responses endpoint for callers that don't otherwise classify (e.g. + # Codex Desktop, whose User-Agent isn't a known codex UA). Without it + # the backend refuses oversized + # requests with a 413 on a compression timeout, which Codex treats as a + # hard connection failure. Mutating scope["headers"] before call_next + # makes every downstream classify_client(headers) read "codex". + if should_stamp_codex_client(path, headers): + request.scope["headers"].append((b"x-client", b"codex")) client = getattr(request, "client", None) client_addr = "" if client is not None: diff --git a/tests/test_codex_client_stamp.py b/tests/test_codex_client_stamp.py new file mode 100644 index 000000000..3925f9960 --- /dev/null +++ b/tests/test_codex_client_stamp.py @@ -0,0 +1,56 @@ +"""Tests for ``should_stamp_codex_client`` — the path-based ``X-Client: codex`` +stamp on the Responses endpoint. + +The stamp fires only for an unidentified caller on the Responses endpoint, so +Codex Desktop (whose User-Agent isn't a known codex UA) takes the codex +fail-open path instead of being refused with a 413 on a compression timeout. +""" + +from __future__ import annotations + +from headroom.proxy.auth_mode import classify_client, should_stamp_codex_client + +CODEX_DESKTOP_UA = ( + "Codex Desktop/0.140.0-alpha.2 (Mac OS 15.7.7; arm64) unknown (Codex Desktop; 26.609.71450)" +) + + +def test_unidentified_codex_desktop_on_responses_is_stamped() -> None: + assert should_stamp_codex_client("/v1/responses", {"user-agent": CODEX_DESKTOP_UA}) + + +def test_stamp_then_classify_yields_codex() -> None: + # End-to-end of what the HTTP middleware and the WS handler both do: + # stamp the header, after which classify_client must read "codex". + headers = {"user-agent": CODEX_DESKTOP_UA} + assert should_stamp_codex_client("/v1/responses", headers) + headers["x-client"] = "codex" + assert classify_client(headers) == "codex" + + +def test_no_user_agent_on_responses_is_stamped() -> None: + assert should_stamp_codex_client("/v1/responses", {}) + + +def test_responses_subpath_is_stamped() -> None: + assert should_stamp_codex_client("/v1/responses/foo", {"user-agent": CODEX_DESKTOP_UA}) + + +def test_other_path_is_not_stamped() -> None: + # Scoped to the Responses endpoint; unknown callers elsewhere are untouched. + assert not should_stamp_codex_client("/v1/chat/completions", {"user-agent": CODEX_DESKTOP_UA}) + + +def test_recognized_non_codex_client_is_not_stamped() -> None: + assert not should_stamp_codex_client("/v1/responses", {"user-agent": "claude-code/1.2.3"}) + + +def test_recognized_codex_cli_is_not_stamped() -> None: + # Already classifies as codex via UA; no stamp needed. + assert not should_stamp_codex_client("/v1/responses", {"user-agent": "codex-cli/0.5"}) + + +def test_explicit_x_client_is_not_stamped() -> None: + assert not should_stamp_codex_client( + "/v1/responses", {"x-client": "aider", "user-agent": CODEX_DESKTOP_UA} + ) diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 2db440d45..d0a202063 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -853,3 +853,29 @@ async def test_ws_upstream_connect_allows_large_frames_and_no_pong_deadline(): assert captured.get("max_size") is None, "upstream frame size must be uncapped" assert captured.get("ping_timeout") is None, "upstream must not impose a pong deadline" + + +@pytest.mark.asyncio +async def test_ws_recognized_client_with_real_path_is_not_restamped(): + """A WS caller that already classifies on a real request path is not stamped.""" + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket(frames=[_first_frame()]) + # A non-empty url path (so the handler does not fall back to the default) + # and a recognized codex UA (so should_stamp_codex_client returns False). + client_ws.url = SimpleNamespace(path="/v1/responses") + client_ws.headers = {"authorization": "Bearer test", "user-agent": "codex-cli/0.5"} + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + # The forwarded handshake headers must not carry a proxy-injected x-client: + # the caller already self-identifies via its User-Agent. + assert "x-client" not in {k.lower() for k in client_ws.headers} + assert handler.ws_sessions.active_count() == 0 diff --git a/tests/test_proxy_codex_route_aliases.py b/tests/test_proxy_codex_route_aliases.py index cf8469fa6..641ec6c75 100644 --- a/tests/test_proxy_codex_route_aliases.py +++ b/tests/test_proxy_codex_route_aliases.py @@ -248,3 +248,53 @@ def test_codex_model_metadata_fetches_codex_registry_for_chatgpt_auth(monkeypatc # Unknown model variants 404 against the dynamic registry. assert unknown_response.status_code == 404 + + +_CODEX_DESKTOP_UA = ( + "Codex Desktop/0.140.0-alpha.2 (Mac OS 15.7.7; arm64) unknown (Codex Desktop; 26.609.71450)" +) + + +def test_responses_middleware_stamps_x_client_codex_for_unidentified_caller(monkeypatch): + # Codex Desktop's User-Agent isn't a known codex UA, so the HTTP middleware + # must stamp X-Client: codex on /v1/responses before the handler classifies + # the caller — otherwise a compression timeout is refused with a 413 that + # Codex treats as a hard connection failure. + seen: dict[str, str | None] = {} + + async def fake_handle(self, request): # type: ignore[no-untyped-def] + seen["x-client"] = request.headers.get("x-client") + return JSONResponse({"ok": True}) + + monkeypatch.setattr(HeadroomProxy, "handle_openai_responses", fake_handle) + + with TestClient(create_app(ProxyConfig())) as client: + response = client.post( + "/v1/responses", + headers={"user-agent": _CODEX_DESKTOP_UA}, + json={"model": "gpt-5.3-codex"}, + ) + + assert response.status_code == 200 + assert seen["x-client"] == "codex" + + +def test_responses_middleware_preserves_explicit_x_client(monkeypatch): + # A caller that already self-identifies is left untouched by the stamp. + seen: dict[str, str | None] = {} + + async def fake_handle(self, request): # type: ignore[no-untyped-def] + seen["x-client"] = request.headers.get("x-client") + return JSONResponse({"ok": True}) + + monkeypatch.setattr(HeadroomProxy, "handle_openai_responses", fake_handle) + + with TestClient(create_app(ProxyConfig())) as client: + response = client.post( + "/v1/responses", + headers={"x-client": "aider", "user-agent": _CODEX_DESKTOP_UA}, + json={"model": "gpt-5.3-codex"}, + ) + + assert response.status_code == 200 + assert seen["x-client"] == "aider"