From 039cd2431aaec7d59fefaf7e97aeda1fd7ab3afa Mon Sep 17 00:00:00 2001 From: JD Davis Date: Wed, 12 Aug 2026 12:27:05 -0500 Subject: [PATCH] fix(proxy): preserve merged session and quarantine contracts (#2943) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description Forward-fixes two integration contracts exposed while auditing the large August 12 merge batch on `main`. The Codex WebSocket request-ID hardening correctly gave every emitted dashboard/feed row a unique ID, but it also changed the human-readable `PERF` prefix from the stable WebSocket session ID to that per-emission ID. That broke operator correlation and the contract documented by the original merge. This PR separates storage identity from log correlation: rows remain unique, while `PERF` lines remain grouped under the session ID. The same audit found two tokenizer quarantine tests still modeling the pre-time-cap behavior. Timeout debt no longer activates quarantine after its deadline expires. The tests now establish a live deadline and therefore continue to exercise the intended fail-open branch without weakening the production guard. ## Changes - Add an optional `RequestOutcome.perf_request_id` correlation field, defaulting to the existing `request_id` behavior for all current callers. - Set that field to the stable session ID for both per-turn and residual Codex WebSocket emissions. - Strengthen the lifecycle regression test to prove the unique feed-row ID is not used as the `PERF` prefix. - Update tokenizer quarantine tests to model an active, time-capped quarantine. This is a forward fix; it does not revert the unique WebSocket request IDs or the time-capped quarantine behavior. ## Merge-batch audit context - Audited 48 squash merges from `12149f74466c08b69be8d5fe751425be63c2fda4` through `941c25d31e6c6e0b436c307cbe212771ff76b45f`. - Reviewed repeated-touch hotspots in the OpenAI/Anthropic handlers and proxy server. - Confirmed the batch has no deleted or renamed paths and no revert/supersession commits. - Reproduced the deterministic failures on the merged head and confirmed the corresponding pre-batch tokenizer tests passed; the new WebSocket assertion was introduced by the batch and exposed the dropped correlation contract. ## Validation - `ruff check` — passed - `ruff format --check` — passed - `mypy headroom` — passed (517 source files) - Focused proxy/outcome/tokenizer suite — 99 passed - CI-equivalent Python suite in four `pytest-split` shards — 10,797 passed, 603 skipped - `cargo fmt --all -- --check` — passed - `cargo clippy --workspace -- -D warnings` — passed - New beta-sticky integration suite — 9 passed - Simulator-backed proxy E2E — 7 passed - Real local Headroom process against a real local FastAPI upstream: - `/livez` healthy - `/readyz` ready (Kompress correctly optional/degraded in minimal passthrough mode) - OpenAI `/v1/chat/completions` round trip passed - Anthropic `/v1/messages` round trip passed - Local `act` exercised the push workflow's change-detection job and entered the real lint job. `act` required local no-op cache/setup shims because its runner action post-hooks/pip pairing are incompatible with this macOS Docker environment; the actual Ruff/mypy validations above were run natively and passed. ## Risk Low. The new field is optional and preserves existing behavior by default. Only Codex WebSocket outcome emitters opt into a separate log correlation ID; dashboard/feed row identity remains unchanged. --------- Co-authored-by: Claude Fable 5 --- headroom/proxy/handlers/openai.py | 4 ++++ headroom/proxy/outcome.py | 8 +++++++- tests/test_openai_codex_ws_lifecycle.py | 9 ++++++++- tests/test_tokenizer_count_offload.py | 9 ++++++++- 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index c3e3ff0a3..ce89c0d2c 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -7922,6 +7922,9 @@ class OpenAIHandlerMixin: RequestOutcome( # Per-emission ids keep dashboard request-log keys unique. request_id=await self._next_request_id(), + # PERF remains grouped under the stable WS + # session id even though feed rows are unique. + perf_request_id=request_id, provider="openai", model=model_for_metrics, original_tokens=max(0, input_delta) + max(0, saved_delta), @@ -8498,6 +8501,7 @@ class OpenAIHandlerMixin: RequestOutcome( # Per-emission ids keep dashboard request-log keys unique. request_id=await self._next_request_id(), + perf_request_id=request_id, provider="openai", model=model_name, original_tokens=residual_input_tokens + residual_tokens_saved, diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index 5a23acc04..6134230c2 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -87,6 +87,12 @@ class RequestOutcome: output_tokens: int tokens_saved: int attempted_input_tokens: int + # Optional correlation id for the human-readable PERF line. Most requests + # use ``request_id`` for both storage identity and log correlation. A + # long-lived WebSocket session is different: each emitted feed row needs a + # unique request id, while operators still need every line from the socket + # under one greppable session prefix. + perf_request_id: str | None = None # Optional so the 18 existing emit sites need no change: a handler that has # no provider count (or whose optimized_tokens is already provider-scaled) # leaves it 0 and billing falls back to optimized_tokens, exactly as before. @@ -565,7 +571,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None: tool_saved = tool_schema_saved_from_tags(outcome.tags or {}) total_saved = headline_tokens_saved(outcome.tokens_saved, outcome.tags or {}) logger.info( - f"[{outcome.request_id}] PERF " + f"[{outcome.perf_request_id or outcome.request_id}] PERF " f"model={outcome.model} msgs={outcome.num_messages} " f"tok_before={outcome.original_tokens} tok_after={outcome.optimized_tokens} " f"tok_saved={outcome.tokens_saved} " diff --git a/tests/test_openai_codex_ws_lifecycle.py b/tests/test_openai_codex_ws_lifecycle.py index b89a81a7e..dd60f5ed7 100644 --- a/tests/test_openai_codex_ws_lifecycle.py +++ b/tests/test_openai_codex_ws_lifecycle.py @@ -1030,8 +1030,15 @@ async def test_ws_session_log_prefix_uses_session_id(caplog: pytest.LogCaptureFi await handler.handle_openai_responses_ws(client_ws) assert handler.logger.entries - assert handler.logger.entries[0].request_id != "req-ws-1" + turn_request_id = handler.logger.entries[0].request_id + assert turn_request_id != "req-ws-1" + # Session lifecycle and PERF lines keep the session id so a session's log + # lines stay greppable together. The dashboard feed row retains its fresh + # per-turn id independently. + assert "[req-ws-1] WS /v1/responses accepted" in caplog.text + assert "[req-ws-1] WS /v1/responses completed" in caplog.text assert "[req-ws-1] PERF" in caplog.text + assert f"[{turn_request_id}] PERF" not in caplog.text @pytest.mark.asyncio diff --git a/tests/test_tokenizer_count_offload.py b/tests/test_tokenizer_count_offload.py index e4d77e4bc..e59b001a4 100644 --- a/tests/test_tokenizer_count_offload.py +++ b/tests/test_tokenizer_count_offload.py @@ -176,8 +176,12 @@ async def test_count_tokens_offloaded_fails_open_on_executor_quarantine() -> Non proxy = _make_proxy() # Record a concurrent compression as timed out so the real executor guard - # quarantines the next call — no mock of the helper itself. + # quarantines the next call — no mock of the helper itself. Since the + # quarantine became time-capped (#2412), standing debt alone no longer + # quarantines: the deadline armed by the fresh timeout must still be in + # the future, so arm it the way a real timeout would. proxy._compression_timed_out_in_flight = 1 + proxy._compression_quarantine_deadline = time.monotonic() + 60.0 tokenizer, tokens = await proxy._count_tokens_offloaded( "qwen2.5-coder", [{"role": "user", "content": "hello world"}] @@ -193,7 +197,10 @@ async def test_count_tokens_offloaded_returns_count_text_capable_tokenizer() -> that need per-fragment accounting.""" proxy = _make_proxy() # Quarantine forces the fail-open branch (an EstimatingTokenCounter). + # Post-#2412 the quarantine is time-capped, so the deadline must be armed + # alongside the standing debt. proxy._compression_timed_out_in_flight = 1 + proxy._compression_quarantine_deadline = time.monotonic() + 60.0 # The empty-messages count is intentionally discarded by that handler # (it sums text parts itself), so only the tokenizer matters here.