fix(proxy): stamp X-Client: codex on Responses endpoint for unidentified callers (#1036)

## Description

Codex Desktop (OpenAI's Codex GUI/IDE app) sends a `User-Agent` of the
form `Codex Desktop/<ver> (...)`, which is not in `CLIENT_UA_MAP`, so
`classify_client` returns `None`. On a compression timeout the backend
only takes the codex fail-open path when the client classifies as
`codex`; for an unidentified client it refuses with HTTP 413
(`compression_refused`), which Codex treats as a hard connection
failure.

This stamps `X-Client: codex` on requests to the Responses endpoint
(`/v1/responses`) only when the caller does not otherwise classify. The
stamp is scoped to the Responses endpoint and skipped for any caller
that already classifies through a recognized user-agent or explicit
`X-Client`, so non-Codex traffic is not relabeled.

## Type of Change

- [x] Bug fix (non-breaking change that fixes an issue)
- [x] Tests only

## Changes Made

- Added `should_stamp_codex_client(path, headers)` in
`headroom.proxy.auth_mode` for narrow Responses-endpoint client
stamping.
- Applied the stamp in HTTP middleware before downstream request
classification.
- Applied the same stamp in the Responses WebSocket handler, which
bypasses HTTP middleware.
- Added unit coverage for the stamp/skip matrix, including Codex
Desktop, explicit clients, recognized user-agents, and WebSocket
behavior.
- Merged current `main` and kept both the new stamp coverage and the
existing Codex WebSocket image-generation regression coverage.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Formatting passes (`ruff format --check`)
- [x] New tests added for new functionality

### Test Output

```text
python -m pytest tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py -q
48 passed in 1.13s

ruff check headroom/proxy/auth_mode.py headroom/proxy/server.py headroom/proxy/handlers/openai.py tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py
All checks passed!

ruff format --check headroom/proxy/auth_mode.py headroom/proxy/server.py headroom/proxy/handlers/openai.py tests/test_codex_client_stamp.py tests/test_auth_mode.py tests/test_openai_codex_ws_lifecycle.py
6 files already formatted
```

## Real Behavior Proof

- Environment: local Windows 11 development checkout, Python 3.13.13,
branch updated from `upstream/main`.
- Exact command / steps: Ran the focused unit suite for the new
client-stamp behavior and the overlapping Codex WebSocket lifecycle
tests, then ran `ruff check` and `ruff format --check` on the changed
modules and tests.
- Observed result: The focused suite passed with 48 tests, lint passed,
and formatting passed. The tests assert that unidentified
`/v1/responses` callers classify as Codex after stamping while explicit
or already-recognized clients are preserved.
- Not tested: a live end-to-end Codex Desktop session through a running
`headroom wrap codex` instance; verification is at the unit/integration
boundary for classification and request routing.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>

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",
 ]

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"

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:

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}
+    )

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

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"