fix(proxy): stop discarding compressed Codex WS later-frame payloads (#2823)
## Description `headroom perf` reports 0 tokens saved for Codex CLI sessions despite real traffic being processed (confirmed via the reporter's live proxy stats in the issue). Root cause: a misplaced `return` statement in the Codex WS later-frame compression path silently discards every compressed payload and skips all token/savings bookkeeping for it. Closes #2819 ## 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 - `headroom/proxy/handlers/openai.py` (`_maybe_compress_response_create_frame`): PR #1579 (2026-07-16) moved a `return (raw_after_store, ...)` statement to the same indentation as the enclosing `except Exception:` block instead of inside it. That made the `return` fire **unconditionally** after every later (2nd+) `response.create` frame in a Codex WS session — success or failure — always forwarding the original pre-compression frame upstream and skipping the entire success-path code below it (correct rewritten-payload return, `tokens_saved`, `attempted_input_tokens_total`, `ws_frames_compressed`). Fixed by moving the `return` back inside the `except` block, restoring the success path. - `tests/test_openai_codex_ws_lifecycle.py`: new regression test `test_ws_later_frame_compression_is_actually_forwarded` — mocks the compressor to report `modified=True` with a distinct rewritten payload on a later frame, asserts the rewritten payload (not the original) is what's actually sent upstream. ## Testing - [x] Unit tests pass (`pytest`) - [ ] Linting passes (`ruff check .`) — not run locally, will confirm via CI - [ ] Type checking passes (`mypy headroom`) — not run locally, will confirm via CI - [x] New tests added for new functionality - [x] Manual testing performed (see Real Behavior Proof) ### Test Output ```text $ .venv/Scripts/python -m pytest tests/test_openai_codex_ws_lifecycle.py tests/test_codex_ws_per_frame_memory.py tests/test_codex_ws_savings_deferral.py tests/test_openai_codex_ws_timings.py -q ............................................... 48 passed in 5.34s $ .venv/Scripts/python -m pytest tests/ -k "openai or codex" -q (wider sweep, unrelated dirs excluded) 968 passed, 3 failed, 77 skipped, 2 errors in 483.21s ``` The 3 failures (`test_client_integration.py::test_auto_detect_openai_optimizer`, `test_compress_route_tokenizer_by_model.py::...[deepseek/deepseek-v4]`, `test_learn/test_integration.py::TestCodexIntegration::test_full_pipeline`) reproduce identically on a clean, unmodified `main` — confirmed by stashing this PR's changes and re-running. They're local-environment issues (a live litellm 503, and this dev box's tool registry missing a `Bash` entry), not caused by this change. ## Real Behavior Proof - Environment: Windows 11, Python 3.14.5, local venv, `headroom._core` rebuilt via `maturin develop --release` against current `main` to rule out stale-build noise - Exact command / steps: (1) `git blame` on the buggy block traced the misplaced `return` to commit `1c50eca8` (PR #1579); (2) wrote a regression test that scripts two `response.create` frames over a fake Codex WS session, with the compressor mock returning `modified=False` for frame 1 and `modified=True` (with a distinct payload) for frame 2; (3) ran the test against the pre-fix code (`git stash` isolating just the source fix, keeping the test) — **failed**, `upstream.sent[-1]` was the untouched original frame; (4) ran the test against the fix — **passed**, `upstream.sent[-1]` is the compressed payload; (5) added a further regression test for the later-frame non-timeout-exception path Codecov flagged as uncovered, confirmed `pytest tests/test_openai_codex_ws_lifecycle.py -q` passes 32/32 - Observed result: confirmed the bug exists and the fix resolves it, at the unit level - Not tested: have not reproduced the full `headroom wrap codex` → `headroom perf` end-to-end flow against a live Codex CLI session (no access to Codex CLI / real OpenAI credentials in this environment) — root cause and fix are verified at the code-path level via the regression test above, not via the reporter's exact repro steps ## 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 - [ ] I have made corresponding changes to the documentation (N/A — internal bugfix, no user-facing behavior/docs change beyond "compression now works as originally intended") - [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 - [ ] I have updated the CHANGELOG.md if applicable (release-please generates this automatically from commit messages) ## Additional Notes **Bonus finding, not just a metrics bug**: because the compressed payload was discarded and the original was always sent, this also means Codex WS sessions with multiple turns were silently getting **zero compression benefit** past the first `response.create` frame — not just wrong dashboards. The fix restores actual compression for those turns, not only correct accounting of it.
This commit is contained in:
Ashish Patel
@@ -7086,6 +7086,11 @@ class OpenAIHandlerMixin:
|
||||
frame_type="response.create",
|
||||
model=str(inner_payload.get("model") or "unknown"),
|
||||
)
|
||||
return (
|
||||
raw_after_store,
|
||||
store_forced,
|
||||
"chatgpt_store_false" if store_forced else "compression_exception",
|
||||
)
|
||||
# Record transform labels even when the frame bytes are
|
||||
# unchanged: control-arm output-shaper labels
|
||||
# (output_shaper:control:*) must reach the outcome
|
||||
@@ -7093,11 +7098,6 @@ class OpenAIHandlerMixin:
|
||||
for t in frame_transforms:
|
||||
if t not in transforms_applied:
|
||||
transforms_applied.append(t)
|
||||
return (
|
||||
raw_after_store,
|
||||
store_forced,
|
||||
"chatgpt_store_false" if store_forced else "compression_exception",
|
||||
)
|
||||
if not modified:
|
||||
reason = frame_reason or "no_compression"
|
||||
_log_ws_passthrough(
|
||||
|
||||
@@ -588,6 +588,131 @@ async def test_ws_first_frame_non_timeout_exception_keeps_generic_reason(
|
||||
assert "reason=compression_exception" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_later_frame_compression_is_actually_forwarded(monkeypatch):
|
||||
"""Regression for issue #2819: a later (2nd+) Codex WS response.create
|
||||
frame whose compressor reports ``modified=True`` must have the REWRITTEN
|
||||
payload sent upstream — not the original raw frame.
|
||||
|
||||
A misplaced ``return`` (introduced in #1579) sat at the same indentation
|
||||
as the surrounding ``except`` block, so it fired unconditionally after
|
||||
every later-frame compression attempt — success or failure — and always
|
||||
forwarded ``raw_after_store`` (the pre-compression frame). Compressed
|
||||
later frames were silently discarded on the wire, and the token/savings
|
||||
accounting that only runs on the (dead) success path never accumulated,
|
||||
which is why ``headroom perf`` showed 0 tokens for Codex sessions with
|
||||
multiple turns.
|
||||
"""
|
||||
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,
|
||||
disconnect_after_n_sends=None,
|
||||
)
|
||||
handler = _DummyOpenAIHandler()
|
||||
handler.config.optimize = True
|
||||
monkeypatch.setattr(openai_module, "COMPRESSION_TIMEOUT_SECONDS", 30.0)
|
||||
|
||||
compressed_inner = {"model": "gpt-5.4", "input": "compressed"}
|
||||
calls = 0
|
||||
|
||||
def _compress(payload, *, model, request_id, timing=None, client=None):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
if calls == 1:
|
||||
# First frame: not modified (exercises the other call site).
|
||||
return payload, False, 0, [], "router_no_compression", 10, 10, 0
|
||||
# Later frame: compressor DID find savings.
|
||||
return compressed_inner, True, 5, ["text"], "compressed", 10, 5, 10
|
||||
|
||||
async def _trigger() -> None:
|
||||
await asyncio.sleep(0.05)
|
||||
client_ws.trigger_disconnect()
|
||||
|
||||
handler._compress_openai_responses_payload = _compress # type: ignore[method-assign]
|
||||
|
||||
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
|
||||
|
||||
# The compressed payload must reach upstream for the later frame — not
|
||||
# the untouched original second_frame.
|
||||
assert upstream.sent[-1] != second_frame
|
||||
assert json.loads(upstream.sent[-1])["response"] == compressed_inner
|
||||
|
||||
# The success-path bookkeeping (tokens_saved / frame count) must run —
|
||||
# proof the "modified" branch executed rather than short-circuiting.
|
||||
modified_frames = [frame for frame in handler.metrics.codex_ws_frames if frame.get("modified")]
|
||||
assert modified_frames, "expected at least one frame recorded as modified=True"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ws_later_frame_non_timeout_exception_falls_back_to_original(caplog, monkeypatch):
|
||||
"""A non-timeout compression exception on a later frame must forward the
|
||||
original frame via the except-block return (the line this PR moved back
|
||||
inside the except), not fall through to the (now correctly gated)
|
||||
success-path handling below it.
|
||||
"""
|
||||
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)
|
||||
|
||||
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 RuntimeError("simulated later-frame compression failure")
|
||||
return fn()
|
||||
|
||||
def _noop_compress(payload, *, model, request_id, timing=None, client=None):
|
||||
return payload, False, 0, [], "test_noop", 10, 10, 0
|
||||
|
||||
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
|
||||
|
||||
# The failed later frame must forward the original, unmodified frame.
|
||||
assert upstream.sent[-1] == second_frame
|
||||
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."""
|
||||
|
||||
Reference in New Issue
Block a user