From d24a3f842551d36c14dc0ec146a9302e256c5c0f Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Sun, 5 Jul 2026 17:02:52 -0400 Subject: [PATCH] fix(proxy): bound Codex WS compression fallback latency (#1802) ## Description Codex `/v1/responses` WebSocket frames could spend the full global compression timeout before falling through unchanged, then report only a generic `compression_exception` reason. That made a recoverable timeout look like an opaque compression failure and left Codex users waiting around 30 seconds for frames that did not produce useful compression. This keeps the existing compression executor, adds a Codex WS-specific compression timeout bound, and records timeout fallback distinctly from other compression exceptions. Closes #922. ## 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 - Bound Codex Responses WebSocket frame compression with a WS-specific timeout. - Pass that timeout through the existing compression executor instead of adding a parallel executor path. - Record timeout passthrough with `compression_timeout` instead of the generic compression exception reason. - Preserve generic `compression_exception` for non-timeout failures. - Add coverage for first-frame timeout bounds, timeout reason logging, generic exception preservation, and later-frame failed metrics. ## Testing - [x] Unit tests pass (`uv run pytest tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_compression_scheduler.py tests/test_compression_observability.py -q`) - [x] Linting passes (`uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_compression_scheduler.py tests/test_compression_observability.py`) - [ ] Type checking passes (`uv run mypy headroom`) - [x] New tests added for new functionality when applicable - [x] Manual testing performed ### Test Output ```text uv run pytest tests/test_openai_codex_ws_lifecycle.py -q 22 passed in 1.38s uv run pytest tests/test_codex_ws_compression_scheduler.py tests/test_compression_observability.py -q 14 passed, 1 skipped in 3.63s uv run pytest tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_compression_scheduler.py tests/test_compression_observability.py -q 36 passed, 1 skipped in 2.09s uv run ruff check headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_compression_scheduler.py tests/test_compression_observability.py All checks passed! uv run ruff format headroom/proxy/handlers/openai.py tests/test_openai_codex_ws_lifecycle.py --check 2 files already formatted ``` ## Real Behavior Proof - Environment: Windows, Python through the project `uv` environment. - Exact command / steps: run the focused Codex WS timeout regressions with small monkeypatched timeout values. - Observed result: base uses the global timeout path or reports only generic `compression_exception`; head passes with Codex WS timeout bounded to the smaller WS cap, logs `compression_timeout` for timeout fallback, preserves `compression_exception` for non-timeout failures, and records failed metrics for later-frame timeout fallback. - Not tested: live Codex Desktop traffic against paid OpenAI credentials. ## 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 - [x] I have commented my code, particularly in hard-to-understand areas - [x] 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 No changelog entry is needed for this request-path bug fix. Type checking was not part of the focused local validation for this Python-only change. Live Codex Desktop validation is not included because the regression is covered at the handler boundary. --- headroom/proxy/handlers/openai.py | 60 ++++++++-- tests/test_openai_codex_ws_lifecycle.py | 153 +++++++++++++++++++++++- 2 files changed, 203 insertions(+), 10 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 480490930..ee39c07e7 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -67,6 +67,13 @@ _OPENAI_RESPONSES_UNIT_PARALLELISM_MAX = 16 _OPENAI_RESPONSES_UNIT_CACHE_INIT_LOCK = threading.RLock() _OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock() _OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None +_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS = 5.0 + + +def _codex_ws_compression_timeout_seconds() -> float: + return min(COMPRESSION_TIMEOUT_SECONDS, _CODEX_WS_COMPRESSION_TIMEOUT_SECONDS) + + _WS_ALLOWED_ORIGINS_ENV = "HEADROOM_WS_ORIGINS" _CORS_ALLOWED_ORIGINS_ENV = "HEADROOM_CORS_ORIGINS" _CODEX_RESPONSES_LITE_HEADER = "x-openai-internal-codex-responses-lite" @@ -1756,6 +1763,7 @@ class OpenAIHandlerMixin: *, model: str, request_id: str, + timeout: float = COMPRESSION_TIMEOUT_SECONDS, ) -> tuple[dict[str, Any], bool, int, list[str], str | None, int, int, int, dict[str, float]]: timing: dict[str, float] = {} @@ -1778,7 +1786,7 @@ class OpenAIHandlerMixin: result = await self._run_compression_in_executor( _compress, - timeout=COMPRESSION_TIMEOUT_SECONDS, + timeout=timeout, ) if len(result) == 8: return (*result, timing) @@ -4864,6 +4872,9 @@ class OpenAIHandlerMixin: _inner, model=_model, request_id=request_id, + timeout=_codex_ws_compression_timeout_seconds() + if client == "codex" + else COMPRESSION_TIMEOUT_SECONDS, ) for _timing_name, _timing_ms in _ws_compression_timing.items(): _record_ws_compression_timing(_timing_name, _timing_ms) @@ -4949,12 +4960,14 @@ class OpenAIHandlerMixin: bytes_before=_ws_frame_bytes, failed=True, ) + _timeout_failure = isinstance(_ce, asyncio.TimeoutError) logger.warning( - f"[{request_id}] WS /v1/responses compression failed " + f"[{request_id}] WS /v1/responses compression " + f"{'timed out' if _timeout_failure else 'failed'} " f"(bytes={_ws_frame_bytes}): {type(_ce).__name__}: {_ce}" ) _log_ws_passthrough( - "compression_exception", + "compression_timeout" if _timeout_failure else "compression_exception", frame_index=1, raw_bytes=_ws_frame_bytes, frame_type="response.create" if body else "unknown", @@ -5175,12 +5188,45 @@ class OpenAIHandlerMixin: inner_payload, model=model_for_frame, request_id=request_id, + timeout=_codex_ws_compression_timeout_seconds() + if client == "codex" + else COMPRESSION_TIMEOUT_SECONDS, ) - for ( - _timing_name, - _timing_ms, - ) in frame_compression_timing.items(): + for _timing_name, _timing_ms in frame_compression_timing.items(): _record_ws_compression_timing(_timing_name, _timing_ms) + except asyncio.TimeoutError as _frame_err: + frame_compression_elapsed_ms = ( + time.perf_counter() - _compression_started + ) * 1000.0 + if frame_compression_elapsed_ms > 0: + record_frame = getattr( + getattr(self, "metrics", None), + "record_codex_ws_frame", + None, + ) + if record_frame is not None: + record_frame( + elapsed_ms=frame_compression_elapsed_ms, + bytes_before=len( + raw_msg.encode("utf-8", errors="replace") + ), + failed=True, + ) + logger.warning( + "[%s] WS /v1/responses frame compression " + "timed out; forwarding original: %s: %s", + request_id, + type(_frame_err).__name__, + _frame_err, + ) + _log_ws_passthrough( + "compression_timeout", + frame_index=frame_index, + raw_bytes=len(raw_msg.encode("utf-8", errors="replace")), + frame_type="response.create", + model=str(inner_payload.get("model") or "unknown"), + ) + return raw_msg, False, "compression_timeout" finally: frame_compression_elapsed_ms = ( time.perf_counter() - _compression_started diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index 10642bd23..dda3203eb 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -9,14 +9,15 @@ from __future__ import annotations import asyncio import json +import logging import sys from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +import headroom.proxy.handlers.openai as openai_module from headroom.proxy.handlers.openai import OpenAIHandlerMixin -from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS from headroom.proxy.ws_session_registry import WebSocketSessionRegistry # --------------------------------------------------------------------------- @@ -38,6 +39,7 @@ class _DummyMetrics: self.stage_timings: list[tuple[str, dict[str, float]]] = [] self.termination_causes: list[str] = [] self.recorded_requests: list[dict] = [] + self.codex_ws_frames: list[dict] = [] async def record_request(self, **kwargs): # pragma: no cover self.recorded_requests.append(dict(kwargs)) @@ -63,6 +65,9 @@ class _DummyMetrics: self.ws_session_durations.append(duration_ms) self.termination_causes.append(cause) + def record_codex_ws_frame(self, **kwargs) -> None: + self.codex_ws_frames.append(dict(kwargs)) + class _DummyOpenAIHandler(OpenAIHandlerMixin): OPENAI_API_URL = "https://api.openai.com" @@ -456,7 +461,7 @@ async def test_ws_output_shaper_holdout_labels_without_rewrite(monkeypatch): @pytest.mark.asyncio -async def test_ws_first_frame_compression_uses_bounded_executor(): +async def test_ws_first_frame_compression_uses_bounded_executor(monkeypatch): """Codex WS compression must not run synchronously on the event loop.""" upstream_events = [ json.dumps({"type": "response.created", "response": {"id": "r_1"}}), @@ -468,6 +473,12 @@ async def test_ws_first_frame_compression_uses_bounded_executor(): client_ws = _FakeWebSocket(frames=[_first_frame()]) handler = _DummyOpenAIHandler() handler.config.optimize = True + monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0) + expected_timeout = getattr( + openai_module, + "_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS", + 5.0, + ) handler._compress_openai_responses_payload = MagicMock( return_value=( {"model": "gpt-5.4", "input": "hi"}, @@ -484,10 +495,146 @@ async def test_ws_first_frame_compression_uses_bounded_executor(): await handler.handle_openai_responses_ws(client_ws) assert handler.compression_executor_calls == 1 - assert handler.compression_executor_timeouts == [COMPRESSION_TIMEOUT_SECONDS] + assert handler.compression_executor_timeouts == [expected_timeout] handler._compress_openai_responses_payload.assert_called_once() +@pytest.mark.asyncio +async def test_ws_first_frame_timeout_uses_timeout_reason(caplog, monkeypatch): + """Codex WS compression timeout must stay bounded and visible.""" + 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()]) + handler = _DummyOpenAIHandler() + handler.config.optimize = True + monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0) + monkeypatch.setattr( + openai_module, + "_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS", + 0.01, + raising=False, + ) + + async def _timeout_run(fn, *, timeout: float): + handler.compression_executor_calls += 1 + handler.compression_executor_timeouts.append(timeout) + raise asyncio.TimeoutError("simulated timeout") + + handler._run_compression_in_executor = _timeout_run # type: ignore[method-assign] + caplog.set_level(logging.INFO, logger="headroom.proxy") + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert handler.compression_executor_timeouts == [0.01] + assert "reason=compression_timeout" in caplog.text + + +@pytest.mark.asyncio +async def test_ws_first_frame_non_timeout_exception_keeps_generic_reason( + caplog, + monkeypatch, +): + """Codex WS non-timeout compression failures still log the generic reason.""" + 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()]) + handler = _DummyOpenAIHandler() + handler.config.optimize = True + monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0) + monkeypatch.setattr( + openai_module, + "_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS", + 0.01, + raising=False, + ) + + async def _error_run(fn, *, timeout: float): + handler.compression_executor_calls += 1 + handler.compression_executor_timeouts.append(timeout) + raise RuntimeError("simulated failure") + + handler._run_compression_in_executor = _error_run # type: ignore[method-assign] + caplog.set_level(logging.INFO, logger="headroom.proxy") + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + await handler.handle_openai_responses_ws(client_ws) + + assert handler.compression_executor_timeouts == [0.01] + assert "reason=compression_exception" in caplog.text + + +@pytest.mark.asyncio +async def test_ws_later_frame_timeout_records_failed_frame(caplog, monkeypatch): + """Later Codex WS compression timeout records failed frame metrics.""" + second_frame = _first_frame() + upstream = _FakeUpstream([], hold_after_events=True) + fake_ws_mod = _make_fake_websockets_module(upstream) + + client_ws = _FakeWebSocket( + frames=[_first_frame(), second_frame], + hold_after_initial=True, + ) + handler = _DummyOpenAIHandler() + handler.config.optimize = True + monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0) + monkeypatch.setattr( + openai_module, + "_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS", + 0.01, + raising=False, + ) + + def _noop_compress(payload, *, model, request_id, timing=None): + return payload, False, 0, [], "test_noop", 10, 10, 0 + + calls = 0 + + async def _run(fn, *, timeout: float): + nonlocal calls + calls += 1 + handler.compression_executor_calls += 1 + handler.compression_executor_timeouts.append(timeout) + if calls == 2: + raise asyncio.TimeoutError("simulated later-frame timeout") + return fn() + + async def _trigger() -> None: + await asyncio.sleep(0.05) + client_ws.trigger_disconnect() + + handler._compress_openai_responses_payload = _noop_compress # type: ignore[method-assign] + handler._run_compression_in_executor = _run # type: ignore[method-assign] + caplog.set_level(logging.INFO, logger="headroom.proxy") + + with patch.dict(sys.modules, {"websockets": fake_ws_mod}): + trigger_task = asyncio.create_task(_trigger()) + try: + await asyncio.wait_for(handler.handle_openai_responses_ws(client_ws), timeout=2.0) + finally: + trigger_task.cancel() + try: + await trigger_task + except asyncio.CancelledError: + pass + + failed_frames = [frame for frame in handler.metrics.codex_ws_frames if frame.get("failed")] + assert handler.compression_executor_timeouts == [0.01, 0.01] + assert upstream.sent[-1] == second_frame + assert failed_frames and failed_frames[-1]["elapsed_ms"] > 0 + assert "reason=compression_timeout" in caplog.text + + @pytest.mark.asyncio async def test_happy_path_registry_empty_after_response_completed(): """Normal session completes — both relay tasks done, registry empty."""