Skip to content

Commit d24a3f8

Browse files
rodboevrodboev
authored
fix(proxy): bound Codex WS compression fallback latency (headroomlabs-ai#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 headroomlabs-ai#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.
1 parent 0a3851b commit d24a3f8

2 files changed

Lines changed: 203 additions & 10 deletions

File tree

‎headroom/proxy/handlers/openai.py‎

Lines changed: 53 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,13 @@
6767
_OPENAI_RESPONSES_UNIT_CACHE_INIT_LOCK = threading.RLock()
6868
_OPENAI_RESPONSES_UNIT_EXECUTOR_LOCK = threading.RLock()
6969
_OPENAI_RESPONSES_UNIT_EXECUTOR: ThreadPoolExecutor | None = None
70+
_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS = 5.0
71+
72+
73+
def _codex_ws_compression_timeout_seconds() -> float:
74+
return min(COMPRESSION_TIMEOUT_SECONDS, _CODEX_WS_COMPRESSION_TIMEOUT_SECONDS)
75+
76+
7077
_WS_ALLOWED_ORIGINS_ENV = "HEADROOM_WS_ORIGINS"
7178
_CORS_ALLOWED_ORIGINS_ENV = "HEADROOM_CORS_ORIGINS"
7279
_CODEX_RESPONSES_LITE_HEADER = "x-openai-internal-codex-responses-lite"
@@ -1756,6 +1763,7 @@ async def _compress_openai_responses_payload_in_executor(
17561763
*,
17571764
model: str,
17581765
request_id: str,
1766+
timeout: float = COMPRESSION_TIMEOUT_SECONDS,
17591767
) -> tuple[dict[str, Any], bool, int, list[str], str | None, int, int, int, dict[str, float]]:
17601768
timing: dict[str, float] = {}
17611769

@@ -1778,7 +1786,7 @@ def _compress(): # noqa: ANN202
17781786

17791787
result = await self._run_compression_in_executor(
17801788
_compress,
1781-
timeout=COMPRESSION_TIMEOUT_SECONDS,
1789+
timeout=timeout,
17821790
)
17831791
if len(result) == 8:
17841792
return (*result, timing)
@@ -4864,6 +4872,9 @@ def _prepare_ws_performance_metrics() -> tuple[float, float, dict[str, float]]:
48644872
_inner,
48654873
model=_model,
48664874
request_id=request_id,
4875+
timeout=_codex_ws_compression_timeout_seconds()
4876+
if client == "codex"
4877+
else COMPRESSION_TIMEOUT_SECONDS,
48674878
)
48684879
for _timing_name, _timing_ms in _ws_compression_timing.items():
48694880
_record_ws_compression_timing(_timing_name, _timing_ms)
@@ -4949,12 +4960,14 @@ def _prepare_ws_performance_metrics() -> tuple[float, float, dict[str, float]]:
49494960
bytes_before=_ws_frame_bytes,
49504961
failed=True,
49514962
)
4963+
_timeout_failure = isinstance(_ce, asyncio.TimeoutError)
49524964
logger.warning(
4953-
f"[{request_id}] WS /v1/responses compression failed "
4965+
f"[{request_id}] WS /v1/responses compression "
4966+
f"{'timed out' if _timeout_failure else 'failed'} "
49544967
f"(bytes={_ws_frame_bytes}): {type(_ce).__name__}: {_ce}"
49554968
)
49564969
_log_ws_passthrough(
4957-
"compression_exception",
4970+
"compression_timeout" if _timeout_failure else "compression_exception",
49584971
frame_index=1,
49594972
raw_bytes=_ws_frame_bytes,
49604973
frame_type="response.create" if body else "unknown",
@@ -5175,12 +5188,45 @@ async def _maybe_compress_response_create_frame(
51755188
inner_payload,
51765189
model=model_for_frame,
51775190
request_id=request_id,
5191+
timeout=_codex_ws_compression_timeout_seconds()
5192+
if client == "codex"
5193+
else COMPRESSION_TIMEOUT_SECONDS,
51785194
)
5179-
for (
5180-
_timing_name,
5181-
_timing_ms,
5182-
) in frame_compression_timing.items():
5195+
for _timing_name, _timing_ms in frame_compression_timing.items():
51835196
_record_ws_compression_timing(_timing_name, _timing_ms)
5197+
except asyncio.TimeoutError as _frame_err:
5198+
frame_compression_elapsed_ms = (
5199+
time.perf_counter() - _compression_started
5200+
) * 1000.0
5201+
if frame_compression_elapsed_ms > 0:
5202+
record_frame = getattr(
5203+
getattr(self, "metrics", None),
5204+
"record_codex_ws_frame",
5205+
None,
5206+
)
5207+
if record_frame is not None:
5208+
record_frame(
5209+
elapsed_ms=frame_compression_elapsed_ms,
5210+
bytes_before=len(
5211+
raw_msg.encode("utf-8", errors="replace")
5212+
),
5213+
failed=True,
5214+
)
5215+
logger.warning(
5216+
"[%s] WS /v1/responses frame compression "
5217+
"timed out; forwarding original: %s: %s",
5218+
request_id,
5219+
type(_frame_err).__name__,
5220+
_frame_err,
5221+
)
5222+
_log_ws_passthrough(
5223+
"compression_timeout",
5224+
frame_index=frame_index,
5225+
raw_bytes=len(raw_msg.encode("utf-8", errors="replace")),
5226+
frame_type="response.create",
5227+
model=str(inner_payload.get("model") or "unknown"),
5228+
)
5229+
return raw_msg, False, "compression_timeout"
51845230
finally:
51855231
frame_compression_elapsed_ms = (
51865232
time.perf_counter() - _compression_started

‎tests/test_openai_codex_ws_lifecycle.py‎

Lines changed: 150 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,15 @@
99

1010
import asyncio
1111
import json
12+
import logging
1213
import sys
1314
from types import SimpleNamespace
1415
from unittest.mock import MagicMock, patch
1516

1617
import pytest
1718

19+
import headroom.proxy.handlers.openai as openai_module
1820
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
19-
from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS
2021
from headroom.proxy.ws_session_registry import WebSocketSessionRegistry
2122

2223
# ---------------------------------------------------------------------------
@@ -38,6 +39,7 @@ def __init__(self) -> None:
3839
self.stage_timings: list[tuple[str, dict[str, float]]] = []
3940
self.termination_causes: list[str] = []
4041
self.recorded_requests: list[dict] = []
42+
self.codex_ws_frames: list[dict] = []
4143

4244
async def record_request(self, **kwargs): # pragma: no cover
4345
self.recorded_requests.append(dict(kwargs))
@@ -63,6 +65,9 @@ def record_ws_session_duration(self, duration_ms: float, cause: str) -> None:
6365
self.ws_session_durations.append(duration_ms)
6466
self.termination_causes.append(cause)
6567

68+
def record_codex_ws_frame(self, **kwargs) -> None:
69+
self.codex_ws_frames.append(dict(kwargs))
70+
6671

6772
class _DummyOpenAIHandler(OpenAIHandlerMixin):
6873
OPENAI_API_URL = "https://api.openai.com"
@@ -456,7 +461,7 @@ async def _record_request_outcome(outcome):
456461

457462

458463
@pytest.mark.asyncio
459-
async def test_ws_first_frame_compression_uses_bounded_executor():
464+
async def test_ws_first_frame_compression_uses_bounded_executor(monkeypatch):
460465
"""Codex WS compression must not run synchronously on the event loop."""
461466
upstream_events = [
462467
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
@@ -468,6 +473,12 @@ async def test_ws_first_frame_compression_uses_bounded_executor():
468473
client_ws = _FakeWebSocket(frames=[_first_frame()])
469474
handler = _DummyOpenAIHandler()
470475
handler.config.optimize = True
476+
monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0)
477+
expected_timeout = getattr(
478+
openai_module,
479+
"_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS",
480+
5.0,
481+
)
471482
handler._compress_openai_responses_payload = MagicMock(
472483
return_value=(
473484
{"model": "gpt-5.4", "input": "hi"},
@@ -484,10 +495,146 @@ async def test_ws_first_frame_compression_uses_bounded_executor():
484495
await handler.handle_openai_responses_ws(client_ws)
485496

486497
assert handler.compression_executor_calls == 1
487-
assert handler.compression_executor_timeouts == [COMPRESSION_TIMEOUT_SECONDS]
498+
assert handler.compression_executor_timeouts == [expected_timeout]
488499
handler._compress_openai_responses_payload.assert_called_once()
489500

490501

502+
@pytest.mark.asyncio
503+
async def test_ws_first_frame_timeout_uses_timeout_reason(caplog, monkeypatch):
504+
"""Codex WS compression timeout must stay bounded and visible."""
505+
upstream_events = [
506+
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
507+
json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
508+
]
509+
upstream = _FakeUpstream(upstream_events)
510+
fake_ws_mod = _make_fake_websockets_module(upstream)
511+
512+
client_ws = _FakeWebSocket(frames=[_first_frame()])
513+
handler = _DummyOpenAIHandler()
514+
handler.config.optimize = True
515+
monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0)
516+
monkeypatch.setattr(
517+
openai_module,
518+
"_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS",
519+
0.01,
520+
raising=False,
521+
)
522+
523+
async def _timeout_run(fn, *, timeout: float):
524+
handler.compression_executor_calls += 1
525+
handler.compression_executor_timeouts.append(timeout)
526+
raise asyncio.TimeoutError("simulated timeout")
527+
528+
handler._run_compression_in_executor = _timeout_run # type: ignore[method-assign]
529+
caplog.set_level(logging.INFO, logger="headroom.proxy")
530+
531+
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
532+
await handler.handle_openai_responses_ws(client_ws)
533+
534+
assert handler.compression_executor_timeouts == [0.01]
535+
assert "reason=compression_timeout" in caplog.text
536+
537+
538+
@pytest.mark.asyncio
539+
async def test_ws_first_frame_non_timeout_exception_keeps_generic_reason(
540+
caplog,
541+
monkeypatch,
542+
):
543+
"""Codex WS non-timeout compression failures still log the generic reason."""
544+
upstream_events = [
545+
json.dumps({"type": "response.created", "response": {"id": "r_1"}}),
546+
json.dumps({"type": "response.completed", "response": {"id": "r_1"}}),
547+
]
548+
upstream = _FakeUpstream(upstream_events)
549+
fake_ws_mod = _make_fake_websockets_module(upstream)
550+
551+
client_ws = _FakeWebSocket(frames=[_first_frame()])
552+
handler = _DummyOpenAIHandler()
553+
handler.config.optimize = True
554+
monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0)
555+
monkeypatch.setattr(
556+
openai_module,
557+
"_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS",
558+
0.01,
559+
raising=False,
560+
)
561+
562+
async def _error_run(fn, *, timeout: float):
563+
handler.compression_executor_calls += 1
564+
handler.compression_executor_timeouts.append(timeout)
565+
raise RuntimeError("simulated failure")
566+
567+
handler._run_compression_in_executor = _error_run # type: ignore[method-assign]
568+
caplog.set_level(logging.INFO, logger="headroom.proxy")
569+
570+
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
571+
await handler.handle_openai_responses_ws(client_ws)
572+
573+
assert handler.compression_executor_timeouts == [0.01]
574+
assert "reason=compression_exception" in caplog.text
575+
576+
577+
@pytest.mark.asyncio
578+
async def test_ws_later_frame_timeout_records_failed_frame(caplog, monkeypatch):
579+
"""Later Codex WS compression timeout records failed frame metrics."""
580+
second_frame = _first_frame()
581+
upstream = _FakeUpstream([], hold_after_events=True)
582+
fake_ws_mod = _make_fake_websockets_module(upstream)
583+
584+
client_ws = _FakeWebSocket(
585+
frames=[_first_frame(), second_frame],
586+
hold_after_initial=True,
587+
)
588+
handler = _DummyOpenAIHandler()
589+
handler.config.optimize = True
590+
monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0)
591+
monkeypatch.setattr(
592+
openai_module,
593+
"_CODEX_WS_COMPRESSION_TIMEOUT_SECONDS",
594+
0.01,
595+
raising=False,
596+
)
597+
598+
def _noop_compress(payload, *, model, request_id, timing=None):
599+
return payload, False, 0, [], "test_noop", 10, 10, 0
600+
601+
calls = 0
602+
603+
async def _run(fn, *, timeout: float):
604+
nonlocal calls
605+
calls += 1
606+
handler.compression_executor_calls += 1
607+
handler.compression_executor_timeouts.append(timeout)
608+
if calls == 2:
609+
raise asyncio.TimeoutError("simulated later-frame timeout")
610+
return fn()
611+
612+
async def _trigger() -> None:
613+
await asyncio.sleep(0.05)
614+
client_ws.trigger_disconnect()
615+
616+
handler._compress_openai_responses_payload = _noop_compress # type: ignore[method-assign]
617+
handler._run_compression_in_executor = _run # type: ignore[method-assign]
618+
caplog.set_level(logging.INFO, logger="headroom.proxy")
619+
620+
with patch.dict(sys.modules, {"websockets": fake_ws_mod}):
621+
trigger_task = asyncio.create_task(_trigger())
622+
try:
623+
await asyncio.wait_for(handler.handle_openai_responses_ws(client_ws), timeout=2.0)
624+
finally:
625+
trigger_task.cancel()
626+
try:
627+
await trigger_task
628+
except asyncio.CancelledError:
629+
pass
630+
631+
failed_frames = [frame for frame in handler.metrics.codex_ws_frames if frame.get("failed")]
632+
assert handler.compression_executor_timeouts == [0.01, 0.01]
633+
assert upstream.sent[-1] == second_frame
634+
assert failed_frames and failed_frames[-1]["elapsed_ms"] > 0
635+
assert "reason=compression_timeout" in caplog.text
636+
637+
491638
@pytest.mark.asyncio
492639
async def test_happy_path_registry_empty_after_response_completed():
493640
"""Normal session completes — both relay tasks done, registry empty."""

0 commit comments

Comments
 (0)