From 89f7b6c2dd0f3503312490b2052f534e76c01ac3 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Thu, 7 May 2026 14:50:03 -0700 Subject: [PATCH] fix: complete /v1/responses compression telemetry, multi-frame WS, frozen-count cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Builds on PR #406 (HTTP /v1/responses PyO3) and PR #410 (WS first-frame PyO3) which landed the binding for `compress_openai_responses_live_zone`. This change closes the remaining gaps so every (provider × endpoint × auth-mode × streaming) combination compresses AND surfaces in the dashboard. Telemetry: extend the PyO3 binding return tuple from `(bytes, modified)` to `(bytes, modified, tokens_saved, transforms_applied)` by adding `CompressionManifest::tokens_saved()` and `transforms_applied()` accessors on the existing manifest. The Python proxy populates request-log telemetry from the binding output instead of recounting tokens. Updates the existing 2-tuple call sites in HTTP and WS first-frame, plus the unpacks in tests. WebSocket multi-frame compression: subscription Codex users keep a long-lived WS open and send multiple `response.create` events per session. PR #410 only compressed the first frame; subsequent frames went raw. Added `_maybe_compress_response_create_frame` closure inside `_client_to_upstream` that runs the same Rust dispatcher on every client→upstream `response.create` text frame, passes other event types (response.cancel, session.update, etc.) through unchanged, and accumulates `tokens_saved` / `transforms_applied` / `ws_frames_compressed` counters across the session. Pre-existing dashboard gap: `streaming.py` and `anthropic.py` write `RequestLog` entries; the non-streaming OpenAI HTTP and WS handlers did not. Result: /transformations/feed was invisible for every Codex turn and every Cline / OpenClaude / Aider turn. Added the same wiring in `handle_openai_chat` (non-streaming), `handle_openai_responses` (non-streaming HTTP), and `handle_openai_responses_ws` (session-end). All three populate `auth_mode` + `endpoint` tags so the dashboard can break compression activity down by client class (PAYG / OAuth / Subscription) and surface (`chat_completions` / `responses_http` / `responses_ws`). The WS metric record is now unconditional — was previously gated on `tokens_saved > 0`, so first-frame no-changes never registered. compute_frozen_count over-freeze for prose-format clients: `compute_frozen_count` walked until it found an unstable `tool_result` / `role: "tool"` block. Cline / OpenClaude / Aider — clients that embed tool calls as XML inside plain text — never produce such a boundary, so the function returned `len(messages)` and the pipeline froze 100% of messages including the brand-new user turn. Live zone empty → `Transform content_router: 16414 → 16414 tokens (saved 0)`. Reported on Discord 2026-05-07 with Cline+DeepSeek. Fix: cap at `max(0, len(messages) - 1)`. Updates 3 existing test assertions whose expected values encoded the old over-freeze. Adds 6 new prose-format invariant tests. CodeQL "clear-text logging of sensitive information" fix: `tests/e2e_real_compression.py` previously stored API keys in local variables in the same scope as diagnostic prints, which CodeQL flagged via data-flow analysis. Refactored to read keys from `os.environ` inside the request helper — the credentials never enter the runner's main scope, so the taint flow never reaches the print. End-to-end verification with real keys (.env): /v1/messages (PAYG, non-stream) tok 14109 → 969 saved 13140 /v1/messages (PAYG, stream) tok 14109 → 969 saved 13140 /v1/chat/completions (PAYG, non-stream) tok 18460 → 1374 saved 17086 /v1/chat/completions (PAYG, stream) tok 18460 → 1374 saved 17086 (cache_hit=100%) /v1/responses HTTP (PAYG, non-stream) bytes 50138 → 597 saved 18391 /v1/responses WS (frame 1) bytes 46429 → 488 saved 16791 /v1/responses WS (frame 2 multi) bytes 46429 → 488 saved 16791 /v1/responses WS (response.cancel) passthrough untouched Tests: cargo workspace + pytest (4846 pass, 0 fail), make ci-precheck passed, two E2E scripts (multi-turn HTTP, WS fake-upstream) all green. --- .claude-plugin/marketplace.json | 4 +- .github/plugin/marketplace.json | 4 +- .../headroom-core/src/transforms/live_zone.rs | 141 +++++++ crates/headroom-py/src/lib.rs | 34 +- headroom/cache/compression_cache.py | 16 +- headroom/proxy/handlers/openai.py | 331 ++++++++++++++- .../.claude-plugin/plugin.json | 2 +- .../.github/plugin/plugin.json | 2 +- tests/e2e_real_compression.py | 390 ++++++++++++++++++ tests/e2e_ws_responses_compression.py | 279 +++++++++++++ tests/test_compression_cache.py | 27 +- tests/test_responses_pyo3_compression.py | 82 +++- tests/test_responses_ws_pyo3_compression.py | 2 +- tests/test_token_headroom_mode.py | 103 ++++- 14 files changed, 1362 insertions(+), 55 deletions(-) create mode 100644 tests/e2e_real_compression.py create mode 100644 tests/e2e_ws_responses_compression.py diff --git a/.claude-plugin/marketplace.json b/.claude-plugin/marketplace.json index 5bd29ddb2..1d6c46d37 100644 --- a/.claude-plugin/marketplace.json +++ b/.claude-plugin/marketplace.json @@ -5,14 +5,14 @@ }, "metadata": { "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.21.4" + "version": "0.21.5" }, "plugins": [ { "name": "headroom", "source": "./plugins/headroom-agent-hooks", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.21.4", + "version": "0.21.5", "author": { "name": "Headroom Contributors", "url": "https://github.com/chopratejas/headroom" diff --git a/.github/plugin/marketplace.json b/.github/plugin/marketplace.json index 5bd29ddb2..1d6c46d37 100644 --- a/.github/plugin/marketplace.json +++ b/.github/plugin/marketplace.json @@ -5,14 +5,14 @@ }, "metadata": { "description": "Headroom marketplace for Claude Code and GitHub Copilot CLI plugins.", - "version": "0.21.4" + "version": "0.21.5" }, "plugins": [ { "name": "headroom", "source": "./plugins/headroom-agent-hooks", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", - "version": "0.21.4", + "version": "0.21.5", "author": { "name": "Headroom Contributors", "url": "https://github.com/chopratejas/headroom" diff --git a/crates/headroom-core/src/transforms/live_zone.rs b/crates/headroom-core/src/transforms/live_zone.rs index 266776898..058a72b49 100644 --- a/crates/headroom-core/src/transforms/live_zone.rs +++ b/crates/headroom-core/src/transforms/live_zone.rs @@ -390,6 +390,42 @@ impl CompressionManifest { .iter() .any(|b| matches!(b.action, BlockAction::Compressed { .. })) } + + /// Aggregate `original_tokens − compressed_tokens` across every + /// `BlockAction::Compressed` outcome. Zero when no block was + /// rewritten. Saturating subtraction guards against the + /// theoretically-impossible case where a `Compressed` variant + /// reports compressed > original (the dispatcher's + /// `RejectedNotSmaller` gate should make this unreachable, but the + /// saturating arithmetic keeps callers panic-free). + pub fn tokens_saved(&self) -> usize { + self.block_outcomes + .iter() + .filter_map(|b| match &b.action { + BlockAction::Compressed { + original_tokens, + compressed_tokens, + .. + } => Some(original_tokens.saturating_sub(*compressed_tokens)), + _ => None, + }) + .sum() + } + + /// Distinct compressor strategies that actually produced rewritten + /// output, in first-seen order. Mirrors what the proxy logs as + /// `transforms_applied`. Empty when no block was rewritten. + pub fn transforms_applied(&self) -> Vec<&'static str> { + let mut seen: Vec<&'static str> = Vec::new(); + for b in &self.block_outcomes { + if let BlockAction::Compressed { strategy, .. } = &b.action { + if !seen.contains(strategy) { + seen.push(*strategy); + } + } + } + seen + } } /// Outcome of dispatching the live zone. @@ -1554,6 +1590,111 @@ mod tests { assert_eq!(manifest.messages_below_frozen_floor, 1); assert_eq!(manifest.latest_user_message_index, None); } + + // ─── Manifest accessor helpers (consumed by PyO3 binding) ───────── + + fn make_manifest(actions: Vec) -> CompressionManifest { + CompressionManifest { + messages_total: actions.len(), + messages_below_frozen_floor: 0, + latest_user_message_index: None, + block_outcomes: actions + .into_iter() + .enumerate() + .map(|(i, a)| BlockOutcome { + message_index: i, + block_index: None, + block_type: "test".to_string(), + action: a, + }) + .collect(), + } + } + + #[test] + fn tokens_saved_zero_for_empty_manifest() { + let m = CompressionManifest::empty(); + assert_eq!(m.tokens_saved(), 0); + assert!(m.transforms_applied().is_empty()); + } + + #[test] + fn tokens_saved_sums_compressed_outcomes_only() { + let m = make_manifest(vec![ + BlockAction::Compressed { + strategy: "smart_crusher", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 100, + compressed_tokens: 30, + }, + BlockAction::NoCompressionApplied { + content_type: "image".to_string(), + }, + BlockAction::Compressed { + strategy: "log_compressor", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 200, + compressed_tokens: 50, + }, + BlockAction::RejectedNotSmaller { + strategy: "smart_crusher", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 80, + compressed_tokens: 90, + }, + ]); + // 70 + 150 = 220; rejected variant must not contribute. + assert_eq!(m.tokens_saved(), 220); + } + + #[test] + fn transforms_applied_dedup_first_seen_order() { + let m = make_manifest(vec![ + BlockAction::Compressed { + strategy: "log_compressor", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 50, + compressed_tokens: 10, + }, + BlockAction::Compressed { + strategy: "smart_crusher", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 50, + compressed_tokens: 10, + }, + BlockAction::Compressed { + strategy: "log_compressor", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 50, + compressed_tokens: 10, + }, + ]); + assert_eq!( + m.transforms_applied(), + vec!["log_compressor", "smart_crusher"] + ); + } + + #[test] + fn tokens_saved_saturates_when_compressed_exceeds_original() { + // Defensive — the dispatcher's RejectedNotSmaller gate should + // make this unreachable, but the helper must not panic if a + // future caller hand-constructs such a manifest. + let m = make_manifest(vec![BlockAction::Compressed { + strategy: "smart_crusher", + original_bytes: 0, + compressed_bytes: 0, + original_tokens: 10, + compressed_tokens: 50, + }]); + assert_eq!(m.tokens_saved(), 0); + } } // ─── OpenAI Chat Completions live-zone dispatcher (Phase C PR-C2) ──────── diff --git a/crates/headroom-py/src/lib.rs b/crates/headroom-py/src/lib.rs index ccbdc09bf..9dac7b9dc 100644 --- a/crates/headroom-py/src/lib.rs +++ b/crates/headroom-py/src/lib.rs @@ -1498,7 +1498,7 @@ fn compress_openai_responses_live_zone( body: &[u8], auth_mode: &str, model: &str, -) -> (Py, bool) { +) -> (Py, bool, u64, Vec) { let mode = match auth_mode.to_ascii_lowercase().as_str() { "payg" => RustLiveZoneAuthMode::Payg, "oauth" => RustLiveZoneAuthMode::OAuth, @@ -1512,17 +1512,41 @@ fn compress_openai_responses_live_zone( }; match rust_compress_openai_responses_live_zone(body, mode, model_str) { - Ok(LiveZoneOutcome::NoChange { .. }) => (PyBytes::new_bound(py, body).unbind(), false), - Ok(LiveZoneOutcome::Modified { new_body, .. }) => { + Ok(LiveZoneOutcome::NoChange { manifest }) => { + let saved = manifest.tokens_saved() as u64; + let transforms: Vec = manifest + .transforms_applied() + .into_iter() + .map(String::from) + .collect(); + ( + PyBytes::new_bound(py, body).unbind(), + false, + saved, + transforms, + ) + } + Ok(LiveZoneOutcome::Modified { new_body, manifest }) => { // `RawValue::get` returns the underlying serialized JSON // as `&str`; bytes are valid UTF-8 by construction. let bytes = new_body.get().as_bytes(); - (PyBytes::new_bound(py, bytes).unbind(), true) + let saved = manifest.tokens_saved() as u64; + let transforms: Vec = manifest + .transforms_applied() + .into_iter() + .map(String::from) + .collect(); + ( + PyBytes::new_bound(py, bytes).unbind(), + true, + saved, + transforms, + ) } Err(_) => { // BodyNotJson / NoMessagesArray are non-fatal: nothing to // compress, fall through to passthrough byte-for-byte. - (PyBytes::new_bound(py, body).unbind(), false) + (PyBytes::new_bound(py, body).unbind(), false, 0, Vec::new()) } } } diff --git a/headroom/cache/compression_cache.py b/headroom/cache/compression_cache.py index c8cbfed3f..a5ec4bf53 100644 --- a/headroom/cache/compression_cache.py +++ b/headroom/cache/compression_cache.py @@ -228,6 +228,18 @@ class CompressionCache: an assistant message with tool_use blocks, or a tool_result whose content hash is already in the cache. The first unstable tool_result (cache miss) stops the count. + + The trailing message is *always* excluded from the frozen prefix + (cap of ``len(messages) - 1``). The trailing message represents + the just-arrived turn — by definition it has not yet been sent + upstream and therefore cannot be in any provider prefix cache. + Without this cap, prose-format clients (Cline, OpenClaude, Aider, + any client that does not use OpenAI-native or Anthropic-native + tool messages) would have every message marked stable, making + the live zone empty and producing zero compression. See issue + observed 2026-05-07 with Cline+DeepSeek (`Pipeline: freezing + first N/N messages` followed by ``Transform content_router: + X -> X tokens (saved 0)`` on every request). """ with self._lock: count = 0 @@ -244,7 +256,9 @@ class CompressionCache: # Regular user/assistant/system messages and assistant+tool_use # are always stable — fall through. count += 1 - return count + # Reserve the trailing message as the live zone. `max(0, ...)` + # handles the empty-list edge case cleanly. + return min(count, max(0, len(messages) - 1)) def apply_cached(self, messages: list[dict]) -> list[dict]: """Return a new list with cached compressions swapped into tool results. diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index a598ca85b..f0036feaf 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -1091,6 +1091,16 @@ class OpenAIHandlerMixin: get_codex_rate_limit_state().update_from_headers(dict(response.headers)) + # Tag the metric/log with auth_mode + endpoint so the + # dashboard can break down by client class (PAYG vs + # subscription vs OAuth) without re-classifying. + _auth_mode_chat = getattr(request.state, "auth_mode", None) + _chat_log_tags = { + **(tags or {}), + "auth_mode": _auth_mode_chat.value if _auth_mode_chat else "payg", + "endpoint": "chat_completions", + } + await self.metrics.record_request( provider="openai", model=model, @@ -1105,6 +1115,41 @@ class OpenAIHandlerMixin: uncached_input_tokens=uncached_input_tokens, ) + # Per-request log entry for /transformations/feed + + # /stats `recent_requests`. Without this the dashboard + # only shows aggregates for non-streaming OpenAI traffic. + # Mirror of the streaming.py + anthropic.py wiring. + if getattr(self, "logger", None) is not None: + from headroom.proxy.helpers import compute_turn_id + from headroom.proxy.models import RequestLog + + self.logger.log( + RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="openai", + model=model, + input_tokens_original=original_tokens, + input_tokens_optimized=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / original_tokens * 100) + if original_tokens > 0 + else 0, + optimization_latency_ms=optimization_latency, + total_latency_ms=total_latency, + tags=_chat_log_tags, + cache_hit=False, + transforms_applied=transforms_applied, + request_messages=body.get("messages") + if getattr(self.config, "log_full_messages", False) + else None, + turn_id=compute_turn_id( + model, body.get("system"), body.get("messages") + ), + ) + ) + if tokens_saved > 0: logger.info( f"[{request_id}] {model}: {original_tokens:,} → {optimized_tokens:,} " @@ -1493,18 +1538,34 @@ class OpenAIHandlerMixin: ) _input_bytes = json.dumps(body).encode("utf-8") - _new_bytes, _modified = _rust_compress_responses( + ( + _new_bytes, + _modified, + _rust_tokens_saved, + _rust_transforms, + ) = _rust_compress_responses( _input_bytes, auth_mode.value, model, ) if _modified: body = json.loads(_new_bytes) - transforms_applied = list(transforms_applied) + ["openai_responses_live_zone"] + tokens_saved = int(_rust_tokens_saved) + optimized_tokens = max(0, original_tokens - tokens_saved) + transforms_applied = [ + "openai_responses_live_zone", + *_rust_transforms, + *list(transforms_applied), + ] logger.info( - f"[{request_id}] /v1/responses compressed " - f"{len(_input_bytes):,}→{len(_new_bytes):,} bytes " - f"(auth_mode={auth_mode.value})" + "[%s] /v1/responses compressed %d→%d bytes " + "(%d tokens saved, auth_mode=%s, transforms=%s)", + request_id, + len(_input_bytes), + len(_new_bytes), + tokens_saved, + auth_mode.value, + transforms_applied, ) except Exception as _e: logger.warning( @@ -1625,6 +1686,12 @@ class OpenAIHandlerMixin: if self.cost_tracker: self.cost_tracker.record_tokens(model, tokens_saved, total_input_tokens) + _resp_log_tags = { + **(tags or {}), + "auth_mode": auth_mode.value if auth_mode else "payg", + "endpoint": "responses_http", + } + await self.metrics.record_request( provider="openai", model=model, @@ -1635,6 +1702,39 @@ class OpenAIHandlerMixin: overhead_ms=optimization_latency, ) + # Per-request log entry for /transformations/feed + + # /stats `recent_requests`. Mirror of streaming.py / + # anthropic.py wiring; without this the dashboard's + # per-request feed misses every Codex HTTP turn. + if getattr(self, "logger", None) is not None: + from headroom.proxy.helpers import compute_turn_id + from headroom.proxy.models import RequestLog + + self.logger.log( + RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="openai", + model=model, + input_tokens_original=original_tokens, + input_tokens_optimized=optimized_tokens, + output_tokens=output_tokens, + tokens_saved=tokens_saved, + savings_percent=(tokens_saved / original_tokens * 100) + if original_tokens > 0 + else 0, + optimization_latency_ms=optimization_latency, + total_latency_ms=total_latency, + tags=_resp_log_tags, + cache_hit=False, + transforms_applied=transforms_applied, + request_messages=messages + if getattr(self.config, "log_full_messages", False) + else None, + turn_id=compute_turn_id(model, body.get("instructions"), messages), + ) + ) + logger.info(f"[{request_id}] /v1/responses {model}: {total_input_tokens:,} tokens") # Capture Codex rate-limit window data from response headers @@ -1707,6 +1807,13 @@ class OpenAIHandlerMixin: # Forward client headers to upstream, adding required OpenAI-Beta header ws_headers = dict(websocket.headers) + # Extract per-request tags from headers up front so the + # session-end RequestLog can attach them. `_extract_tags` is + # the same helper the HTTP handlers use; on a WebSocket the + # tags come from `x-headroom-tag-*` headers in the upgrade + # handshake. Returns `{}` when no tags are present. + _extract_ws_tags = getattr(self, "_extract_tags", None) + ws_tags = _extract_ws_tags(ws_headers) if callable(_extract_ws_tags) else {} # Extract subprotocol from client — this is an application-level negotiation # that MUST be forwarded end-to-end (unlike sec-websocket-key which is per-connection). @@ -1921,10 +2028,14 @@ class OpenAIHandlerMixin: # via uvicorn). Hot-fix follow-up to PR #406: the first frame # is now compressed via the inline PyO3 binding right before # upstream send (see the compression block further below). - # Subsequent client→upstream frames in the relay loop remain - # unmodified — multi-frame compression is a separate follow-up. + # Subsequent client→upstream frames are now ALSO compressed + # via `_maybe_compress_response_create_frame` in + # `_client_to_upstream` so long-lived subscription Codex + # sessions get savings on every turn, not just the first. body: dict[str, Any] = {} tokens_saved = 0 + transforms_applied: list[str] = [] + ws_frames_compressed = 0 try: body = json.loads(first_msg_raw) except json.JSONDecodeError: @@ -2111,7 +2222,12 @@ class OpenAIHandlerMixin: _model = (_inner.get("model") if isinstance(_inner, dict) else None) or "" _inner_bytes = json.dumps(_inner).encode("utf-8") - _new_bytes, _modified = _rust_compress_responses( + ( + _new_bytes, + _modified, + _ws_rust_saved, + _ws_rust_transforms, + ) = _rust_compress_responses( _inner_bytes, _ws_auth_mode.value, _model, @@ -2127,10 +2243,23 @@ class OpenAIHandlerMixin: else: _send_body = _new_inner first_msg_raw = json.dumps(_send_body) + tokens_saved += int(_ws_rust_saved) + for _t in ( + "openai_responses_ws_live_zone", + *list(_ws_rust_transforms), + ): + if _t not in transforms_applied: + transforms_applied.append(_t) logger.info( - f"[{request_id}] WS /v1/responses compressed " - f"{len(_inner_bytes):,}→{len(_new_bytes):,} bytes " - f"(auth_mode={_ws_auth_mode.value})" + "[%s] WS /v1/responses compressed " + "%d→%d bytes (%d tokens saved, " + "auth_mode=%s, transforms=%s)", + request_id, + len(_inner_bytes), + len(_new_bytes), + int(_ws_rust_saved), + _ws_auth_mode.value, + transforms_applied, ) except Exception as _ce: logger.warning( @@ -2192,11 +2321,107 @@ class OpenAIHandlerMixin: upstream_relay_error: BaseException | None = None client_relay_error: BaseException | None = None + async def _maybe_compress_response_create_frame( + raw_msg: str, + ) -> str: + """Compress a single client→upstream frame + when its `type` is `response.create`. Other + event types (response.cancel, session.update, + etc.) pass through unchanged. Errors are + warned and the original frame is returned — + fail loud in logs, fail safe on the wire. + Updates outer-scope ``tokens_saved``, + ``transforms_applied``, and + ``ws_frames_compressed`` so the session-end + log reports cumulative savings across all + frames in the WS session. + """ + nonlocal tokens_saved, transforms_applied + nonlocal ws_frames_compressed + if not self.config.optimize: + return raw_msg + try: + parsed_frame = json.loads(raw_msg) + except json.JSONDecodeError: + return raw_msg + if ( + not isinstance(parsed_frame, dict) + or parsed_frame.get("type") != "response.create" + ): + return raw_msg + wrapped_frame = isinstance(parsed_frame.get("response"), dict) + inner_payload = ( + parsed_frame["response"] if wrapped_frame else parsed_frame + ) + if not isinstance(inner_payload, dict): + return raw_msg + try: + from headroom._core import ( + compress_openai_responses_live_zone as _ws_frame_compress, + ) + + inner_bytes = json.dumps(inner_payload).encode("utf-8") + model_for_frame = inner_payload.get("model") or "" + _frame_auth_mode = classify_auth_mode(ws_headers) + ( + new_inner_bytes, + modified, + rust_saved, + rust_transforms, + ) = _ws_frame_compress( + inner_bytes, + _frame_auth_mode.value, + model_for_frame, + ) + except Exception as _frame_err: + logger.warning( + "[%s] WS /v1/responses frame compression " + "failed; forwarding original: %s: %s", + request_id, + type(_frame_err).__name__, + _frame_err, + ) + return raw_msg + if not modified: + return raw_msg + try: + new_inner = json.loads(new_inner_bytes) + except json.JSONDecodeError: + return raw_msg + if not isinstance(new_inner, dict): + return raw_msg + if wrapped_frame: + parsed_frame["response"] = new_inner + rewritten = json.dumps(parsed_frame) + else: + rewritten = json.dumps(new_inner) + tokens_saved += int(rust_saved) + for t in ( + "openai_responses_ws_live_zone", + *list(rust_transforms), + ): + if t not in transforms_applied: + transforms_applied.append(t) + ws_frames_compressed += 1 + logger.info( + "[%s] WS /v1/responses frame compressed " + "%d→%d bytes (%d tokens saved, " + "auth_mode=%s, frame=%d)", + request_id, + len(inner_bytes), + len(new_inner_bytes), + int(rust_saved), + _frame_auth_mode.value, + ws_frames_compressed, + ) + return rewritten + async def _client_to_upstream() -> None: nonlocal client_relay_error try: while True: msg = await websocket.receive_text() + msg = await _maybe_compress_response_create_frame(msg) await upstream.send(msg) except asyncio.CancelledError: # Explicit cancel from the outer @@ -2577,16 +2802,80 @@ class OpenAIHandlerMixin: websocket, body, first_msg_raw, upstream_headers, request_id ) - # Record metrics - if tokens_saved > 0: - model_name = body.get("model", "unknown") if isinstance(body, dict) else "unknown" - await self.metrics.record_request( - provider="openai", - model=model_name, - input_tokens=0, - output_tokens=0, - tokens_saved=tokens_saved, - latency_ms=0, + # ── WS session-end metric + RequestLog ────────────────── + # + # Unconditional (was previously gated on `tokens_saved>0`, + # which made first-frame no-changes invisible). We record + # one entry per WS session that aggregates `tokens_saved` + # across every `response.create` frame compressed by the + # first-frame block + `_maybe_compress_response_create_frame`. + # The RequestLog entry mirrors the streaming.py / + # anthropic.py shape so /transformations/feed surfaces + # Codex WS turns. + ws_session_duration_ms = (time.perf_counter() - session_started_at) * 1000.0 + ws_inner_for_telemetry: dict[str, Any] = ( + body.get("response", body) if isinstance(body, dict) else {} + ) + if not isinstance(ws_inner_for_telemetry, dict): + ws_inner_for_telemetry = {} + model_name = ( + ws_inner_for_telemetry.get("model") + or (body.get("model") if isinstance(body, dict) else None) + or "unknown" + ) + _final_auth_mode = classify_auth_mode(ws_headers) + ws_session_tags = { + **(ws_tags or {}), + "auth_mode": _final_auth_mode.value, + "endpoint": "responses_ws", + "ws_frames_compressed": str(ws_frames_compressed), + } + await self.metrics.record_request( + provider="openai", + model=model_name, + input_tokens=0, + output_tokens=0, + tokens_saved=tokens_saved, + latency_ms=ws_session_duration_ms, + ) + if getattr(self, "logger", None) is not None: + from headroom.proxy.helpers import compute_turn_id + from headroom.proxy.models import RequestLog + + ws_messages_for_log: list[dict[str, Any]] = [] + ws_input_for_log = ws_inner_for_telemetry.get("input") + ws_instructions_for_log = ws_inner_for_telemetry.get("instructions") + if isinstance(ws_instructions_for_log, str) and ws_instructions_for_log: + ws_messages_for_log.append( + {"role": "system", "content": ws_instructions_for_log} + ) + if isinstance(ws_input_for_log, str) and ws_input_for_log: + ws_messages_for_log.append({"role": "user", "content": ws_input_for_log}) + self.logger.log( + RequestLog( + request_id=request_id, + timestamp=datetime.now().isoformat(), + provider="openai", + model=model_name, + input_tokens_original=0, + input_tokens_optimized=0, + output_tokens=0, + tokens_saved=tokens_saved, + savings_percent=0.0, + optimization_latency_ms=0.0, + total_latency_ms=ws_session_duration_ms, + tags=ws_session_tags, + cache_hit=False, + transforms_applied=transforms_applied, + request_messages=ws_messages_for_log + if getattr(self.config, "log_full_messages", False) + else None, + turn_id=compute_turn_id( + model_name, + ws_instructions_for_log, + ws_messages_for_log, + ), + ) ) except Exception as e: diff --git a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json index 1ca7cfe16..eda3fff15 100644 --- a/plugins/headroom-agent-hooks/.claude-plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "headroom", - "version": "0.21.4", + "version": "0.21.5", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", "author": { "name": "Headroom Contributors", diff --git a/plugins/headroom-agent-hooks/.github/plugin/plugin.json b/plugins/headroom-agent-hooks/.github/plugin/plugin.json index a33cddbd2..7aa1a7019 100644 --- a/plugins/headroom-agent-hooks/.github/plugin/plugin.json +++ b/plugins/headroom-agent-hooks/.github/plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "headroom", - "version": "0.21.4", + "version": "0.21.5", "description": "Headroom startup hooks for Claude Code and GitHub Copilot CLI.", "author": { "name": "Headroom Contributors", diff --git a/tests/e2e_real_compression.py b/tests/e2e_real_compression.py new file mode 100644 index 000000000..6816bf2ca --- /dev/null +++ b/tests/e2e_real_compression.py @@ -0,0 +1,390 @@ +"""End-to-end compression verification with realistic multi-turn payloads. + +Headroom only compresses content the model has already seen — assistant +turns, tool results, and Responses-API output items. Fresh user prompts +are *intentionally* skipped (the model needs them verbatim, and they're +in the live-zone tail anyway). A conversation that contains nothing but +a single user prompt will produce 0 tokens saved by design — that is +not a bug; it's the live-zone-only invariant. + +This script exercises every (provider × endpoint × streaming) combination +with a payload large enough to trigger compression. Pass criteria: + + * tokens_saved > 0 for at least one chat-completions case + * tokens_saved > 0 for at least one /v1/messages case + * tokens_saved > 0 for the /v1/responses case + * tokens_saved > 0 for streaming variants + * No proxy errors, no compression-failed warnings on happy paths + +Reads keys from .env. Run via: + + .venv/bin/python tests/e2e_real_compression.py + +# Note on auth-header construction +# The API keys are read from `os.environ` *inside* `_post` and never +# stored as local variables in the test runner's main scope. This +# breaks the CodeQL taint flow that would otherwise flag every +# diagnostic `print()` in the loop as +# `py/clear-text-logging-sensitive-data` because credentials live in +# the same scope. +""" + +from __future__ import annotations + +import json +import os +import re +import socket +import subprocess +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +def load_env_into_environ() -> None: + """Read REPO_ROOT/.env and merge into os.environ. Keys are never + returned to the caller — see module docstring.""" + p = REPO_ROOT / ".env" + if not p.exists(): + return + for line in p.read_text().splitlines(): + line = line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + k, v = line.split("=", 1) + k = k.strip() + v = v.strip().strip('"').strip("'") + os.environ.setdefault(k, v) + + +def have_required_keys() -> tuple[bool, str]: + """Sentinel check without exposing the keys themselves to local scope.""" + missing = [n for n in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY") if not os.environ.get(n)] + if missing: + return False, ", ".join(missing) + return True, "" + + +def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def wait_ready(port: int, timeout_s: float = 60.0) -> None: + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r: + if r.status == 200: + return + except Exception: + time.sleep(0.5) + raise TimeoutError("proxy not ready") + + +def _post(url: str, body: dict, *, provider: str, stream: bool = False) -> tuple[int, Any]: + """Make a POST request, building auth headers from os.environ at + call time. The credential never appears in the caller's local + scope, which keeps CodeQL's taint analysis happy.""" + if provider == "openai": + headers = { + "Authorization": "Bearer " + (os.environ.get("OPENAI_API_KEY") or ""), + "Content-Type": "application/json", + } + elif provider == "anthropic": + headers = { + "x-api-key": os.environ.get("ANTHROPIC_API_KEY") or "", + "anthropic-version": "2023-06-01", + "Content-Type": "application/json", + } + else: + raise ValueError(f"unknown provider: {provider!r}") + + req = urllib.request.Request( + url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST" + ) + try: + with urllib.request.urlopen(req, timeout=120) as r: + raw = r.read() + if stream: + return r.status, raw.decode("utf-8", errors="replace") + try: + return r.status, json.loads(raw) + except json.JSONDecodeError: + return r.status, raw.decode("utf-8", errors="replace") + except urllib.error.HTTPError as e: + try: + return e.code, json.loads(e.read()) + except Exception: + return e.code, str(e) + + +# ── Payload builders ────────────────────────────────────────────────────── + + +def long_build_log() -> str: + """~24 KB of structured BuildOutput-style content. The Rust + LogCompressor recognizes this and compresses aggressively.""" + return "".join( + f"[2024-01-01 00:00:{i % 60:02d}] INFO compile.rs:42 building module foo_{i} " + f"(crate=workspace-{i // 10}, deps=[serde={i}, tokio={i}, regex={i % 7}])\n" + for i in range(400) + ) + + +def anthropic_messages_payload(streaming: bool = False) -> dict: + return { + "model": "claude-haiku-4-5-20251001", + "max_tokens": 30, + "stream": streaming, + "tools": [ + { + "name": "shell", + "description": "Run a shell command", + "input_schema": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + } + ], + "messages": [ + {"role": "user", "content": "Run cargo build and tell me if it succeeded."}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Running it now."}, + { + "type": "tool_use", + "id": "toolu_e2e_1", + "name": "shell", + "input": {"command": "cargo build --release"}, + }, + ], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_e2e_1", + "content": long_build_log(), + } + ], + }, + {"role": "user", "content": "One word: pass or fail?"}, + ], + } + + +def openai_chat_payload(streaming: bool = False) -> dict: + return { + "model": "gpt-4o-mini", + "max_tokens": 30, + "stream": streaming, + "messages": [ + {"role": "user", "content": "Run cargo build and report if it succeeded."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_e2e_1", + "type": "function", + "function": { + "name": "shell", + "arguments": '{"command": "cargo build --release"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_e2e_1", + "content": long_build_log(), + }, + {"role": "user", "content": "One word: pass or fail?"}, + ], + } + + +def openai_responses_payload(streaming: bool = False) -> dict: + return { + "model": "gpt-4o-mini", + "max_output_tokens": 30, + "stream": streaming, + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "Run cargo build and report."}], + }, + { + "type": "function_call", + "call_id": "call_resp_e2e_1", + "name": "shell", + "arguments": '{"command": "cargo build --release"}', + }, + { + "type": "function_call_output", + "call_id": "call_resp_e2e_1", + "output": long_build_log(), + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "One word: pass or fail?"}], + }, + ], + "instructions": "You read shell output and reply tersely.", + } + + +# ── Test runner ─────────────────────────────────────────────────────────── + + +def main() -> int: + load_env_into_environ() + ok, missing = have_required_keys() + if not ok: + print(f"FAIL: missing keys: {missing}", file=sys.stderr) + return 1 + + port = free_port() + print(f"[e2e] starting proxy on :{port}") + log_fp = open("/tmp/e2e_real_proxy.log", "w") + proc = subprocess.Popen( + [ + str(REPO_ROOT / ".venv/bin/headroom"), + "proxy", + "--port", + str(port), + "--no-telemetry", + ], + env={**os.environ, "HEADROOM_REQUIRE_RUST_CORE": "true"}, + stdout=log_fp, + stderr=subprocess.STDOUT, + cwd=str(REPO_ROOT), + ) + + failures: list[str] = [] + try: + wait_ready(port) + print("[e2e] proxy ready") + + # Cases carry only structural info: name, path, provider tag, + # body, stream. Auth headers are built inside `_post` from + # os.environ — see module docstring. + cases: list[tuple[str, str, str, dict, bool]] = [ + ( + "anthropic_messages_nonstream", + "/v1/messages", + "anthropic", + anthropic_messages_payload(streaming=False), + False, + ), + ( + "anthropic_messages_stream", + "/v1/messages", + "anthropic", + anthropic_messages_payload(streaming=True), + True, + ), + ( + "openai_chat_nonstream", + "/v1/chat/completions", + "openai", + openai_chat_payload(streaming=False), + False, + ), + ( + "openai_chat_stream", + "/v1/chat/completions", + "openai", + openai_chat_payload(streaming=True), + True, + ), + ( + "openai_responses_nonstream", + "/v1/responses", + "openai", + openai_responses_payload(streaming=False), + False, + ), + ] + + for name, path, provider, body, stream in cases: + url = f"http://127.0.0.1:{port}{path}" + print(f"[e2e] {name}: POST {path}") + status, _ = _post(url, body, provider=provider, stream=stream) + if status != 200: + failures.append(f"{name}: HTTP {status}") + continue + print(" ok status=200") + + # ── Scrape proxy log for compression evidence ──────────── + time.sleep(1.5) + canonical_log = Path.home() / ".headroom" / "logs" / "proxy.log" + if canonical_log.exists(): + log_lines = canonical_log.read_text(errors="replace").splitlines()[-5000:] + else: + log_lines = Path("/tmp/e2e_real_proxy.log").read_text(errors="replace").splitlines() + + compressed_evidence = [ + line + for line in log_lines + if "compressed" in line and ("tokens" in line.lower() or "bytes" in line.lower()) + ][-30:] + if compressed_evidence: + print("\n[e2e] compression evidence (last 10 lines):") + for line in compressed_evidence[-10:]: + idx = line.find("] ") + print(" ", line[idx + 2 :] if idx > 0 else line) + else: + print("\n[e2e] no compression evidence in canonical log") + + saved_pattern = re.compile(r"saved (\d[\d,]*) tokens?", re.IGNORECASE) + total_saved = 0 + for line in log_lines[-2000:]: + m = saved_pattern.search(line) + if m: + num = int(m.group(1).replace(",", "")) + if num > 0: + total_saved += num + + print( + f"\n[e2e] aggregate tokens saved across cases (~last 2000 log lines): {total_saved:,}" + ) + if total_saved == 0: + failures.append(f"no compression evidence — check {canonical_log}") + + joined = "\n".join(log_lines) + if "compression failed" in joined: + failures.append("proxy log contains 'compression failed' — see canonical log") + + finally: + print("\n[e2e] terminating proxy") + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + log_fp.close() + + if failures: + print("\n=== E2E FAILURES ===") + for f in failures: + print(" -", f) + return 1 + print("\n=== E2E ALL GREEN ===") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/e2e_ws_responses_compression.py b/tests/e2e_ws_responses_compression.py new file mode 100644 index 000000000..a95d625eb --- /dev/null +++ b/tests/e2e_ws_responses_compression.py @@ -0,0 +1,279 @@ +"""End-to-end verification that /v1/responses WebSocket compression fires. + +We can't reach OpenAI's WS endpoint without the `responses_websockets` +beta enabled on the test API key, so this test does the next-best +thing: it spins up a *fake upstream* WebSocket server, points the +proxy at it via OPENAI_API_URL, and connects a client to the proxy. + +Verifies: + 1. First-frame compression: the client sends a `response.create` + event with a 24 KB output_item; the fake upstream receives the + COMPRESSED frame (much smaller than what was sent). + 2. Multi-frame compression: a second `response.create` on the same + WS session is also compressed (the new behavior — was previously + first-frame-only). + 3. Other event types (e.g. `response.cancel`) pass through + unchanged. + 4. Proxy log surfaces both compression events with token-saved + numbers. + +Run via: + + .venv/bin/python tests/e2e_ws_responses_compression.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import socket +import subprocess +import sys +import time +import urllib.request +from pathlib import Path + +import websockets + +REPO_ROOT = Path(__file__).resolve().parent.parent + + +def free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def wait_ready(port: int, timeout_s: float = 60.0) -> None: + deadline = time.time() + timeout_s + while time.time() < deadline: + try: + with urllib.request.urlopen(f"http://127.0.0.1:{port}/livez", timeout=2) as r: + if r.status == 200: + return + except Exception: + time.sleep(0.5) + raise TimeoutError("proxy not ready") + + +def long_build_log() -> str: + return "".join( + f"[2024-01-01 00:00:{i % 60:02d}] INFO compile.rs:42 building module foo_{i} " + f"(crate=workspace-{i // 10}, deps=[serde={i}, tokio={i}])\n" + for i in range(400) + ) + + +def make_response_create_payload(turn_no: int) -> dict: + """A wire-shape `response.create` envelope, with a long + function_call_output that the dispatcher will compress.""" + return { + "type": "response.create", + "response": { + "model": "gpt-4o-mini", + "input": [ + { + "type": "message", + "role": "user", + "content": [ + { + "type": "input_text", + "text": f"Turn {turn_no} — please summarize the build output.", + } + ], + }, + { + "type": "function_call", + "call_id": f"call_e2e_ws_{turn_no}", + "name": "shell", + "arguments": '{"command": "cargo build --release"}', + }, + { + "type": "function_call_output", + "call_id": f"call_e2e_ws_{turn_no}", + "output": long_build_log(), + }, + ], + "instructions": "You read shell output and reply tersely.", + "max_output_tokens": 30, + }, + } + + +# ── Fake upstream WS server ───────────────────────────────────────────── +# +# Captures every text frame the proxy forwards so we can assert what +# arrived upstream actually got compressed. + + +class FakeUpstream: + def __init__(self) -> None: + self.received_frames: list[str] = [] + self.server: websockets.server.WebSocketServer | None = None + self.port: int = 0 + + async def _handler(self, ws): + try: + async for msg in ws: + if isinstance(msg, str): + self.received_frames.append(msg) + # Echo a minimal completion event so the proxy doesn't + # think upstream is hung. + await ws.send( + json.dumps( + { + "type": "response.completed", + "response": {"id": "fake_resp", "output": []}, + } + ) + ) + except websockets.exceptions.ConnectionClosed: + pass + + async def start(self) -> int: + self.port = free_port() + self.server = await websockets.serve(self._handler, "127.0.0.1", self.port) + return self.port + + async def stop(self) -> None: + if self.server: + self.server.close() + await self.server.wait_closed() + + +async def main_async() -> int: + fake = FakeUpstream() + upstream_port = await fake.start() + upstream_url = f"http://127.0.0.1:{upstream_port}" + + proxy_port = free_port() + print(f"[ws-e2e] fake upstream at ws://127.0.0.1:{upstream_port}") + print(f"[ws-e2e] starting proxy on :{proxy_port}") + + log_fp = open("/tmp/e2e_ws_proxy.log", "w") + proc = subprocess.Popen( + [ + str(REPO_ROOT / ".venv/bin/headroom"), + "proxy", + "--port", + str(proxy_port), + "--no-telemetry", + # Point /v1/responses upstream at our fake server instead + # of api.openai.com. + "--openai-api-url", + upstream_url, + ], + env={ + **os.environ, + # Need *some* OpenAI key value so the proxy doesn't refuse; + # the fake upstream ignores it. + "OPENAI_API_KEY": os.environ.get("OPENAI_API_KEY", "sk-fake-for-test"), + "ANTHROPIC_API_KEY": os.environ.get("ANTHROPIC_API_KEY", "sk-ant-fake"), + "HEADROOM_REQUIRE_RUST_CORE": "true", + }, + stdout=log_fp, + stderr=subprocess.STDOUT, + cwd=str(REPO_ROOT), + ) + + failures: list[str] = [] + try: + wait_ready(proxy_port) + print("[ws-e2e] proxy ready") + + # ── Connect WS client to the proxy ─────────────────────── + proxy_ws_url = f"ws://127.0.0.1:{proxy_port}/v1/responses" + async with websockets.connect( + proxy_ws_url, + additional_headers={ + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY', 'sk-fake')}", + "OpenAI-Beta": "responses_websockets=2026-02-06", + }, + ) as ws: + # Frame 1: response.create with large content + payload_1 = make_response_create_payload(1) + payload_1_bytes = len(json.dumps(payload_1).encode("utf-8")) + print(f"[ws-e2e] sending frame 1 ({payload_1_bytes:,} bytes)") + await ws.send(json.dumps(payload_1)) + + # Wait for the fake upstream to receive (or timeout) + await asyncio.sleep(2.0) + + # Frame 2: ANOTHER response.create on the same session + payload_2 = make_response_create_payload(2) + payload_2_bytes = len(json.dumps(payload_2).encode("utf-8")) + print(f"[ws-e2e] sending frame 2 ({payload_2_bytes:,} bytes)") + await ws.send(json.dumps(payload_2)) + await asyncio.sleep(2.0) + + # Frame 3: a non-response.create event — should pass through + cancel = {"type": "response.cancel"} + print("[ws-e2e] sending frame 3 (response.cancel — passthrough)") + await ws.send(json.dumps(cancel)) + await asyncio.sleep(1.0) + + # ── Inspect what arrived at the fake upstream ──────────── + print(f"\n[ws-e2e] fake upstream received {len(fake.received_frames)} frames") + for i, frame in enumerate(fake.received_frames): + print(f" frame {i + 1}: {len(frame.encode()):,} bytes") + + # First two frames should be MUCH smaller than what we sent. + # The third (response.cancel) should be a small fixed size. + if len(fake.received_frames) < 2: + failures.append(f"expected ≥2 frames at upstream, got {len(fake.received_frames)}") + else: + f1 = len(fake.received_frames[0].encode()) + f2 = len(fake.received_frames[1].encode()) + if f1 >= payload_1_bytes // 2: + failures.append( + f"frame 1 not compressed: arrived {f1:,} bytes (sent {payload_1_bytes:,})" + ) + if f2 >= payload_2_bytes // 2: + failures.append( + f"frame 2 not compressed (multi-frame regression): " + f"arrived {f2:,} bytes (sent {payload_2_bytes:,})" + ) + + # ── Scrape proxy log for compression evidence ──────────── + await asyncio.sleep(1.0) + canonical = Path.home() / ".headroom" / "logs" / "proxy.log" + log_lines = canonical.read_text(errors="replace").splitlines()[-1000:] + ws_compressed = [ + line for line in log_lines if "WS /v1/responses" in line and "compressed" in line + ] + print("\n[ws-e2e] WS compression log lines (last few):") + for line in ws_compressed[-6:]: + idx = line.find("] ") + print(" ", line[idx + 2 :] if idx > 0 else line) + if len(ws_compressed) < 2: + failures.append( + f"expected ≥2 WS compression log entries (first frame + multi-frame), " + f"saw {len(ws_compressed)}" + ) + + finally: + print("\n[ws-e2e] terminating proxy") + proc.terminate() + try: + proc.wait(timeout=10) + except subprocess.TimeoutExpired: + proc.kill() + log_fp.close() + await fake.stop() + + if failures: + print("\n=== WS E2E FAILURES ===") + for f in failures: + print(" -", f) + return 1 + print("\n=== WS E2E ALL GREEN ===") + return 0 + + +def main() -> int: + return asyncio.run(main_async()) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_compression_cache.py b/tests/test_compression_cache.py index 116c0ce19..66833fe8d 100644 --- a/tests/test_compression_cache.py +++ b/tests/test_compression_cache.py @@ -113,15 +113,20 @@ class TestCompressionCacheFrozenCount: def test_empty_cache_returns_zero(self, cache: CompressionCache) -> None: assert cache.compute_frozen_count([]) == 0 - def test_user_assistant_always_stable(self, cache: CompressionCache) -> None: + def test_user_assistant_stable_with_live_zone_cap(self, cache: CompressionCache) -> None: + """Plain user/assistant turns are individually stable, but the + trailing message is reserved as the live zone — the new turn + cannot be in any provider prefix cache. See docstring on + ``CompressionCache.compute_frozen_count``.""" messages = [ {"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi there"}, {"role": "user", "content": "how are you"}, ] - assert cache.compute_frozen_count(messages) == 3 + # 3 messages structurally stable; cap clamps to len-1 = 2. + assert cache.compute_frozen_count(messages) == 2 - def test_tool_result_with_cache_hit_is_stable(self, cache: CompressionCache) -> None: + def test_tool_result_with_cache_hit_capped_at_live_zone(self, cache: CompressionCache) -> None: tool_content = "tool output data" h = CompressionCache.content_hash(tool_content) cache.store_compressed(h, "compressed tool output", tokens_saved=5) @@ -137,7 +142,9 @@ class TestCompressionCacheFrozenCount: "content": [{"type": "tool_result", "tool_use_id": "t1", "content": tool_content}], }, ] - assert cache.compute_frozen_count(messages) == 3 + # All 3 stable; cap clamps to len-1 = 2 (trailing tool_result is + # the live zone). + assert cache.compute_frozen_count(messages) == 2 def test_tool_result_cache_miss_stops_frozen(self, cache: CompressionCache) -> None: messages = [ @@ -188,9 +195,10 @@ class TestCompressionCacheFrozenCount: }, {"role": "user", "content": "follow up"}, ] - # Without mark_stable, this would stop at msg[1] → frozen=1. - # With stable hash, the walk continues past msg[1] → frozen=3. - assert cache.compute_frozen_count(messages) == 3 + # Without mark_stable, the walk would stop at msg[1] → frozen=1. + # With stable hash, the walk continues past msg[1]; structural + # count = 3, then capped at len-1 = 2 (live-zone reservation). + assert cache.compute_frozen_count(messages) == 2 def test_update_from_result_identical_content_marks_stable( self, cache: CompressionCache @@ -217,7 +225,8 @@ class TestCompressionCacheFrozenCount: h = CompressionCache.content_hash(tool_content) assert h in cache._stable_hashes - # Frozen count should now walk past this tool_result + # Frozen count walks past this tool_result (its hash is stable), + # but the trailing message is still reserved as live zone. messages = [ {"role": "user", "content": "hello"}, { @@ -226,7 +235,7 @@ class TestCompressionCacheFrozenCount: }, {"role": "user", "content": "more stuff"}, ] - assert cache.compute_frozen_count(messages) == 3 + assert cache.compute_frozen_count(messages) == 2 def test_mark_stable_from_messages(self, cache: CompressionCache) -> None: """mark_stable_from_messages records hashes for tool_results.""" diff --git a/tests/test_responses_pyo3_compression.py b/tests/test_responses_pyo3_compression.py index 3e7dc4a4a..e3ffd6fe2 100644 --- a/tests/test_responses_pyo3_compression.py +++ b/tests/test_responses_pyo3_compression.py @@ -48,21 +48,21 @@ class TestPassthroughCases: def test_not_json_passthrough(self): compress = _ensure_binding() body = b"this is not JSON at all" - out, modified = compress(body, "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini") assert out == body assert modified is False def test_no_input_array_passthrough(self): compress = _ensure_binding() body = json.dumps({"model": "gpt-4o-mini"}).encode() - out, modified = compress(body, "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini") assert out == body assert modified is False def test_empty_input_array_passthrough(self): compress = _ensure_binding() body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode() - out, modified = compress(body, "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini") assert out == body assert modified is False @@ -76,7 +76,7 @@ class TestPassthroughCases: "input": [{"type": "message", "role": "user", "content": "hi"}], } ).encode() - out, modified = compress(body, "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(body, "payg", "gpt-4o-mini") assert modified is False # Body should be byte-equal (passthrough, not re-serialized). assert out == body @@ -94,7 +94,7 @@ class TestAuthModeAccepted: compress = _ensure_binding() body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode() # Should not raise on any string input. - out, modified = compress(body, auth_mode, "gpt-4o-mini") + out, modified, _saved, _transforms = compress(body, auth_mode, "gpt-4o-mini") assert isinstance(out, bytes) assert modified is False @@ -105,7 +105,7 @@ class TestModelDefault: def test_empty_model_uses_default(self): compress = _ensure_binding() body = json.dumps({"input": []}).encode() - out, modified = compress(body, "payg", "") + out, modified, _saved, _transforms = compress(body, "payg", "") assert isinstance(out, bytes) assert modified is False @@ -118,12 +118,78 @@ class TestNoExceptionsLeak: def test_garbage_bytes_no_raise(self): compress = _ensure_binding() - out, modified = compress(b"\xff\xfe\x00\xff", "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(b"\xff\xfe\x00\xff", "payg", "gpt-4o-mini") assert modified is False assert out == b"\xff\xfe\x00\xff" def test_empty_body_no_raise(self): compress = _ensure_binding() - out, modified = compress(b"", "payg", "gpt-4o-mini") + out, modified, _saved, _transforms = compress(b"", "payg", "gpt-4o-mini") assert modified is False assert out == b"" + + +class TestTelemetryFields: + """The 4-tuple return surfaces ``tokens_saved`` (sum of + `original_tokens − compressed_tokens` across the manifest's + Compressed outcomes) and ``transforms_applied`` (deduplicated list + of compressor strategy names). The Python proxy uses these to + populate /transformations/feed and the dashboard's per-request log + without recounting tokens. See `crates/headroom-core/src/transforms/ + live_zone.rs::CompressionManifest::tokens_saved` / + `::transforms_applied`.""" + + def test_no_change_returns_zero_savings_and_empty_transforms(self): + compress = _ensure_binding() + body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode() + out, modified, saved, transforms = compress(body, "payg", "gpt-4o-mini") + assert modified is False + assert out == body + assert saved == 0 + assert transforms == [] + + def test_field_types(self): + """Pin the wire shape so downstream callers don't break.""" + compress = _ensure_binding() + body = json.dumps({"model": "gpt-4o-mini", "input": []}).encode() + result = compress(body, "payg", "gpt-4o-mini") + assert isinstance(result, tuple) + assert len(result) == 4 + out, modified, saved, transforms = result + assert isinstance(out, bytes) + assert isinstance(modified, bool) + assert isinstance(saved, int) + assert isinstance(transforms, list) + assert all(isinstance(t, str) for t in transforms) + + def test_large_local_shell_output_compresses_with_telemetry(self): + """End-to-end check: a payload large enough to clear the + per-item byte threshold produces ``modified=True`` plus a + non-zero ``tokens_saved`` and a populated ``transforms`` + list. Mirrors the shape in the Rust crate's + ``large_log_output_compressed`` test.""" + compress = _ensure_binding() + log_body = "".join( + f"[2024-01-01 00:00:00] INFO compile.rs:42 building module foo_{i}\n" + for i in range(400) + ) + assert len(log_body) > 2048 + body = json.dumps( + { + "model": "gpt-4o", + "input": [ + { + "type": "local_shell_call_output", + "call_id": "c1", + "output": log_body, + } + ], + } + ).encode() + out, modified, saved, transforms = compress(body, "payg", "gpt-4o") + assert modified is True + assert saved > 0 + assert transforms, "expected at least one strategy in transforms" + new_doc = json.loads(out) + assert new_doc["input"][0]["type"] == "local_shell_call_output" + assert len(new_doc["input"][0]["output"]) < len(log_body) diff --git a/tests/test_responses_ws_pyo3_compression.py b/tests/test_responses_ws_pyo3_compression.py index 759e9e259..632856f9e 100644 --- a/tests/test_responses_ws_pyo3_compression.py +++ b/tests/test_responses_ws_pyo3_compression.py @@ -60,7 +60,7 @@ def _ws_compress_first_frame( model = (inner.get("model") if isinstance(inner, dict) else None) or "" inner_bytes = json.dumps(inner).encode("utf-8") - new_bytes, modified = compress(inner_bytes, auth_mode_value, model) + new_bytes, modified, _saved, _transforms = compress(inner_bytes, auth_mode_value, model) if not modified: return first_msg_raw, False diff --git a/tests/test_token_headroom_mode.py b/tests/test_token_headroom_mode.py index ef63b50fb..300d8a48a 100644 --- a/tests/test_token_headroom_mode.py +++ b/tests/test_token_headroom_mode.py @@ -87,8 +87,13 @@ class TestMultiTurnCompression: _make_user_msg("now edit it"), ] frozen = cache.compute_frozen_count(messages) - # All 4 messages stable (user, tool_use, tool_result cached, user) - assert frozen == 4 + # First 3 stable; trailing user message ("now edit it") is the + # live zone by construction — it has not been sent upstream + # before, so it cannot be in any provider prefix cache. Cap at + # len - 1 prevents the over-freeze pattern that produced 0 % + # compression for prose-format clients (issue observed + # 2026-05-07 with Cline+DeepSeek). + assert frozen == 3 # apply_cached should swap the content result = cache.apply_cached(messages) @@ -129,9 +134,10 @@ class TestMultiTurnCompression: # Now cache C too cache.store_compressed(CompressionCache.content_hash(code_c), "cc", tokens_saved=100) - # Turn 3: all cached + # Turn 3: all 6 messages structurally stable, but the trailing + # message is reserved as live zone. Frozen prefix = 5. frozen = cache.compute_frozen_count(messages) - assert frozen == 6 # all stable + assert frozen == 5 class TestNoMessageInjection: @@ -305,3 +311,92 @@ class TestUpdateFromResult: msg = _make_user_msg("same content") cache.update_from_result([msg], [msg]) assert cache.get_stats()["entries"] == 0 + + +class TestProseFormatLiveZoneInvariant: + """Cline / OpenClaude / Aider — prose-format clients send tool calls + embedded in plain assistant text and tool results pasted into plain + user messages. There are no `tool_use`, `tool_result`, or + ``role: "tool"`` blocks anywhere in the conversation. + + Pre-fix: ``compute_frozen_count`` walked all messages and found no + "unstable" boundary, returning ``len(messages)``. The pipeline then + froze every message — including the brand-new user turn — leaving + the live zone empty. ContentRouter saw ``saved 0`` on every request. + Bug observed 2026-05-07 with Cline+DeepSeek over /v1/chat/completions + in token mode. + + Post-fix: cap at ``len(messages) - 1`` always reserves the trailing + message as the live zone. These tests lock that invariant. + """ + + def test_pure_user_assistant_turns_leave_live_zone(self): + cache = CompressionCache() + # A 6-turn Cline-shaped conversation: alternating user/assistant + # plain-text. No tool blocks of any kind. + messages = [ + _make_user_msg("system instructions baked into first user msg"), + _make_assistant_msg("ls"), + _make_user_msg("[tool_result]\nfile1.py\nfile2.py\n[/tool_result]"), + _make_assistant_msg("file1.py"), + _make_user_msg("[tool_result]\n\n[/tool_result]"), + _make_user_msg("now please refactor it"), + ] + frozen = cache.compute_frozen_count(messages) + # Pre-fix would return 6 (every plain message is "stable"). + # Post-fix: 6 messages stable, capped at len-1 = 5. + assert frozen == 5 + assert frozen < len(messages), ( + "Live zone must never be empty; trailing user message must " + "always be available for compression" + ) + + def test_single_message_yields_zero_frozen(self): + # Edge case: only the user's first message, nothing to freeze. + cache = CompressionCache() + messages = [_make_user_msg("first turn")] + assert cache.compute_frozen_count(messages) == 0 + + def test_empty_messages_yields_zero(self): + cache = CompressionCache() + assert cache.compute_frozen_count([]) == 0 + + def test_two_messages_first_is_frozen_second_is_live(self): + cache = CompressionCache() + messages = [ + _make_user_msg("turn 1 content"), + _make_user_msg("turn 2 — the live zone"), + ] + # First message structurally stable; trailing is live → 1 frozen. + assert cache.compute_frozen_count(messages) == 1 + + def test_anthropic_format_last_tool_result_is_still_live(self): + """Even when the trailing message is a tool_result whose content + IS in the cache, it stays in the live zone. Trailing == live by + construction; the cache-read decision is upstream's job.""" + cache = CompressionCache() + code = _large_code_content(50) + cache.store_compressed( + CompressionCache.content_hash(code), "compressed code", tokens_saved=200 + ) + messages = [ + _make_user_msg("hi"), + _make_assistant_msg("ok"), + _make_tool_result_msg("t1", code), + ] + frozen = cache.compute_frozen_count(messages) + # Walk gets to 3, cap clamps to 2 (= len-1). + assert frozen == 2 + + def test_openai_format_last_tool_msg_is_live(self): + cache = CompressionCache() + content = "tool output " * 100 + cache.store_compressed( + CompressionCache.content_hash(content), "compressed", tokens_saved=300 + ) + messages = [ + _make_user_msg("run cmd"), + _make_openai_tool_msg("tc1", content), + ] + # Walk: user (stable, 1), tool (cached, 2). Cap → 1. + assert cache.compute_frozen_count(messages) == 1