From eaf5980b4ac48c909a1fa2ef1ace460752ec0918 Mon Sep 17 00:00:00 2001 From: Tejas Chopra Date: Sat, 9 May 2026 13:47:53 -0700 Subject: [PATCH] fix: stabilize codex compression, stats, and proxy lifecycle --- .gitignore | 1 + .../headroom-core/src/transforms/live_zone.rs | 316 +++++-- crates/headroom-core/src/transforms/mod.rs | 5 +- .../src/compression/live_zone_responses.rs | 6 +- crates/headroom-proxy/src/responses_items.rs | 13 +- crates/headroom-proxy/tests/sse_framing.rs | 15 +- crates/headroom-py/src/lib.rs | 14 +- headroom/_version.py | 2 +- headroom/cache/compression_store.py | 150 +++- headroom/cli/proxy.py | 22 + headroom/cli/wrap.py | 233 +++++- headroom/dashboard/templates/dashboard.html | 47 +- headroom/proxy/cost.py | 25 +- headroom/proxy/handlers/openai.py | 791 +++++++++++++++--- headroom/proxy/handlers/streaming.py | 66 +- headroom/proxy/helpers.py | 154 +++- headroom/proxy/prometheus_metrics.py | 56 ++ headroom/proxy/server.py | 126 +-- headroom/transforms/compression_units.py | 176 ++++ headroom/transforms/content_router.py | 11 +- headroom/transforms/pipeline.py | 27 +- tests/test_canonical_pipeline.py | 45 + tests/test_cli/test_unwrap_claude.py | 128 +++ tests/test_cli/test_wrap_codex.py | 157 +++- tests/test_cli/test_wrap_openclaw.py | 22 + tests/test_compression_store.py | 34 +- tests/test_compression_units.py | 160 ++++ tests/test_openai_codex_routing.py | 20 + tests/test_openai_codex_ws_lifecycle.py | 1 + ...test_openai_responses_compression_units.py | 194 +++++ tests/test_proxy_ccr.py | 44 + tests/test_proxy_dashboard_stats_cache.py | 86 +- tests/test_proxy_handler_helpers.py | 12 + tests/test_proxy_openai_responses_bypass.py | 109 +++ ...test_proxy_openai_responses_integration.py | 7 +- tests/test_proxy_streaming_request_logger.py | 75 ++ tests/test_responses_pyo3_compression.py | 29 +- tests/test_responses_ws_pyo3_compression.py | 39 +- 38 files changed, 3002 insertions(+), 416 deletions(-) create mode 100644 headroom/transforms/compression_units.py create mode 100644 tests/test_cli/test_unwrap_claude.py create mode 100644 tests/test_compression_units.py create mode 100644 tests/test_openai_responses_compression_units.py create mode 100644 tests/test_proxy_openai_responses_bypass.py diff --git a/.gitignore b/.gitignore index 09075e8b5..b3173fc3a 100644 --- a/.gitignore +++ b/.gitignore @@ -187,6 +187,7 @@ benchmark_results/ .deepeval/ # Headroom specific +.headroom/ headroom.db headroom_*.db *.jsonl diff --git a/crates/headroom-core/src/transforms/live_zone.rs b/crates/headroom-core/src/transforms/live_zone.rs index 058a72b49..d621d7165 100644 --- a/crates/headroom-core/src/transforms/live_zone.rs +++ b/crates/headroom-core/src/transforms/live_zone.rs @@ -94,7 +94,7 @@ //! parameter in the signature now means later PRs are pure //! implementation swaps, not signature redesigns. -use std::sync::OnceLock; +use std::{collections::HashSet, sync::OnceLock}; use serde::Deserialize; use serde_json::value::RawValue; @@ -148,27 +148,27 @@ pub const DEFAULT_MODEL: &str = "claude-3-5-sonnet-20241022"; // Pinned as `const` rather than a hard-coded `match` so the values are // grep-able and reviewable in one place. -/// JSON-array tool_results below this size route to no-op (1 KiB). -const THRESHOLD_JSON_ARRAY: usize = 1024; +/// JSON-array tool_results below this size route to no-op. +const THRESHOLD_JSON_ARRAY: usize = 512; /// Build / log output below this size routes to no-op (512 B). Logs /// are the most repetitive content type so the threshold is the /// lowest of the bunch. const THRESHOLD_BUILD_OUTPUT: usize = 512; -/// Search-result blocks below this size route to no-op (1 KiB). -const THRESHOLD_SEARCH_RESULTS: usize = 1024; -/// Git-diff blocks below this size route to no-op (1 KiB). -const THRESHOLD_GIT_DIFF: usize = 1024; -/// Source-code blocks below this size route to no-op (2 KiB). Pinned +/// Search-result blocks below this size route to no-op. +const THRESHOLD_SEARCH_RESULTS: usize = 512; +/// Git-diff blocks below this size route to no-op. +const THRESHOLD_GIT_DIFF: usize = 512; +/// Source-code blocks below this size route to no-op. Pinned /// for the future Rust code-compressor port — currently unused /// because `ContentType::SourceCode` short-circuits to no-op above /// the dispatch (see `dispatch_compressor`). -const THRESHOLD_SOURCE_CODE: usize = 2048; -/// Plain-text blocks below this size route to no-op (5 KiB). Pinned +const THRESHOLD_SOURCE_CODE: usize = 512; +/// Plain-text blocks below this size route to no-op. Pinned /// for the future Kompress wiring (PR-B7 follow-up); currently unused. -const THRESHOLD_PLAIN_TEXT: usize = 5120; +const THRESHOLD_PLAIN_TEXT: usize = 512; /// HTML blocks have no compressor; threshold matches plain text so /// when an HTML compressor lands the value is already pinned. -const THRESHOLD_HTML: usize = 5120; +const THRESHOLD_HTML: usize = 512; /// Map a content type to its byte threshold. Returning `usize` rather /// than an `Option` because every variant has a sensible default; @@ -428,6 +428,59 @@ impl CompressionManifest { } } +/// Summarize why a Responses live-zone dispatch made no changes. +/// +/// The proxy uses this to log stable, grep-able reasons instead of the +/// generic `rust_no_compression` bucket. The classification is +/// intentionally coarse: operators want to know whether the dispatcher +/// saw no eligible items, hit a size floor, rejected output as not +/// smaller, or encountered a compressor error. +pub fn summarize_openai_responses_no_change_reason(manifest: &CompressionManifest) -> &'static str { + if manifest.block_outcomes.is_empty() { + return "no_eligible_items"; + } + + let mut saw_no_compression_applied = false; + let mut saw_excluded = false; + let mut saw_below_output_floor = false; + let mut saw_below_plain_text_floor = false; + let mut saw_rejected_not_smaller = false; + let mut saw_compressor_error = false; + + for outcome in &manifest.block_outcomes { + match &outcome.action { + BlockAction::CompressorError { .. } => saw_compressor_error = true, + BlockAction::RejectedNotSmaller { .. } => saw_rejected_not_smaller = true, + BlockAction::BelowByteThreshold { content_type, .. } => { + if *content_type == "output_item" { + saw_below_output_floor = true; + } else { + saw_below_plain_text_floor = true; + } + } + BlockAction::NoCompressionApplied { .. } => saw_no_compression_applied = true, + BlockAction::Excluded { .. } => saw_excluded = true, + BlockAction::Compressed { .. } => {} + } + } + + if saw_compressor_error { + "compressor_error" + } else if saw_rejected_not_smaller { + "rejected_not_smaller" + } else if saw_below_output_floor { + "below_output_floor" + } else if saw_below_plain_text_floor { + "below_plain_text_floor" + } else if saw_excluded { + "excluded_live_zone" + } else if saw_no_compression_applied { + "no_compressible_content" + } else { + "no_change" + } +} + /// Outcome of dispatching the live zone. #[derive(Debug)] pub enum LiveZoneOutcome { @@ -1482,7 +1535,7 @@ mod tests { let out = compress_anthropic_live_zone(&b, 0, AuthMode::Payg, DEFAULT_MODEL).unwrap(); let actions = outcome_block_actions(&out); assert_eq!(actions.len(), 3); - // tool_result with tiny content → BelowByteThreshold (1 byte < 5 KiB plain-text threshold). + // tool_result with tiny content → BelowByteThreshold. assert!(matches!(actions[0], BlockAction::BelowByteThreshold { .. })); assert!(matches!( actions[1], @@ -1490,7 +1543,7 @@ mod tests { reason: ExclusionReason::HotZoneBlockType } )); - // text block with "ok" → BelowByteThreshold (2 bytes < 5 KiB plain-text threshold). + // text block with "ok" → BelowByteThreshold. assert!(matches!(actions[2], BlockAction::BelowByteThreshold { .. })); } @@ -1506,7 +1559,7 @@ mod tests { }; assert_eq!(manifest.block_outcomes.len(), 1); assert_eq!(manifest.block_outcomes[0].block_type, "string_content"); - // 13 bytes of plain text is well below the 5 KiB plain-text threshold. + // 13 bytes of plain text is well below the plain-text threshold. assert!(matches!( manifest.block_outcomes[0].action, BlockAction::BelowByteThreshold { .. } @@ -2153,14 +2206,14 @@ mod openai_chat_tests { // records a `NoCompressionApplied` outcome but never plans a // replacement. // -// Output items must additionally clear a 2 KiB minimum (per spec line +// Output items must additionally clear a 512-byte minimum // 167) before the per-content-type byte threshold even runs. /// Output-item floor below which the Responses dispatcher does not -/// even attempt compression. Per spec PR-C3 §scope. Matches +/// even attempt compression. Matches /// `responses_items::OUTPUT_ITEM_MIN_BYTES`; pinned here too because /// `headroom-core` is independent of the proxy crate. -const RESPONSES_OUTPUT_MIN_BYTES: usize = 2 * 1024; +const RESPONSES_OUTPUT_MIN_BYTES: usize = 512; /// Compress live-zone blocks of an OpenAI Responses request. /// @@ -2182,11 +2235,13 @@ const RESPONSES_OUTPUT_MIN_BYTES: usize = 2 * 1024; /// } /// ``` /// -/// Live zone = the latest item of each compressible kind: -/// `function_call_output`, `local_shell_call_output`, -/// `apply_patch_call_output`, plus the latest `message` text. Earlier -/// items of those kinds are frozen (cached prefix); never rewritten. -/// All other item types pass through verbatim. +/// Live zone = every current-frame output item with a byte-safe +/// string payload (`function_call_output`, `local_shell_call_output`, +/// `apply_patch_call_output`), except CCR retrieval outputs that must +/// reach the model byte-for-byte. +/// Codex commonly batches parallel tool results in one `response.create` +/// frame; those sibling outputs are all live input for the next model +/// turn. All other item types pass through verbatim. /// /// Cache-safety invariant matches the Anthropic / Chat dispatchers: /// bytes outside the rewritten ranges are *literally copied* from the @@ -2217,54 +2272,46 @@ pub fn compress_openai_responses_live_zone( let items_total = items.len(); - // Walk items from the back, tagging the first occurrence of each - // compressible kind. This naturally yields "latest" semantics. - let mut latest_function_output: Option = None; - let mut latest_local_shell_output: Option = None; - let mut latest_apply_patch_output: Option = None; - let mut latest_message: Option = None; - - for (idx, item) in items.iter().enumerate().rev() { - let type_tag = item.get("type").and_then(Value::as_str).unwrap_or(""); - match type_tag { - "function_call_output" if latest_function_output.is_none() => { - latest_function_output = Some(idx); - } - "local_shell_call_output" if latest_local_shell_output.is_none() => { - latest_local_shell_output = Some(idx); - } - "apply_patch_call_output" if latest_apply_patch_output.is_none() => { - latest_apply_patch_output = Some(idx); - } - // Only consider user-role messages for compression. - // Assistant messages are part of the cache hot zone - // (next-turn continuation context). - "message" - if latest_message.is_none() - && item.get("role").and_then(Value::as_str) == Some("user") => - { - latest_message = Some(idx); - } - _ => {} + // Output items in the current Responses frame are live deltas, not + // cached history. Codex often sends several sibling tool outputs + // after parallel local commands; compressing only the last one + // leaves large same-frame payloads untouched. + let mut headroom_retrieve_call_ids: HashSet<&str> = HashSet::new(); + for item in items { + if item.get("type").and_then(Value::as_str) != Some("function_call") { + continue; } - // Early-exit if we've found everything we care about. - if latest_function_output.is_some() - && latest_local_shell_output.is_some() - && latest_apply_patch_output.is_some() - && latest_message.is_some() - { - break; + let name = item.get("name").and_then(Value::as_str).unwrap_or(""); + if name == "headroom_retrieve" || name.ends_with("__headroom_retrieve") { + if let Some(call_id) = item.get("call_id").and_then(Value::as_str) { + headroom_retrieve_call_ids.insert(call_id); + } } } - let candidates: &[(Option, &str)] = &[ - (latest_function_output, "function_call_output"), - (latest_local_shell_output, "local_shell_call_output"), - (latest_apply_patch_output, "apply_patch_call_output"), - (latest_message, "message"), - ]; + let mut output_candidates: Vec<(usize, &str)> = Vec::new(); + let latest_message: Option = None; - if candidates.iter().all(|(idx, _)| idx.is_none()) { + for (idx, item) in items.iter().enumerate() { + let type_tag = item.get("type").and_then(Value::as_str).unwrap_or(""); + match type_tag { + "function_call_output" | "local_shell_call_output" | "apply_patch_call_output" => { + let call_id = item.get("call_id").and_then(Value::as_str); + if call_id.is_some_and(|id| headroom_retrieve_call_ids.contains(id)) { + continue; + } + output_candidates.push((idx, type_tag)); + } + _ => {} + } + } + + let mut candidates = output_candidates; + if let Some(idx) = latest_message { + candidates.push((idx, "message")); + } + + if candidates.is_empty() { return Ok(LiveZoneOutcome::NoChange { manifest: CompressionManifest { messages_total: items_total, @@ -2279,10 +2326,9 @@ pub fn compress_openai_responses_live_zone( // one slot (output items have a single string field; messages // have a single text content slot). let mut all_slots: Vec<(usize, ResponsesPlanSlot)> = Vec::new(); - for (maybe_idx, kind_tag) in candidates { - let Some(idx) = maybe_idx else { continue }; - match plan_responses_item(body_raw, *idx, kind_tag) { - Ok(Some(slot)) => all_slots.push((*idx, slot)), + for (idx, kind_tag) in candidates { + match plan_responses_item(body_raw, idx, kind_tag) { + Ok(Some(slot)) => all_slots.push((idx, slot)), Ok(None) => {} Err(_) => { // Body shape doesn't match what we expect for this @@ -2308,7 +2354,7 @@ pub fn compress_openai_responses_live_zone( let mut replacements: Vec = Vec::new(); for (msg_idx, slot) in all_slots { - // Output items must clear the 2 KiB floor BEFORE the + // Output items must clear the response-output floor BEFORE the // per-content-type threshold even runs. This is on top of the // existing per-block byte-threshold gate. if slot.is_output_item && slot.content_text.len() < RESPONSES_OUTPUT_MIN_BYTES { @@ -2368,7 +2414,7 @@ pub fn compress_openai_responses_live_zone( /// Per-kind plan slot for the Responses dispatcher. Mirrors /// `OpenAiPlanSlot` but tracks whether the slot is an `*_output` item -/// (so the 2 KiB floor only applies there, not to `message` text). +/// (so the response-output floor only applies there, not to `message` text). struct ResponsesPlanSlot { block_index: Option, block_type: String, @@ -2376,7 +2422,7 @@ struct ResponsesPlanSlot { content_byte_range: (usize, usize), /// True when the slot is one of `function_call_output`, /// `local_shell_call_output`, `apply_patch_call_output`. Used to - /// gate the 2 KiB output-item floor. + /// gate the response-output floor. is_output_item: bool, } @@ -2578,9 +2624,9 @@ mod openai_responses_tests { } #[test] - fn output_below_2kb_skipped() { - // 1 KB output → below the output-item floor. - let small = "x".repeat(1024); + fn output_below_512b_skipped() { + // 256 B output → below the output-item floor. + let small = "x".repeat(256); let b = body(json!({ "model": "gpt-4o", "input": [ @@ -2598,8 +2644,8 @@ mod openai_responses_tests { threshold_bytes, } => { assert_eq!(*content_type, "output_item"); - assert_eq!(*byte_count, 1024); - assert_eq!(*threshold_bytes, 2048); + assert_eq!(*byte_count, 256); + assert_eq!(*threshold_bytes, RESPONSES_OUTPUT_MIN_BYTES); } other => panic!("expected BelowByteThreshold, got {other:?}"), } @@ -2609,10 +2655,10 @@ mod openai_responses_tests { } #[test] - fn picks_latest_function_output_only() { - // Two function_call_output items; only the latest is in the - // live zone. Both small, so neither compresses, but the - // manifest must show one slot, not two. + fn plans_all_same_frame_function_outputs() { + // Codex can batch parallel tool results in a single + // response.create frame. They are all current-frame live + // inputs, so each byte-safe output string gets a slot. let b = body(json!({ "model": "gpt-4o", "input": [ @@ -2631,8 +2677,45 @@ mod openai_responses_tests { .iter() .filter(|b| b.block_type == "function_call_output") .collect(); - assert_eq!(outputs.len(), 1); - assert_eq!(outputs[0].message_index, 2); + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].message_index, 0); + assert_eq!(outputs[1].message_index, 2); + } + + #[test] + fn compresses_multiple_same_frame_outputs() { + let mut first = String::new(); + let mut second = String::new(); + for i in 0..400 { + first.push_str(&format!( + "./src/foo_{i}.rs:12: error[E0308]: mismatched types in module foo_{i}\n" + )); + second.push_str(&format!( + "./tests/bar_{i}.rs:44: warning: unused variable in test bar_{i}\n" + )); + } + let b = body(json!({ + "model": "gpt-4o", + "input": [ + {"type": "function_call_output", "call_id": "c1", "output": first}, + {"type": "function_call", "call_id": "c2", "name": "f", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "c2", "output": second}, + ] + })); + let out = compress_openai_responses_live_zone(&b, AuthMode::Payg, "gpt-4o").unwrap(); + let manifest = match &out { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => manifest, + }; + let compressed_outputs = manifest + .block_outcomes + .iter() + .filter(|b| { + b.block_type == "function_call_output" + && matches!(b.action, BlockAction::Compressed { .. }) + }) + .count(); + assert_eq!(compressed_outputs, 2, "{manifest:?}"); } #[test] @@ -2702,7 +2785,7 @@ mod openai_responses_tests { } #[test] - fn message_user_content_planned() { + fn message_user_content_not_in_live_zone() { let b = body(json!({ "model": "gpt-4o", "input": [ @@ -2713,8 +2796,35 @@ mod openai_responses_tests { let out = compress_openai_responses_live_zone(&b, AuthMode::Payg, "gpt-4o").unwrap(); match &out { LiveZoneOutcome::NoChange { manifest } => { - assert_eq!(manifest.block_outcomes.len(), 1); - assert_eq!(manifest.block_outcomes[0].block_type, "message_input_text"); + assert!(manifest.block_outcomes.is_empty()); + } + _ => panic!("expected NoChange"), + } + } + + #[test] + fn headroom_retrieve_output_not_in_live_zone() { + let retrieved = "retrieved original content ".repeat(100); + let b = body(json!({ + "model": "gpt-4o", + "input": [ + { + "type": "function_call", + "call_id": "call_retrieve", + "name": "mcp__headroom__headroom_retrieve", + "arguments": "{}" + }, + { + "type": "function_call_output", + "call_id": "call_retrieve", + "output": retrieved + } + ] + })); + let out = compress_openai_responses_live_zone(&b, AuthMode::Payg, "gpt-4o").unwrap(); + match &out { + LiveZoneOutcome::NoChange { manifest } => { + assert!(manifest.block_outcomes.is_empty()); } _ => panic!("expected NoChange"), } @@ -2739,4 +2849,36 @@ mod openai_responses_tests { _ => panic!("expected NoChange"), } } + + #[test] + fn no_change_reason_empty_input_is_no_eligible_items() { + let manifest = CompressionManifest::empty(); + assert_eq!( + summarize_openai_responses_no_change_reason(&manifest), + "no_eligible_items" + ); + } + + #[test] + fn no_change_reason_prefers_output_floor() { + let manifest = CompressionManifest { + messages_total: 1, + messages_below_frozen_floor: 0, + latest_user_message_index: Some(0), + block_outcomes: vec![BlockOutcome { + message_index: 0, + block_index: None, + block_type: "function_call_output".to_string(), + action: BlockAction::BelowByteThreshold { + content_type: "output_item", + byte_count: 1024, + threshold_bytes: RESPONSES_OUTPUT_MIN_BYTES, + }, + }], + }; + assert_eq!( + summarize_openai_responses_no_change_reason(&manifest), + "below_output_floor" + ); + } } diff --git a/crates/headroom-core/src/transforms/mod.rs b/crates/headroom-core/src/transforms/mod.rs index 5546d45b8..6c909d96f 100644 --- a/crates/headroom-core/src/transforms/mod.rs +++ b/crates/headroom-core/src/transforms/mod.rs @@ -40,8 +40,9 @@ pub use diff_compressor::{ }; pub use live_zone::{ compress_anthropic_live_zone, compress_openai_chat_live_zone, - compress_openai_responses_live_zone, AuthMode, BlockAction, BlockOutcome, CompressionManifest, - ExclusionReason, LiveZoneError, LiveZoneOutcome, + compress_openai_responses_live_zone, summarize_openai_responses_no_change_reason, AuthMode, + BlockAction, BlockOutcome, CompressionManifest, ExclusionReason, LiveZoneError, + LiveZoneOutcome, }; pub use log_compressor::{ LogCompressionResult, LogCompressor, LogCompressorConfig, LogCompressorStats, LogFormat, diff --git a/crates/headroom-proxy/src/compression/live_zone_responses.rs b/crates/headroom-proxy/src/compression/live_zone_responses.rs index 213f1fd3f..ef4e58378 100644 --- a/crates/headroom-proxy/src/compression/live_zone_responses.rs +++ b/crates/headroom-proxy/src/compression/live_zone_responses.rs @@ -35,7 +35,8 @@ use bytes::Bytes; use headroom_core::auth_mode::AuthMode as RequestAuthMode; use headroom_core::transforms::live_zone::DEFAULT_MODEL; use headroom_core::transforms::{ - compress_openai_responses_live_zone, BlockAction, LiveZoneError, LiveZoneOutcome, + compress_openai_responses_live_zone, summarize_openai_responses_no_change_reason, BlockAction, + LiveZoneError, LiveZoneOutcome, }; use serde_json::Value; @@ -146,6 +147,7 @@ pub fn compress_openai_responses_request( // the rationale — same wiring on the OpenAI Responses path. match compress_openai_responses_live_zone(&dispatch_body, auth_mode.into(), model) { Ok(LiveZoneOutcome::NoChange { manifest }) => { + let reason = summarize_openai_responses_no_change_reason(&manifest); tracing::info!( event = "compression_decision", request_id = %request_id, @@ -153,7 +155,7 @@ pub fn compress_openai_responses_request( method = "POST", compression_mode = mode.as_str(), decision = "no_change", - reason = "no_block_compressed", + reason = reason, body_bytes = body.len(), items_total = manifest.messages_total, latest_user_message_index = ?manifest.latest_user_message_index, diff --git a/crates/headroom-proxy/src/responses_items.rs b/crates/headroom-proxy/src/responses_items.rs index 4aabd710a..0b6ec999c 100644 --- a/crates/headroom-proxy/src/responses_items.rs +++ b/crates/headroom-proxy/src/responses_items.rs @@ -25,7 +25,7 @@ //! (`function_call_output`, `local_shell_call_output`, //! `apply_patch_call_output`) are eligible for live-zone //! compression — but only the *latest* of each kind, only above the -//! 2 KiB output-item floor, and only when the per-content-type +//! output-item floor, and only when the per-content-type //! compressor agrees the result shrinks the token count. //! - **Unknown item types** are logged at warn level and preserved //! byte-for-byte via `serde_json::value::RawValue`. This is the @@ -122,7 +122,7 @@ pub struct ApplyPatchOperation<'a> { /// the wire; never parse it as JSON inside the proxy. The model /// built it; the model parses it. /// - `output` on `*_output` items is a string. Compressors run only -/// on the latest of each kind, and only above the 2 KiB output-item +/// on the latest of each kind, and only above the output-item /// floor (see [`OUTPUT_ITEM_MIN_BYTES`]). /// - String fields use `Cow<'a, str>` so escape-bearing JSON values /// (e.g. `"{\"q\":\"hello\"}"`) succeed without allocation on the @@ -180,7 +180,7 @@ pub enum ResponseItem<'a> { /// Function tool output. `output` is the string the proxy may /// compress when this is the latest `function_call_output` and - /// the bytes exceed the 2 KiB floor. + /// the bytes exceed the output-item floor. #[serde(rename = "function_call_output")] FunctionCallOutput { #[serde(default, borrow)] @@ -379,11 +379,10 @@ impl<'a> ResponseItem<'a> { } /// Per-item-type minimum bytes before the live-zone dispatcher even -/// inspects an `*_output` payload. Per spec PR-C3 §scope: 2 KiB. +/// inspects an `*_output` payload. /// Per-content-type thresholds from `transforms::live_zone` still -/// apply on top of this floor (e.g. logs at 512 B → still skipped -/// because output items must clear 2 KiB first). -pub const OUTPUT_ITEM_MIN_BYTES: usize = 2 * 1024; +/// apply on top of this floor. +pub const OUTPUT_ITEM_MIN_BYTES: usize = 512; /// Two-pass result: a typed view alongside the byte-faithful raw /// slice. Lifetime ties to the underlying request body. Always diff --git a/crates/headroom-proxy/tests/sse_framing.rs b/crates/headroom-proxy/tests/sse_framing.rs index e04208859..bf487e4cb 100644 --- a/crates/headroom-proxy/tests/sse_framing.rs +++ b/crates/headroom-proxy/tests/sse_framing.rs @@ -147,20 +147,17 @@ fn chunk_boundary_inside_event_name() { // ───────────────────────────── property test ───────────────────────── // // The framer must NEVER panic on arbitrary byte input. TCP can hand us -// anything — partial codepoints, NUL bytes, fuzz noise. 100K random -// inputs is the project default for "no panic" parser invariants per -// `feedback_realignment_build_constraints.md`. The test does not assert -// what the framer yields, only that pushing arbitrary bytes and pulling -// events terminates without panic. +// anything — partial codepoints, NUL bytes, fuzz noise. Keep this in the +// same order of magnitude as the other Rust parser fuzz tests so +// `cargo test --workspace` remains practical in CI. use proptest::prelude::*; proptest! { #![proptest_config(ProptestConfig { - cases: 100_000, - // Disable shrinking timeout — 100K trivial cases run fast on - // CI and we want the fuzzer to actually shrink any panic it - // finds (vs. give up early). + cases: 4_096, + // We want the fuzzer to shrink any panic it finds instead of + // giving up early. max_shrink_iters: 1024, ..ProptestConfig::default() })] diff --git a/crates/headroom-py/src/lib.rs b/crates/headroom-py/src/lib.rs index 9dac7b9dc..e47c9b649 100644 --- a/crates/headroom-py/src/lib.rs +++ b/crates/headroom-py/src/lib.rs @@ -30,6 +30,7 @@ use headroom_core::transforms::tag_protector::{ use headroom_core::transforms::{ compress_openai_responses_live_zone as rust_compress_openai_responses_live_zone, detect as rust_detect_chain, is_json_array_of_dicts as rust_is_json_array_of_dicts, + summarize_openai_responses_no_change_reason as rust_summarize_openai_responses_no_change_reason, AuthMode as RustLiveZoneAuthMode, ContentType as RustContentType, DetectionResult as RustDetectionResult, DiffCompressionResult, DiffCompressor, DiffCompressorConfig, DiffCompressorStats, LiveZoneOutcome, @@ -1498,7 +1499,7 @@ fn compress_openai_responses_live_zone( body: &[u8], auth_mode: &str, model: &str, -) -> (Py, bool, u64, Vec) { +) -> (Py, bool, u64, Vec, Option) { let mode = match auth_mode.to_ascii_lowercase().as_str() { "payg" => RustLiveZoneAuthMode::Payg, "oauth" => RustLiveZoneAuthMode::OAuth, @@ -1519,11 +1520,13 @@ fn compress_openai_responses_live_zone( .into_iter() .map(String::from) .collect(); + let reason = rust_summarize_openai_responses_no_change_reason(&manifest).to_string(); ( PyBytes::new_bound(py, body).unbind(), false, saved, transforms, + Some(reason), ) } Ok(LiveZoneOutcome::Modified { new_body, manifest }) => { @@ -1541,12 +1544,19 @@ fn compress_openai_responses_live_zone( true, saved, transforms, + None, ) } Err(_) => { // BodyNotJson / NoMessagesArray are non-fatal: nothing to // compress, fall through to passthrough byte-for-byte. - (PyBytes::new_bound(py, body).unbind(), false, 0, Vec::new()) + ( + PyBytes::new_bound(py, body).unbind(), + false, + 0, + Vec::new(), + Some("dispatch_error".to_string()), + ) } } } diff --git a/headroom/_version.py b/headroom/_version.py index bbe9da5e6..07d29900a 100644 --- a/headroom/_version.py +++ b/headroom/_version.py @@ -1,3 +1,3 @@ """Package version metadata.""" -__version__ = "0.5.25" +__version__ = "0.9.1" diff --git a/headroom/cache/compression_store.py b/headroom/cache/compression_store.py index 810d1a07c..31d6d042b 100644 --- a/headroom/cache/compression_store.py +++ b/headroom/cache/compression_store.py @@ -409,12 +409,7 @@ class CompressionStore: if entry is None: return [] - try: - items = json.loads(entry.original_content) - if not isinstance(items, list): - return [] - except json.JSONDecodeError: - return [] + items = self._search_items_from_original(entry.original_content) if not items: return [] @@ -450,6 +445,149 @@ class CompressionStore: return results + def _search_items_from_original(self, original_content: str) -> list[Any]: + """Normalize cached originals into searchable items. + + CCR producers store different shapes: + - SmartCrusher/search-style paths usually store JSON arrays. + - Kompress stores the original plain text. + - Some callers store JSON objects or scalar JSON values. + + Search should work for all of them. Preserve the legacy JSON-array + result shape, but fall back to structured text chunks for everything + else so `headroom_retrieve(hash, query=...)` can find plain-text + originals. + """ + + try: + parsed = json.loads(original_content) + except json.JSONDecodeError: + return self._plain_text_search_items(original_content) + + if isinstance(parsed, list): + return parsed + if isinstance(parsed, dict): + return self._json_object_search_items(parsed) + if isinstance(parsed, str): + return self._plain_text_search_items(parsed) + if parsed is None: + return [] + return [{"type": "json_scalar", "value": parsed}] + + def _json_object_search_items(self, value: dict[str, Any]) -> list[dict[str, Any]]: + """Return searchable leaf records for a JSON object.""" + + items: list[dict[str, Any]] = [] + + def walk(node: Any, path: str) -> None: + if isinstance(node, dict): + for key, child in node.items(): + child_path = f"{path}.{key}" if path else str(key) + walk(child, child_path) + return + if isinstance(node, list): + for idx, child in enumerate(node): + walk(child, f"{path}[{idx}]") + return + if node is None: + return + items.append({"type": "json_leaf", "path": path, "value": node}) + + walk(value, "") + if items: + return items + return [{"type": "json_object", "value": value}] + + def _plain_text_search_items(self, text: str) -> list[dict[str, Any]]: + """Chunk arbitrary text into searchable records. + + Line-aware chunks work well for logs/source. Word-window chunks handle + Kompress originals, which are often long single-line text blobs. + """ + + if not text or not text.strip(): + return [] + + normalized = text.replace("\r\n", "\n").replace("\r", "\n") + lines = normalized.split("\n") + if len(lines) > 1: + return self._line_text_search_items(lines) + + words = normalized.split() + if not words: + return [] + max_words = 350 + overlap_words = 50 + if len(words) <= max_words: + return [ + { + "type": "text", + "text": normalized, + "chunk_index": 0, + "word_start": 1, + "word_end": len(words), + } + ] + + items: list[dict[str, Any]] = [] + start = 0 + chunk_index = 0 + step = max_words - overlap_words + while start < len(words): + end = min(len(words), start + max_words) + items.append( + { + "type": "text", + "text": " ".join(words[start:end]), + "chunk_index": chunk_index, + "word_start": start + 1, + "word_end": end, + } + ) + if end == len(words): + break + start += step + chunk_index += 1 + return items + + @staticmethod + def _line_text_search_items(lines: list[str]) -> list[dict[str, Any]]: + max_chars = 2000 + items: list[dict[str, Any]] = [] + current: list[str] = [] + line_start = 1 + char_count = 0 + + for idx, line in enumerate(lines, start=1): + line_len = len(line) + 1 + if current and char_count + line_len > max_chars: + items.append( + { + "type": "text", + "text": "\n".join(current), + "chunk_index": len(items), + "line_start": line_start, + "line_end": idx - 1, + } + ) + current = [] + line_start = idx + char_count = 0 + current.append(line) + char_count += line_len + + if current: + items.append( + { + "type": "text", + "text": "\n".join(current), + "chunk_index": len(items), + "line_start": line_start, + "line_end": len(lines), + } + ) + return items + def _get_entry_for_search( self, hash_key: str, diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index f0a8989d9..afc473a58 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -6,6 +6,7 @@ from typing import Any import click +from headroom import paths as _paths from headroom.providers.registry import resolve_api_overrides, resolve_api_targets from headroom.proxy.modes import PROXY_MODE_TOKEN, normalize_proxy_mode @@ -191,6 +192,19 @@ def _get_env_bool(name: str, default: bool) -> bool: is_flag=True, help="Enable full message logging (request/response content stored for live feed)", ) +@click.option( + "--codex-wire-debug", + is_flag=True, + help="Enable local Codex wire snapshots and matching proxy.log frame traces.", +) +@click.option( + "--codex-wire-debug-dir", + default=None, + help=( + "Directory for Codex wire snapshots (default: " + "~/.headroom/logs/codex_wire or workspace .headroom/logs/codex_wire)." + ), +) @click.option( "--budget", type=float, @@ -393,6 +407,8 @@ def proxy( anthropic_pre_upstream_memory_context_timeout_seconds: float | None, log_file: str | None, log_messages: bool, + codex_wire_debug: bool, + codex_wire_debug_dir: str | None, budget: float | None, code_aware_flag: bool | None, code_graph: bool, @@ -498,6 +514,12 @@ def proxy( if no_telemetry: os.environ["HEADROOM_TELEMETRY"] = "off" + if codex_wire_debug or codex_wire_debug_dir: + os.environ["HEADROOM_CODEX_WIRE_DEBUG"] = "1" + os.environ["HEADROOM_CODEX_WIRE_DEBUG_DIR"] = codex_wire_debug_dir or str( + _paths.codex_wire_debug_dir() + ) + # Stateless mode: suppress TOIN filesystem persistence if is_stateless: os.environ["HEADROOM_TOIN_BACKEND"] = "none" diff --git a/headroom/cli/wrap.py b/headroom/cli/wrap.py index ea54e4d76..fc0eb0c2a 100644 --- a/headroom/cli/wrap.py +++ b/headroom/cli/wrap.py @@ -188,6 +188,7 @@ def _start_proxy( stdout=log_file, stderr=log_file, env=proxy_env, + start_new_session=os.name == "posix", ) # Wait for proxy to be ready (up to 45 seconds). @@ -242,6 +243,74 @@ def _setup_rtk(verbose: bool = False) -> Path | None: return rtk_path +def _remove_claude_rtk_hooks(settings_path: Path | None = None) -> bool: + """Remove Headroom/rtk-managed Claude hook entries from settings.json. + + `rtk init --global --auto-patch` installs a Claude PreToolUse hook that + points at an ``rtk-rewrite`` script. Unwrap should remove that hook without + touching unrelated Claude settings or user-authored hooks. + """ + + path = settings_path or (Path.home() / ".claude" / "settings.json") + if not path.exists(): + return False + + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return False + if not isinstance(payload, dict): + return False + + hooks = payload.get("hooks") + if not isinstance(hooks, dict): + return False + + changed = False + for event, entries in list(hooks.items()): + if not isinstance(entries, list): + continue + retained_entries: list[Any] = [] + for entry in entries: + if not isinstance(entry, dict): + retained_entries.append(entry) + continue + hook_items = entry.get("hooks") + if not isinstance(hook_items, list): + retained_entries.append(entry) + continue + retained_hooks = [ + item + for item in hook_items + if not ( + isinstance(item, dict) and "rtk-rewrite" in str(item.get("command", "")).lower() + ) + ] + if len(retained_hooks) != len(hook_items): + changed = True + if retained_hooks: + retained_entries.append({**entry, "hooks": retained_hooks}) + elif len(retained_hooks) == len(hook_items): + retained_entries.append(entry) + else: + changed = True + if retained_entries: + hooks[event] = retained_entries + else: + del hooks[event] + changed = True + + if not changed: + return False + + if hooks: + payload["hooks"] = hooks + else: + payload.pop("hooks", None) + path.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return True + + def _setup_headroom_mcp( registrar: Any, port: int, *, verbose: bool = False, force: bool = False ) -> None: @@ -454,6 +523,8 @@ _MEMORY_AGENTS_MARKER = "" # Codex config injection markers _CODEX_TOP_LEVEL_MARKER = "# --- Headroom proxy (auto-injected by headroom wrap codex) ---" _CODEX_END_MARKER = "# --- end Headroom ---" +_CODEX_MCP_MARKER = "# --- Headroom MCP server ---" +_CODEX_MCP_END = "# --- end Headroom MCP server ---" # File name used for the pre-wrap snapshot of ~/.codex/config.toml. The # snapshot lets `headroom unwrap codex` restore the exact prior state, even # if the user had their own `model_provider` / `[model_providers.*]` config @@ -469,7 +540,7 @@ def _codex_config_paths() -> tuple[Path, Path]: return config_file, backup_file -def _strip_codex_headroom_blocks(content: str) -> str: +def _strip_codex_headroom_blocks(content: str, *, remove_mcp: bool = False) -> str: """Remove all Headroom-managed blocks from a Codex ``config.toml`` string. Returns the cleaned content. Safe to call on content that never contained @@ -478,19 +549,25 @@ def _strip_codex_headroom_blocks(content: str) -> str: """ import re - # Remove any top-level-marker → end-marker span, possibly repeated. - while _CODEX_TOP_LEVEL_MARKER in content and _CODEX_END_MARKER in content: - start = content.index(_CODEX_TOP_LEVEL_MARKER) - end_idx = content.index(_CODEX_END_MARKER, start) - if end_idx < start: - break - end = end_idx + len(_CODEX_END_MARKER) - content = content[:start].rstrip("\n") + "\n" + content[end:].lstrip("\n") + def _remove_marker_span(text: str, start_marker: str, end_marker: str) -> str: + while start_marker in text and end_marker in text: + start = text.index(start_marker) + end_idx = text.index(end_marker, start) + if end_idx < start: + break + end = end_idx + len(end_marker) + text = text[:start].rstrip("\n") + "\n" + text[end:].lstrip("\n") + text = text.replace(start_marker + "\n", "") + text = text.replace(end_marker + "\n", "") + return text - # Remove any stale top-level marker or end marker that lost its partner - # (e.g. a crashed prior wrap). - content = content.replace(_CODEX_TOP_LEVEL_MARKER + "\n", "") - content = content.replace(_CODEX_END_MARKER + "\n", "") + # Remove any top-level-marker → end-marker span, possibly repeated. + content = _remove_marker_span(content, _CODEX_TOP_LEVEL_MARKER, _CODEX_END_MARKER) + + if remove_mcp: + # Remove Headroom-managed MCP blocks written by `wrap codex`. + content = _remove_marker_span(content, _CODEX_MCP_MARKER, _CODEX_MCP_END) + content = _remove_marker_span(content, _MEMORY_MCP_MARKER, _MEMORY_MCP_END) # Strip any leftover top-level keys that older (or crashed) versions of # `wrap codex` may have written outside the marker block. @@ -673,7 +750,7 @@ def _restore_codex_provider_config() -> tuple[str, Path]: if config_file.exists(): original = config_file.read_text() if _CODEX_TOP_LEVEL_MARKER in original or _CODEX_END_MARKER in original: - cleaned = _strip_codex_headroom_blocks(original) + cleaned = _strip_codex_headroom_blocks(original, remove_mcp=True) if not cleaned.strip(): # Nothing left but Headroom content — remove the file entirely # so Codex falls back to its default config. @@ -845,6 +922,56 @@ def _kill_proxy_by_pid(pid: int, port: int) -> bool: return False +def _stop_local_proxy_for_unwrap(port: int) -> str: + """Stop a local Headroom proxy for durable unwrap commands. + + Returns a status string: + * ``"stopped"``: a Headroom proxy was identified and stopped. + * ``"not_running"``: nothing is listening on the requested port. + * ``"unidentified"``: something is listening, but it did not expose + Headroom's health/config payload, so we did not kill it. + * ``"no_pid"``: the service looked like Headroom but did not expose a PID. + * ``"failed"``: a PID was found but the port stayed bound after stop. + """ + + if not _check_proxy(port): + return "not_running" + + running_config = _query_proxy_config(port) + if running_config is None: + return "unidentified" + + proxy_pid = running_config.get("pid") + if proxy_pid is None: + return "no_pid" + + try: + pid = int(proxy_pid) + except (TypeError, ValueError): + return "no_pid" + + return "stopped" if _kill_proxy_by_pid(pid, port) else "failed" + + +def _echo_unwrap_proxy_stop_status(status: str, port: int) -> None: + """Print a human-readable proxy stop result for unwrap commands.""" + + if status == "stopped": + click.echo(f" Stopped local Headroom proxy on port {port}.") + elif status == "not_running": + click.echo(f" No local Headroom proxy detected on port {port}.") + elif status == "unidentified": + click.echo( + f" Warning: port {port} is in use, but it did not look like Headroom; left it running." + ) + elif status == "no_pid": + click.echo( + f" Warning: Headroom proxy on port {port} did not expose a PID; left it running." + ) + else: + click.echo(f" Warning: failed to stop Headroom proxy on port {port}; stop it manually.") + + def _find_persistent_manifest(port: int) -> Any: """Return a matching persistent deployment manifest for the requested port.""" from headroom.install.state import list_manifests @@ -1069,6 +1196,12 @@ def _make_cleanup(proxy_proc_holder: list, port: int = 8787) -> Any: return cleanup +def _ignore_child_sigint(signum: int | None = None, frame: Any = None) -> None: + """Keep the wrapper alive when Ctrl-C is intended for the child CLI.""" + + return None + + def _launch_tool( binary: str, args: tuple, @@ -1090,7 +1223,7 @@ def _launch_tool( """Common logic: start proxy, launch tool, clean up.""" proxy_holder: list[subprocess.Popen | None] = [None] cleanup = _make_cleanup(proxy_holder, port) - signal.signal(signal.SIGINT, cleanup) + signal.signal(signal.SIGINT, _ignore_child_sigint) signal.signal(signal.SIGTERM, cleanup) try: @@ -1431,7 +1564,7 @@ def claude( # Setup rtk before launching (Claude-specific) proxy_holder: list[subprocess.Popen | None] = [None] cleanup = _make_cleanup(proxy_holder, port) - signal.signal(signal.SIGINT, cleanup) + signal.signal(signal.SIGINT, _ignore_child_sigint) signal.signal(signal.SIGTERM, cleanup) # Memory sync BEFORE proxy startup — sync headroom DB ↔ Claude's files @@ -1526,6 +1659,62 @@ def claude( cleanup() +# ============================================================================= +# Claude Code (unwrap) +# ============================================================================= + + +@unwrap.command("claude") +@click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") +@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy") +@click.option("--keep-mcp", is_flag=True, help="Keep Headroom MCP registrations") +@click.option("--keep-rtk", is_flag=True, help="Keep rtk Claude hooks") +def unwrap_claude( + port: int, + no_stop_proxy: bool, + keep_mcp: bool, + keep_rtk: bool, +) -> None: + """Undo durable setup from ``headroom wrap claude``.""" + click.echo() + click.echo(" ╔═══════════════════════════════════════════════╗") + click.echo(" ║ HEADROOM UNWRAP: CLAUDE ║") + click.echo(" ╚═══════════════════════════════════════════════╝") + click.echo() + + if not keep_mcp: + from headroom.mcp_registry import ClaudeRegistrar + + registrar = ClaudeRegistrar() + if registrar.detect(): + removed_headroom = registrar.unregister_server("headroom") + removed_code_graph = registrar.unregister_server(_CBM_MCP_SERVER_NAME) + if removed_headroom: + click.echo(" Removed Headroom MCP retrieve tool from Claude.") + else: + click.echo(" Headroom MCP retrieve tool was not registered in Claude.") + if removed_code_graph: + click.echo(" Removed code graph MCP server from Claude.") + else: + click.echo(" Claude Code not detected; skipped MCP cleanup.") + else: + click.echo(" Kept Claude MCP registrations (--keep-mcp).") + + if not keep_rtk: + if _remove_claude_rtk_hooks(): + click.echo(" Removed rtk Claude hook from settings.json.") + else: + click.echo(" No rtk Claude hook found in settings.json.") + else: + click.echo(" Kept rtk Claude hooks (--keep-rtk).") + + click.echo() + click.echo("✓ Claude is no longer durably wrapped by Headroom.") + if not no_stop_proxy: + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) + click.echo() + + # ============================================================================= # GitHub Copilot CLI # ============================================================================= @@ -2320,11 +2509,15 @@ def openclaw( @unwrap.command("openclaw") +@click.option("--proxy-port", default=8787, type=int, help="Headroom proxy port") +@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy") @click.option("--no-restart", is_flag=True, help="Do not restart OpenClaw gateway at the end") @click.option("--verbose", "-v", is_flag=True, help="Verbose output") @click.option("--prepare-only", is_flag=True, hidden=True) @click.option("--existing-entry-json", default=None, hidden=True) def unwrap_openclaw( + proxy_port: int, + no_stop_proxy: bool, no_restart: bool, verbose: bool, prepare_only: bool, @@ -2384,6 +2577,8 @@ def unwrap_openclaw( click.echo("✓ OpenClaw Headroom wrap removed.") click.echo(" Plugin: headroom (installed, disabled)") click.echo(" Slot: plugins.slots.contextEngine = legacy") + if not no_stop_proxy: + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(proxy_port), proxy_port) click.echo() @@ -2393,7 +2588,9 @@ def unwrap_openclaw( @unwrap.command("codex") -def unwrap_codex() -> None: +@click.option("--port", "-p", default=8787, type=int, help="Proxy port (default: 8787)") +@click.option("--no-stop-proxy", is_flag=True, help="Do not stop the local Headroom proxy") +def unwrap_codex(port: int, no_stop_proxy: bool) -> None: """Undo ``headroom wrap codex`` edits to ``~/.codex/config.toml``. Behaviour: @@ -2430,4 +2627,6 @@ def unwrap_codex() -> None: click.echo() click.echo("✓ Codex is no longer routed through the Headroom proxy.") + if not no_stop_proxy: + _echo_unwrap_proxy_stop_status(_stop_local_proxy_for_unwrap(port), port) click.echo() diff --git a/headroom/dashboard/templates/dashboard.html b/headroom/dashboard/templates/dashboard.html index a1db508b6..7d44b366a 100644 --- a/headroom/dashboard/templates/dashboard.html +++ b/headroom/dashboard/templates/dashboard.html @@ -99,29 +99,34 @@