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.
This commit is contained in:
Rod Boev 2026-07-05 17:02:52 -04:00 committed by GitHub
parent 0a3851b240
commit d24a3f8425
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 203 additions and 10 deletions

View file

@ -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

View file

@ -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."""