From 5d3803a21c53907e2fea900524e48b510dd59d7a Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Mon, 29 Jun 2026 11:00:05 -0400 Subject: [PATCH] fix(proxy): strip Codex lite header from OpenAI WebSockets (#1543) ## Description Codex WebSocket traffic through Headroom can forward `X-OpenAI-Internal-Codex-Responses-Lite` upstream. OpenAI tightened enforcement of that header on 2026-06-26 for `gpt-5.5`, `gpt-5.4`, and `gpt-5.4-mini`, so the same Codex setup can fail through Headroom with `unsupported_value` while succeeding when Headroom is bypassed. The OpenAI Responses WS handler strips only `x-headroom-*` internal headers today, so this Codex client header survives into both the direct upstream WebSocket connect and the WS HTTP fallback path. This change strips `X-OpenAI-Internal-Codex-Responses-Lite` from the upstream header copy inside `handle_openai_responses_ws` after routing resolution and before the upstream request is sent. `_ws_http_fallback(...)` reuses that same header dict, so the fallback path inherits the fix without a second guard. `headroom/proxy/helpers.py` stays unchanged; the shared helper contract remains `x-headroom-*` only. Closes #1525 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Add a narrow case-insensitive strip for `X-OpenAI-Internal-Codex-Responses-Lite` in `headroom/proxy/handlers/openai.py` after `_resolve_codex_routing_headers(...)` and before `websockets.connect(...)`. - Keep `_strip_internal_headers(...)` in `headroom/proxy/helpers.py` unchanged so the documented `x-headroom-*` stripping scope does not widen. - Extend `tests/test_openai_codex_ws_lifecycle.py` to capture `additional_headers`, prove the direct WS leak on base, prove the fix on head, prove `_ws_http_fallback(...)` inherits sanitized headers, and prove adjacent non-lite headers still survive. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_openai_codex_ws_lifecycle.py -k "codex_responses_lite or without_codex_lite" -v`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text Base proof before the fix: uv run pytest tests/test_openai_codex_ws_lifecycle.py -k "test_ws_codex_responses_lite_header_is_not_forwarded_upstream" -v FAILED tests/test_openai_codex_ws_lifecycle.py::test_ws_codex_responses_lite_header_is_not_forwarded_upstream AssertionError: assert 'X-OpenAI-Internal-Codex-Responses-Lite' not in { 'authorization': 'Bearer test', 'X-OpenAI-Internal-Codex-Responses-Lite': 'true', 'X-OpenAI-Debug': 'keep-me', 'ChatGPT-Account-ID': 'acct-123', 'x-client': 'codex', 'OpenAI-Beta': 'responses_websockets=2026-02-06' } Focused regression proof after the fix: uv run pytest tests/test_openai_codex_ws_lifecycle.py -k "codex_responses_lite or without_codex_lite" -v tests/test_openai_codex_ws_lifecycle.py::test_ws_codex_responses_lite_header_is_not_forwarded_upstream PASSED [ 33%] tests/test_openai_codex_ws_lifecycle.py::test_ws_codex_responses_lite_header_is_not_forwarded_to_fallback PASSED [ 66%] tests/test_openai_codex_ws_lifecycle.py::test_ws_without_codex_lite_preserves_adjacent_headers_and_api_key_route PASSED [100%] ====================== 3 passed, 16 deselected in 0.62s ======================= uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py All checks passed! ``` ## Real Behavior Proof - Environment: local pytest async lifecycle harness in `tests/test_openai_codex_ws_lifecycle.py`, no live provider required. - Exact command / steps: `uv run pytest tests/test_openai_codex_ws_lifecycle.py -k "codex_responses_lite or without_codex_lite" -v` - Observed result: before the fix, the new direct-leak test failed because `websockets.connect(..., additional_headers=...)` still contained `X-OpenAI-Internal-Codex-Responses-Lite`. After the fix, the focused rerun passed and proved that both direct WS connect and forced `_ws_http_fallback(...)` receive sanitized headers while adjacent non-lite headers still survive. - Not tested: live Codex traffic against OpenAI with real credentials, unless that is added during implementation. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes `uv run mypy headroom` is outside the focused proof for this small WS-path fix and may remain unchecked if the coding pass keeps the validation surface to targeted pytest plus ruff. `CHANGELOG.md` remains unchanged because the resolved repo config says Headroom's release pipeline generates changelog entries from conventional commits. --- headroom/proxy/handlers/openai.py | 7 ++ tests/test_openai_codex_ws_lifecycle.py | 111 +++++++++++++++++++++++- 2 files changed, 117 insertions(+), 1 deletion(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 720eb722c..caa8f9539 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -69,6 +69,7 @@ _OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock() _OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None _WS_ALLOWED_ORIGINS_ENV = "HEADROOM_WS_ORIGINS" _CORS_ALLOWED_ORIGINS_ENV = "HEADROOM_CORS_ORIGINS" +_CODEX_RESPONSES_LITE_HEADER = "x-openai-internal-codex-responses-lite" def _header_get(headers: dict[str, str], name: str) -> str | None: @@ -3780,6 +3781,12 @@ class OpenAIHandlerMixin: ) upstream_headers, is_chatgpt_auth = _resolve_codex_routing_headers(upstream_headers) + # OpenAI rejects newer Codex models when this client-only lite header leaks upstream. + upstream_headers = { + key: value + for key, value in upstream_headers.items() + if key.lower() != _CODEX_RESPONSES_LITE_HEADER + } _lower_headers = {k.lower(): v for k, v in upstream_headers.items()} # Build upstream WebSocket URL based on auth mode diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index d0a202063..059811903 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -119,11 +119,12 @@ class _FakeWebSocket: self, frames: list[str] | None = None, *, + headers: dict[str, str] | None = None, disconnect_after_n_sends: int | None = None, hold_after_initial: bool = False, call_log: list[str] | None = None, ) -> None: - self.headers = {"authorization": "Bearer test"} + self.headers = dict(headers or {"authorization": "Bearer test"}) self._frames = list(frames or []) self._hold_after_initial = hold_after_initial self._disconnect_after_n_sends = disconnect_after_n_sends @@ -253,6 +254,7 @@ def _make_fake_websockets_module( upstream: _FakeUpstream | None, *, call_log: list[str] | None = None, + connect_calls: list[tuple[tuple, dict]] | None = None, connect_error: Exception | None = None, ): """Build a fake ``websockets`` module. @@ -267,6 +269,8 @@ def _make_fake_websockets_module( async def _connect(*args, **kwargs): if call_log is not None: call_log.append("connect") + if connect_calls is not None: + connect_calls.append((args, dict(kwargs))) if connect_error is not None: raise connect_error return upstream @@ -285,6 +289,17 @@ def _first_frame() -> str: ) +def _codex_lite_headers(*, chatgpt: bool) -> dict[str, str]: + headers = { + "authorization": "Bearer test", + "X-OpenAI-Internal-Codex-Responses-Lite": "true", + "X-OpenAI-Debug": "keep-me", + } + if chatgpt: + headers["ChatGPT-Account-ID"] = "acct-123" + return headers + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -675,6 +690,100 @@ async def test_ws_connect_failure_falls_back_to_http(): assert handler.ws_sessions.active_count() == 0 +@pytest.mark.asyncio +async def test_ws_codex_responses_lite_header_is_not_forwarded_upstream(): + """The WS upstream handshake must drop the Codex lite header only.""" + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + connect_calls: list[tuple[tuple, dict]] = [] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream, connect_calls=connect_calls) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + headers=_codex_lite_headers(chatgpt=True), + ) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(connect_calls) == 1 + connect_args, connect_kwargs = connect_calls[0] + assert connect_args[0] == "wss://chatgpt.com/backend-api/codex/responses" + forwarded_headers = connect_kwargs["additional_headers"] + assert "X-OpenAI-Internal-Codex-Responses-Lite" not in forwarded_headers + assert forwarded_headers["ChatGPT-Account-ID"] == "acct-123" + assert forwarded_headers["X-OpenAI-Debug"] == "keep-me" + + +@pytest.mark.asyncio +async def test_ws_codex_responses_lite_header_is_not_forwarded_to_fallback(): + """HTTP fallback must inherit the sanitized upstream header copy.""" + fake_ws_mod = _make_fake_websockets_module( + None, + connect_error=RuntimeError("HTTP 500 from upstream"), + ) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + headers=_codex_lite_headers(chatgpt=True), + ) + handler = _DummyOpenAIHandler() + + fallback_calls: list[dict[str, str]] = [] + + async def _fallback(websocket, body, first_msg_raw, upstream_headers, request_id): + fallback_calls.append(dict(upstream_headers)) + + handler._ws_http_fallback = _fallback # type: ignore[assignment] + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(fallback_calls) == 1 + forwarded_headers = fallback_calls[0] + assert "X-OpenAI-Internal-Codex-Responses-Lite" not in forwarded_headers + assert forwarded_headers["ChatGPT-Account-ID"] == "acct-123" + assert forwarded_headers["X-OpenAI-Debug"] == "keep-me" + + +@pytest.mark.asyncio +async def test_ws_without_codex_lite_preserves_adjacent_headers_and_api_key_route(): + """Requests without the lite header keep adjacent OpenAI headers intact.""" + upstream_events = [ + json.dumps({"type": "response.created", "response": {"id": "r_1"}}), + json.dumps({"type": "response.completed", "response": {"id": "r_1"}}), + ] + connect_calls: list[tuple[tuple, dict]] = [] + upstream = _FakeUpstream(upstream_events) + fake_ws_mod = _make_fake_websockets_module(upstream, connect_calls=connect_calls) + + client_ws = _FakeWebSocket( + frames=[_first_frame()], + headers={ + "authorization": "Bearer test", + "OpenAI-Beta": "responses=v1", + "X-OpenAI-Debug": "keep-me", + }, + ) + handler = _DummyOpenAIHandler() + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert len(connect_calls) == 1 + connect_args, connect_kwargs = connect_calls[0] + assert connect_args[0] == "wss://api.openai.com/v1/responses" + forwarded_headers = connect_kwargs["additional_headers"] + assert "responses=v1" in forwarded_headers["OpenAI-Beta"] + assert "responses_websockets=2026-02-06" in forwarded_headers["OpenAI-Beta"] + assert forwarded_headers["X-OpenAI-Debug"] == "keep-me" + assert "ChatGPT-Account-ID" not in forwarded_headers + + @pytest.mark.asyncio async def test_ws_connect_happens_before_accept(): """The upstream connect must complete before the client 101 is sent,