mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix: stabilize codex compression, stats, and proxy lifecycle
This commit is contained in:
parent
ac1d11c9a2
commit
eaf5980b4a
38 changed files with 3002 additions and 416 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -187,6 +187,7 @@ benchmark_results/
|
|||
.deepeval/
|
||||
|
||||
# Headroom specific
|
||||
.headroom/
|
||||
headroom.db
|
||||
headroom_*.db
|
||||
*.jsonl
|
||||
|
|
|
|||
|
|
@ -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<usize> = None;
|
||||
let mut latest_local_shell_output: Option<usize> = None;
|
||||
let mut latest_apply_patch_output: Option<usize> = None;
|
||||
let mut latest_message: Option<usize> = 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<usize>, &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<usize> = 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<Replacement> = 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<usize>,
|
||||
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"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
})]
|
||||
|
|
|
|||
|
|
@ -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<PyBytes>, bool, u64, Vec<String>) {
|
||||
) -> (Py<PyBytes>, bool, u64, Vec<String>, Option<String>) {
|
||||
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()),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
"""Package version metadata."""
|
||||
|
||||
__version__ = "0.5.25"
|
||||
__version__ = "0.9.1"
|
||||
|
|
|
|||
150
headroom/cache/compression_store.py
vendored
150
headroom/cache/compression_store.py
vendored
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 = "<!-- headroom:memory-instructions -->"
|
|||
# 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()
|
||||
|
|
|
|||
|
|
@ -99,29 +99,34 @@
|
|||
<main class="p-6 max-w-7xl mx-auto">
|
||||
<template x-if="viewMode === 'session'">
|
||||
<div>
|
||||
<!-- Hero Metrics (reordered: Savings $ -> Tokens Saved % -> Quality Confidence -> Overhead) -->
|
||||
<!-- Hero Metrics (reordered: Savings $ -> Compression % -> Quality Confidence -> Overhead) -->
|
||||
<div class="grid grid-cols-2 md:grid-cols-4 gap-4 mb-6">
|
||||
<!-- Savings ($) - Compression only, priced at model list rate -->
|
||||
<!-- Savings ($) - proxy compression only, priced at model list rate -->
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Compression Savings</div>
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Proxy $ Saved</div>
|
||||
<div class="flex items-baseline gap-2">
|
||||
<span class="text-3xl font-light tabular-nums text-emerald-400" x-text="'$' + formatCurrency(stats.cost?.savings_usd || 0)"></span>
|
||||
</div>
|
||||
<div class="mt-2 text-xs text-gray-500">
|
||||
<span x-show="stats.cost?.savings_usd > 0"
|
||||
x-text="formatNumber(stats.tokens?.saved || 0) + ' tokens at model list price'"></span>
|
||||
x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' proxy tokens only; RTK excluded from $'"></span>
|
||||
<span x-show="!(stats.cost?.savings_usd > 0)"
|
||||
x-text="formatNumber(stats.requests?.total || 0) + ' requests processed'"></span>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Tokens Saved (%) -->
|
||||
<!-- Token Savings (%) -->
|
||||
<div class="bg-surface rounded-lg p-4 border border-border">
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Tokens Saved</div>
|
||||
<div class="text-xs text-gray-500 uppercase tracking-wide mb-1">Token Savings</div>
|
||||
<div class="flex items-baseline gap-2">
|
||||
<span class="text-3xl font-light tabular-nums text-accent" x-text="formatNumber(stats.tokens?.saved || 0)"></span>
|
||||
<span class="text-sm text-accent" x-text="(stats.tokens?.savings_percent || 0).toFixed(1) + '%'"></span>
|
||||
</div>
|
||||
<div class="mt-1 text-xs text-gray-500 leading-relaxed">
|
||||
<span x-text="'Proxy ' + formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' (' + proxyShareOfTotal.toFixed(1) + '%)'"></span>
|
||||
<span class="mx-1 text-gray-600">/</span>
|
||||
<span x-text="'RTK ' + formatNumber(stats.tokens?.rtk_saved || 0) + ' (' + rtkShareOfTotal.toFixed(1) + '%)'"></span>
|
||||
</div>
|
||||
<div class="mt-2 h-8">
|
||||
<svg class="w-full h-full" viewBox="0 0 100 32" preserveAspectRatio="none">
|
||||
<defs>
|
||||
|
|
@ -166,7 +171,7 @@
|
|||
<div>
|
||||
<div class="text-xs text-gray-500 mb-1">Compression</div>
|
||||
<div class="text-2xl font-light tabular-nums text-emerald-400" x-text="'$' + formatCurrency(stats.cost?.compression_savings_usd || 0)"></div>
|
||||
<div class="text-xs text-gray-500" x-text="formatNumber(stats.tokens?.saved || 0) + ' tokens removed'"></div>
|
||||
<div class="text-xs text-gray-500" x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0) + ' proxy tokens removed'"></div>
|
||||
</div>
|
||||
<template x-if="(stats.cost?.cache_savings_usd || 0) > 0">
|
||||
<div>
|
||||
|
|
@ -692,6 +697,14 @@
|
|||
<span class="text-sm text-gray-400">Before Compression</span>
|
||||
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.total_before_compression || 0)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">RTK Filtered</span>
|
||||
<span class="font-mono text-sm text-emerald-400" x-text="formatNumber(stats.tokens?.rtk_saved || 0)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">Proxy Removed</span>
|
||||
<span class="font-mono text-sm text-accent" x-text="formatNumber(stats.tokens?.proxy_compression_saved || 0)"></span>
|
||||
</div>
|
||||
<div class="flex justify-between items-center">
|
||||
<span class="text-sm text-gray-400">After Compression (sent)</span>
|
||||
<span class="font-mono text-sm" x-text="formatNumber(stats.tokens?.input || 0)"></span>
|
||||
|
|
@ -1755,10 +1768,26 @@
|
|||
return Math.min((tokens / max) * 100, 100);
|
||||
},
|
||||
|
||||
get compressionTotalBefore() {
|
||||
return this.stats.tokens?.total_before_compression || 0;
|
||||
},
|
||||
|
||||
get proxyShareOfTotal() {
|
||||
const total = this.compressionTotalBefore;
|
||||
if (total <= 0) return 0;
|
||||
return (this.stats.tokens?.proxy_compression_saved || 0) / total * 100;
|
||||
},
|
||||
|
||||
get rtkShareOfTotal() {
|
||||
const total = this.compressionTotalBefore;
|
||||
if (total <= 0) return 0;
|
||||
return (this.stats.tokens?.rtk_saved || 0) / total * 100;
|
||||
},
|
||||
|
||||
// --- Compression Confidence ---
|
||||
|
||||
get confidenceLevel() {
|
||||
const saved = this.stats.tokens?.saved || 0;
|
||||
const saved = this.stats.tokens?.proxy_compression_saved || 0;
|
||||
if (saved === 0) return 'none';
|
||||
const signals = this.stats.waste_signals || {};
|
||||
const totalWaste = Object.values(signals).reduce((a, b) => a + b, 0);
|
||||
|
|
@ -1785,7 +1814,7 @@
|
|||
},
|
||||
|
||||
get confidenceDetail() {
|
||||
const saved = this.stats.tokens?.saved || 0;
|
||||
const saved = this.stats.tokens?.proxy_compression_saved || 0;
|
||||
if (saved === 0) return 'No compression yet';
|
||||
const signals = this.stats.waste_signals || {};
|
||||
const totalWaste = Object.values(signals).reduce((a, b) => a + b, 0);
|
||||
|
|
|
|||
|
|
@ -387,7 +387,9 @@ def build_session_summary(
|
|||
best_compression = best["savings_pct"]
|
||||
best_detail = f"{best['original']:,} → {best['optimized']:,} tokens"
|
||||
|
||||
# Cost summary — savings_usd is compression savings at model list price (monotonic)
|
||||
# Cost summary — dollar savings are proxy-compression only at model list
|
||||
# price. rtk tokens are counted in token savings but have no model-specific
|
||||
# price because they never reached the proxy request.
|
||||
cost_stats = proxy.cost_tracker.stats() if proxy.cost_tracker else {}
|
||||
cost_with = cost_stats.get("cost_with_headroom_usd", 0.0)
|
||||
compression_savings = cost_stats.get("savings_usd", 0.0)
|
||||
|
|
@ -412,6 +414,9 @@ def build_session_summary(
|
|||
"best_compression_pct": best_compression,
|
||||
"best_detail": best_detail,
|
||||
"total_tokens_removed": metrics.tokens_saved_total,
|
||||
"rtk_tokens_avoided": cli_tokens_avoided,
|
||||
"total_tokens_saved_with_rtk": metrics.tokens_saved_total + cli_tokens_avoided,
|
||||
"total_tokens_before_with_rtk": total_tokens_before,
|
||||
},
|
||||
"uncompressed_requests": {k: v for k, v in uncompressed_reasons.items() if v > 0},
|
||||
"cost": {
|
||||
|
|
@ -422,6 +427,11 @@ def build_session_summary(
|
|||
"breakdown": {
|
||||
"cache_savings_usd": round(cache_net, 2),
|
||||
"compression_savings_usd": round(compression_savings, 2),
|
||||
"rtk_savings_usd": None,
|
||||
"rtk_savings_note": (
|
||||
"rtk tokens are included in token savings only; dollar savings "
|
||||
"use proxy compression tokens at model list price."
|
||||
),
|
||||
},
|
||||
},
|
||||
}
|
||||
|
|
@ -470,6 +480,19 @@ class CostTracker:
|
|||
self._api_cache_write_1h_by_model: dict[str, int] = {}
|
||||
self._api_uncached_by_model: dict[str, int] = {}
|
||||
|
||||
def reset_runtime(self) -> None:
|
||||
"""Reset in-memory cost/token counters for local test/debug use."""
|
||||
self._costs.clear()
|
||||
self._last_prune_time = datetime.now()
|
||||
self._tokens_saved_by_model.clear()
|
||||
self._tokens_sent_by_model.clear()
|
||||
self._requests_by_model.clear()
|
||||
self._api_cache_read_by_model.clear()
|
||||
self._api_cache_write_by_model.clear()
|
||||
self._api_cache_write_5m_by_model.clear()
|
||||
self._api_cache_write_1h_by_model.clear()
|
||||
self._api_uncached_by_model.clear()
|
||||
|
||||
# Cache resolved model names to avoid repeated litellm lookups.
|
||||
# This is critical: litellm.cost_per_token() is synchronous and can block
|
||||
# the async event loop if it triggers I/O (lazy model info download).
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -230,11 +230,38 @@ class StreamingMixin:
|
|||
|
||||
elif provider == "openai":
|
||||
chunk_usage = data.get("usage")
|
||||
if chunk_usage:
|
||||
usage_found["input_tokens"] = chunk_usage.get("prompt_tokens", 0)
|
||||
usage_found["output_tokens"] = chunk_usage.get("completion_tokens", 0)
|
||||
details = chunk_usage.get("prompt_tokens_details") or {}
|
||||
usage_found["cache_read_input_tokens"] = details.get("cached_tokens", 0)
|
||||
if not isinstance(chunk_usage, dict):
|
||||
response = data.get("response")
|
||||
if isinstance(response, dict):
|
||||
chunk_usage = response.get("usage")
|
||||
if isinstance(chunk_usage, dict):
|
||||
|
||||
def _usage_int(value: Any) -> int:
|
||||
try:
|
||||
return max(int(value), 0)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
# Chat Completions streams report prompt/completion tokens.
|
||||
# Responses streams report input/output tokens under
|
||||
# response.usage on response.completed.
|
||||
input_tokens = chunk_usage.get("prompt_tokens")
|
||||
if input_tokens is None:
|
||||
input_tokens = chunk_usage.get("input_tokens", 0)
|
||||
output_tokens = chunk_usage.get("completion_tokens")
|
||||
if output_tokens is None:
|
||||
output_tokens = chunk_usage.get("output_tokens", 0)
|
||||
usage_found["input_tokens"] = _usage_int(input_tokens)
|
||||
usage_found["output_tokens"] = _usage_int(output_tokens)
|
||||
details = (
|
||||
chunk_usage.get("prompt_tokens_details")
|
||||
or chunk_usage.get("input_tokens_details")
|
||||
or {}
|
||||
)
|
||||
if isinstance(details, dict):
|
||||
usage_found["cache_read_input_tokens"] = _usage_int(
|
||||
details.get("cached_tokens")
|
||||
)
|
||||
|
||||
elif provider == "gemini":
|
||||
usage_meta = data.get("usageMetadata")
|
||||
|
|
@ -569,11 +596,24 @@ class StreamingMixin:
|
|||
f"estimating {output_tokens} from {stream_state['total_bytes']} bytes"
|
||||
)
|
||||
|
||||
provider_input_tokens = stream_state.get("input_tokens")
|
||||
effective_optimized_tokens = optimized_tokens
|
||||
effective_original_tokens = original_tokens
|
||||
if (
|
||||
provider == "openai"
|
||||
and isinstance(provider_input_tokens, int)
|
||||
and provider_input_tokens > 0
|
||||
):
|
||||
effective_optimized_tokens = provider_input_tokens
|
||||
effective_original_tokens = max(original_tokens, provider_input_tokens + tokens_saved)
|
||||
|
||||
cache_read_tokens = stream_state["cache_read_input_tokens"] or 0
|
||||
cache_write_tokens = stream_state["cache_creation_input_tokens"] or 0
|
||||
cache_write_5m_tokens = stream_state["cache_creation_ephemeral_5m_input_tokens"] or 0
|
||||
cache_write_1h_tokens = stream_state["cache_creation_ephemeral_1h_input_tokens"] or 0
|
||||
uncached_input_tokens = max(optimized_tokens - cache_read_tokens - cache_write_tokens, 0)
|
||||
uncached_input_tokens = max(
|
||||
effective_optimized_tokens - cache_read_tokens - cache_write_tokens, 0
|
||||
)
|
||||
|
||||
num_msgs = len(body.get("messages", []))
|
||||
cache_hit_pct = (
|
||||
|
|
@ -584,7 +624,7 @@ class StreamingMixin:
|
|||
logger.info(
|
||||
f"[{request_id}] PERF "
|
||||
f"model={model} msgs={num_msgs} "
|
||||
f"tok_before={original_tokens} tok_after={optimized_tokens} "
|
||||
f"tok_before={effective_original_tokens} tok_after={effective_optimized_tokens} "
|
||||
f"tok_saved={tokens_saved} "
|
||||
f"cache_read={cache_read_tokens} cache_write={cache_write_tokens} "
|
||||
f"cache_hit_pct={cache_hit_pct} "
|
||||
|
|
@ -622,7 +662,7 @@ class StreamingMixin:
|
|||
self.cost_tracker.record_tokens(
|
||||
model,
|
||||
tokens_saved,
|
||||
optimized_tokens,
|
||||
effective_optimized_tokens,
|
||||
cache_read_tokens=cache_read_tokens,
|
||||
cache_write_tokens=cache_write_tokens,
|
||||
cache_write_5m_tokens=cache_write_5m_tokens,
|
||||
|
|
@ -634,7 +674,7 @@ class StreamingMixin:
|
|||
await self.metrics.record_request(
|
||||
provider=provider,
|
||||
model=model,
|
||||
input_tokens=optimized_tokens,
|
||||
input_tokens=effective_optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
latency_ms=total_latency,
|
||||
|
|
@ -662,12 +702,12 @@ class StreamingMixin:
|
|||
timestamp=datetime.now().isoformat(),
|
||||
provider=provider,
|
||||
model=model,
|
||||
input_tokens_original=original_tokens,
|
||||
input_tokens_optimized=optimized_tokens,
|
||||
input_tokens_original=effective_original_tokens,
|
||||
input_tokens_optimized=effective_optimized_tokens,
|
||||
output_tokens=output_tokens,
|
||||
tokens_saved=tokens_saved,
|
||||
savings_percent=(tokens_saved / original_tokens * 100)
|
||||
if original_tokens > 0
|
||||
savings_percent=(tokens_saved / effective_original_tokens * 100)
|
||||
if effective_original_tokens > 0
|
||||
else 0,
|
||||
optimization_latency_ms=optimization_latency,
|
||||
total_latency_ms=total_latency,
|
||||
|
|
|
|||
|
|
@ -100,6 +100,32 @@ def _safe_event_name(event: str) -> str:
|
|||
return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in event)[:80]
|
||||
|
||||
|
||||
def _wire_debug_preview(value: Any, *, max_chars: int = 900) -> str:
|
||||
"""Return a compact, human-readable preview for proxy.log.
|
||||
|
||||
This is intentionally lossy: the full redacted payload is already written
|
||||
to the wire-debug JSON file. The log line should be short enough to scan
|
||||
live without flooding the proxy log.
|
||||
"""
|
||||
|
||||
try:
|
||||
if isinstance(value, bytes):
|
||||
text = safe_decode_for_logging(value, max_bytes=max_chars)
|
||||
elif isinstance(value, str):
|
||||
text = value
|
||||
elif value is None:
|
||||
return ""
|
||||
else:
|
||||
text = json.dumps(value, ensure_ascii=False, default=str, separators=(",", ":"))
|
||||
except Exception:
|
||||
text = repr(value)
|
||||
|
||||
text = " ".join(text.split())
|
||||
if len(text) > max_chars:
|
||||
return text[: max_chars - 1] + "…"
|
||||
return text
|
||||
|
||||
|
||||
def capture_codex_wire_debug(
|
||||
event: str,
|
||||
*,
|
||||
|
|
@ -159,6 +185,21 @@ def capture_codex_wire_debug(
|
|||
request_id or "",
|
||||
event,
|
||||
)
|
||||
preview_source = redact_for_wire_debug(body) if body is not None else raw_text
|
||||
preview = _wire_debug_preview(preview_source)
|
||||
meta_keys = ",".join(sorted((metadata or {}).keys()))
|
||||
logger.info(
|
||||
"event=codex_wire_debug_frame request_id=%s session_id=%s wire_event=%s "
|
||||
"transport=%s direction=%s status_code=%s meta_keys=%s preview=%s",
|
||||
request_id or "",
|
||||
session_id or "",
|
||||
event,
|
||||
transport,
|
||||
direction,
|
||||
status_code if status_code is not None else "",
|
||||
meta_keys,
|
||||
preview,
|
||||
)
|
||||
return path
|
||||
except Exception as exc: # pragma: no cover - debug path must never break traffic
|
||||
logger.warning("event=codex_wire_debug_capture_failed error=%s", exc)
|
||||
|
|
@ -250,6 +291,21 @@ def get_python_forwarder_mode() -> PythonForwarderMode:
|
|||
)
|
||||
|
||||
|
||||
def _headroom_bypass_enabled(headers: Any) -> bool:
|
||||
"""Return True when inbound headers request full Headroom passthrough.
|
||||
|
||||
This is transport-neutral policy: HTTP and WebSocket handlers both call
|
||||
it on original inbound headers before request-body mutation.
|
||||
"""
|
||||
|
||||
try:
|
||||
bypass = str(headers.get("x-headroom-bypass", "")).strip().lower() == "true"
|
||||
passthrough = str(headers.get("x-headroom-mode", "")).strip().lower() == "passthrough"
|
||||
except AttributeError:
|
||||
return False
|
||||
return bypass or passthrough
|
||||
|
||||
|
||||
def serialize_body_canonical(body: dict[str, Any]) -> bytes:
|
||||
"""Re-serialize a request body deterministically with cache-stable formatting.
|
||||
|
||||
|
|
@ -517,6 +573,11 @@ _rtk_stats_cache: dict[str, Any] = {
|
|||
"has_value": False,
|
||||
"value": None,
|
||||
}
|
||||
_rtk_session_baseline: dict[str, Any] = {
|
||||
"initialized": False,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
}
|
||||
|
||||
# Maximum request body size (100MB - increased to support image-heavy requests)
|
||||
MAX_REQUEST_BODY_SIZE = 100 * 1024 * 1024
|
||||
|
|
@ -742,39 +803,20 @@ def _setup_file_logging() -> None:
|
|||
pass
|
||||
|
||||
|
||||
def _get_rtk_stats() -> dict[str, Any] | None:
|
||||
"""Get rtk (Rust Token Killer) savings stats if rtk is installed.
|
||||
def _read_rtk_lifetime_stats() -> dict[str, Any] | None:
|
||||
"""Read rtk's current project-level lifetime stats."""
|
||||
|
||||
Reads from rtk's tracking database via `rtk gain --format json`.
|
||||
Results are memoized briefly so dashboard polling does not spawn a new
|
||||
subprocess on every refresh.
|
||||
"""
|
||||
import subprocess as _sp
|
||||
|
||||
from headroom.rtk import get_rtk_path
|
||||
|
||||
now = time.monotonic()
|
||||
with _rtk_stats_cache_lock:
|
||||
if _rtk_stats_cache["has_value"] and now < float(_rtk_stats_cache["expires_at"]):
|
||||
return cast(dict[str, Any] | None, _rtk_stats_cache["value"])
|
||||
|
||||
payload: dict[str, Any] | None
|
||||
rtk_path = get_rtk_path()
|
||||
if not rtk_path:
|
||||
payload = None
|
||||
with _rtk_stats_cache_lock:
|
||||
_rtk_stats_cache.update(
|
||||
{
|
||||
"expires_at": time.monotonic() + RTK_STATS_CACHE_TTL_SECONDS,
|
||||
"has_value": True,
|
||||
"value": payload,
|
||||
}
|
||||
)
|
||||
return payload
|
||||
return None
|
||||
|
||||
try:
|
||||
result = _sp.run(
|
||||
[str(rtk_path), "gain", "--format", "json"],
|
||||
[str(rtk_path), "gain", "--project", "--format", "json"],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=5,
|
||||
|
|
@ -789,21 +831,83 @@ def _get_rtk_stats() -> dict[str, Any] | None:
|
|||
"avg_savings_pct": summary.get("avg_savings_pct", 0.0),
|
||||
}
|
||||
else:
|
||||
payload = {
|
||||
return {
|
||||
"installed": True,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
"avg_savings_pct": 0.0,
|
||||
}
|
||||
except Exception:
|
||||
payload = {
|
||||
return {
|
||||
"installed": True,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
"avg_savings_pct": 0.0,
|
||||
}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def initialize_rtk_session_baseline() -> None:
|
||||
"""Pin the current rtk counters as the proxy-session baseline."""
|
||||
|
||||
payload = _read_rtk_lifetime_stats()
|
||||
with _rtk_stats_cache_lock:
|
||||
_rtk_session_baseline.update(
|
||||
{
|
||||
"initialized": True,
|
||||
"total_commands": int((payload or {}).get("total_commands", 0) or 0),
|
||||
"tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0),
|
||||
}
|
||||
)
|
||||
_rtk_stats_cache.update(
|
||||
{
|
||||
"expires_at": 0.0,
|
||||
"has_value": False,
|
||||
"value": None,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _get_rtk_stats() -> dict[str, Any] | None:
|
||||
"""Get rtk savings for the current Headroom proxy session.
|
||||
|
||||
rtk persists project-level lifetime counters. Dashboard stats should be
|
||||
session-local, so we subtract the counter snapshot captured at proxy
|
||||
startup instead of resetting rtk's own history.
|
||||
"""
|
||||
|
||||
now = time.monotonic()
|
||||
with _rtk_stats_cache_lock:
|
||||
if _rtk_stats_cache["has_value"] and now < float(_rtk_stats_cache["expires_at"]):
|
||||
return cast(dict[str, Any] | None, _rtk_stats_cache["value"])
|
||||
|
||||
payload = _read_rtk_lifetime_stats()
|
||||
with _rtk_stats_cache_lock:
|
||||
if not _rtk_session_baseline["initialized"]:
|
||||
_rtk_session_baseline.update(
|
||||
{
|
||||
"initialized": True,
|
||||
"total_commands": int((payload or {}).get("total_commands", 0) or 0),
|
||||
"tokens_saved": int((payload or {}).get("tokens_saved", 0) or 0),
|
||||
}
|
||||
)
|
||||
|
||||
if payload is not None:
|
||||
payload = {
|
||||
**payload,
|
||||
"total_commands": max(
|
||||
int(payload.get("total_commands", 0) or 0)
|
||||
- int(_rtk_session_baseline["total_commands"]),
|
||||
0,
|
||||
),
|
||||
"tokens_saved": max(
|
||||
int(payload.get("tokens_saved", 0) or 0)
|
||||
- int(_rtk_session_baseline["tokens_saved"]),
|
||||
0,
|
||||
),
|
||||
}
|
||||
|
||||
_rtk_stats_cache.update(
|
||||
{
|
||||
"expires_at": time.monotonic() + RTK_STATS_CACHE_TTL_SECONDS,
|
||||
|
|
|
|||
|
|
@ -193,6 +193,62 @@ class PrometheusMetrics:
|
|||
self._stage_timing_lock = threading.Lock()
|
||||
self._otel_metrics = otel_metrics
|
||||
|
||||
async def reset_runtime(self) -> None:
|
||||
"""Reset in-memory request/compression counters for local test/debug use."""
|
||||
async with self._lock:
|
||||
self.requests_total = 0
|
||||
self.requests_by_provider.clear()
|
||||
self.requests_by_model.clear()
|
||||
self.requests_by_stack.clear()
|
||||
self.requests_cached = 0
|
||||
self.requests_rate_limited = 0
|
||||
self.requests_failed = 0
|
||||
|
||||
self.tokens_input_total = 0
|
||||
self.tokens_output_total = 0
|
||||
self.tokens_saved_total = 0
|
||||
|
||||
self.compressions_by_strategy.clear()
|
||||
self.tokens_saved_by_strategy.clear()
|
||||
|
||||
self.latency_sum_ms = 0.0
|
||||
self.latency_min_ms = float("inf")
|
||||
self.latency_max_ms = 0.0
|
||||
self.latency_count = 0
|
||||
|
||||
self.overhead_sum_ms = 0.0
|
||||
self.overhead_min_ms = float("inf")
|
||||
self.overhead_max_ms = 0.0
|
||||
self.overhead_count = 0
|
||||
|
||||
self.ttfb_sum_ms = 0.0
|
||||
self.ttfb_min_ms = float("inf")
|
||||
self.ttfb_max_ms = 0.0
|
||||
self.ttfb_count = 0
|
||||
|
||||
self.transform_timing_sum.clear()
|
||||
self.transform_timing_count.clear()
|
||||
self.transform_timing_max.clear()
|
||||
|
||||
self.waste_signals_total.clear()
|
||||
self.cache_by_provider.clear()
|
||||
self._cache_requests_by_model.clear()
|
||||
|
||||
self.prefix_freeze_busts_avoided = 0
|
||||
self.prefix_freeze_tokens_preserved = 0
|
||||
self.prefix_freeze_compression_foregone = 0
|
||||
self.cache_bust_tokens_lost = 0
|
||||
self.cache_bust_count = 0
|
||||
self.savings_history = []
|
||||
|
||||
with self._stage_timing_lock:
|
||||
self.stage_timing_sum.clear()
|
||||
self.stage_timing_count.clear()
|
||||
self.stage_timing_max.clear()
|
||||
self.ws_session_duration_sum_ms.clear()
|
||||
self.ws_session_duration_count.clear()
|
||||
self.ws_session_duration_max_ms.clear()
|
||||
|
||||
def _get_otel_metrics(self) -> HeadroomOtelMetrics:
|
||||
return self._otel_metrics or get_otel_metrics()
|
||||
|
||||
|
|
|
|||
|
|
@ -74,9 +74,7 @@ from headroom.ccr import (
|
|||
)
|
||||
from headroom.config import (
|
||||
CacheAlignerConfig,
|
||||
CCRConfig,
|
||||
ReadLifecycleConfig,
|
||||
SmartCrusherConfig,
|
||||
)
|
||||
from headroom.dashboard import get_dashboard_html
|
||||
from headroom.observability import (
|
||||
|
|
@ -122,6 +120,7 @@ from headroom.proxy.helpers import (
|
|||
_get_rtk_stats, # noqa: F401
|
||||
_read_request_json, # noqa: F401
|
||||
_setup_file_logging, # noqa: F401
|
||||
initialize_rtk_session_baseline,
|
||||
is_anthropic_auth, # noqa: F401
|
||||
jitter_delay_ms,
|
||||
)
|
||||
|
|
@ -157,7 +156,6 @@ from headroom.transforms import (
|
|||
CodeCompressorConfig,
|
||||
ContentRouter,
|
||||
ContentRouterConfig,
|
||||
SmartCrusher,
|
||||
TransformPipeline,
|
||||
is_tree_sitter_available,
|
||||
)
|
||||
|
|
@ -346,41 +344,22 @@ class HeadroomProxy(
|
|||
# Reported via metrics as `_context_manager_status = "passthrough"`.
|
||||
self._context_manager_status = "passthrough"
|
||||
|
||||
if config.smart_routing:
|
||||
# Smart routing: ContentRouter handles all content types intelligently
|
||||
# It lazy-loads compressors only when needed
|
||||
router_config = ContentRouterConfig(
|
||||
enable_code_aware=config.code_aware_enabled,
|
||||
tool_profiles=config.tool_profiles,
|
||||
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
|
||||
)
|
||||
# Token mode: allow compression of older excluded-tool results
|
||||
if is_token_mode(config.mode):
|
||||
router_config.protect_recent_reads_fraction = 0.3
|
||||
transforms = [
|
||||
CacheAligner(CacheAlignerConfig(enabled=False)),
|
||||
ContentRouter(router_config, observer=self.metrics),
|
||||
]
|
||||
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
|
||||
else:
|
||||
# Legacy mode: sequential pipeline
|
||||
transforms = [
|
||||
CacheAligner(CacheAlignerConfig(enabled=False)),
|
||||
SmartCrusher(
|
||||
SmartCrusherConfig( # type: ignore[arg-type]
|
||||
enabled=True,
|
||||
min_tokens_to_crush=config.min_tokens_to_crush,
|
||||
max_items_after_crush=config.max_items_after_crush,
|
||||
),
|
||||
ccr_config=CCRConfig(
|
||||
enabled=config.ccr_inject_tool,
|
||||
inject_retrieval_marker=config.ccr_inject_tool, # Add CCR markers
|
||||
),
|
||||
observer=self.metrics,
|
||||
),
|
||||
]
|
||||
# Add CodeAware if enabled and available
|
||||
self._code_aware_status = self._setup_code_aware(config, transforms)
|
||||
# ContentRouter is the single proxy routing surface. Provider handlers
|
||||
# normalize their request shapes into messages or CompressionUnits, and
|
||||
# the router chooses SmartCrusher, log/search/diff/code, or Kompress.
|
||||
router_config = ContentRouterConfig(
|
||||
enable_code_aware=config.code_aware_enabled,
|
||||
tool_profiles=config.tool_profiles,
|
||||
read_lifecycle=ReadLifecycleConfig(enabled=config.read_lifecycle),
|
||||
)
|
||||
# Token mode: allow compression of older excluded-tool results.
|
||||
if is_token_mode(config.mode):
|
||||
router_config.protect_recent_reads_fraction = 0.3
|
||||
transforms = [
|
||||
CacheAligner(CacheAlignerConfig(enabled=False)),
|
||||
ContentRouter(router_config, observer=self.metrics),
|
||||
]
|
||||
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
|
||||
|
||||
self.anthropic_pipeline = TransformPipeline(
|
||||
transforms=transforms,
|
||||
|
|
@ -862,11 +841,7 @@ class HeadroomProxy(
|
|||
self.anthropic_pre_upstream_memory_context_timeout_seconds,
|
||||
)
|
||||
|
||||
# Smart routing status
|
||||
if self.config.smart_routing:
|
||||
logger.info("Smart Routing: ENABLED (intelligent content detection)")
|
||||
else:
|
||||
logger.info("Smart Routing: DISABLED (legacy sequential mode)")
|
||||
logger.info("Smart Routing: ENABLED (ContentRouter is always active)")
|
||||
|
||||
# Eagerly load ALL compressors, parsers, and detectors at startup
|
||||
# This eliminates cold-start latency spikes on first requests.
|
||||
|
|
@ -1364,6 +1339,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
app.state.started_at = time.time()
|
||||
app.state.ready = False
|
||||
app.state.startup_error = None
|
||||
initialize_rtk_session_baseline()
|
||||
|
||||
try:
|
||||
try:
|
||||
|
|
@ -1734,8 +1710,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
# compression and rtk both remove tokens before they reach model
|
||||
# context, so dashboard-facing compression savings combines them.
|
||||
proxy_compression_tokens = m.tokens_saved_total
|
||||
compression_tokens = proxy_compression_tokens + cli_tokens_avoided
|
||||
total_tokens_before = m.tokens_input_total + compression_tokens
|
||||
all_layers_tokens_saved = proxy_compression_tokens + cli_tokens_avoided
|
||||
total_tokens_before = m.tokens_input_total + all_layers_tokens_saved
|
||||
proxy_total_before_compression = m.tokens_input_total + proxy_compression_tokens
|
||||
|
||||
# Build human-readable summary
|
||||
summary = _build_session_summary(
|
||||
|
|
@ -1780,7 +1757,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
|
||||
# Build unified savings summary (all layers)
|
||||
cache_net_usd = prefix_cache_stats.get("totals", {}).get("net_savings_usd", 0.0)
|
||||
total_tokens_all_layers = compression_tokens
|
||||
total_tokens_all_layers = all_layers_tokens_saved
|
||||
persistent_savings = m.savings_tracker.stats_preview()
|
||||
display_session = persistent_savings.get("display_session", {})
|
||||
|
||||
|
|
@ -1791,18 +1768,20 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"by_layer": {
|
||||
"cli_filtering": {
|
||||
"tokens": cli_tokens_avoided,
|
||||
"included_in": "compression",
|
||||
"included_in": "tokens.saved",
|
||||
"description": (
|
||||
"Tokens avoided by CLI output filtering (rtk) before reaching context. "
|
||||
"Included in dashboard compression savings."
|
||||
"Included in dashboard token savings, but not in dollar savings."
|
||||
),
|
||||
},
|
||||
"compression": {
|
||||
"tokens": compression_tokens,
|
||||
"tokens": proxy_compression_tokens,
|
||||
"proxy_tokens": proxy_compression_tokens,
|
||||
"rtk_tokens": cli_tokens_avoided,
|
||||
"all_layers_tokens": all_layers_tokens_saved,
|
||||
"description": (
|
||||
"Tokens removed before model context by proxy compression plus rtk CLI filtering."
|
||||
"Tokens removed by Headroom proxy compression. "
|
||||
"Dashboard token savings also includes rtk CLI filtering."
|
||||
),
|
||||
},
|
||||
"prefix_cache": {
|
||||
|
|
@ -1827,13 +1806,27 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
"tokens": {
|
||||
"input": m.tokens_input_total,
|
||||
"output": m.tokens_output_total,
|
||||
"saved": compression_tokens,
|
||||
"saved": all_layers_tokens_saved,
|
||||
"proxy_compression_saved": proxy_compression_tokens,
|
||||
"rtk_saved": cli_tokens_avoided,
|
||||
"cli_tokens_avoided": cli_tokens_avoided,
|
||||
"proxy_total_before_compression": proxy_total_before_compression,
|
||||
"total_before_compression": total_tokens_before,
|
||||
"all_layers_saved": all_layers_tokens_saved,
|
||||
"proxy_savings_percent": round(
|
||||
(proxy_compression_tokens / proxy_total_before_compression * 100)
|
||||
if proxy_total_before_compression > 0
|
||||
else 0,
|
||||
2,
|
||||
),
|
||||
"savings_percent": round(
|
||||
(compression_tokens / total_tokens_before * 100)
|
||||
(all_layers_tokens_saved / total_tokens_before * 100)
|
||||
if total_tokens_before > 0
|
||||
else 0,
|
||||
2,
|
||||
),
|
||||
"all_layers_savings_percent": round(
|
||||
(all_layers_tokens_saved / total_tokens_before * 100)
|
||||
if total_tokens_before > 0
|
||||
else 0,
|
||||
2,
|
||||
|
|
@ -1957,6 +1950,18 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
return await _get_cached_stats_payload()
|
||||
return await _build_stats_payload()
|
||||
|
||||
@app.post("/stats/reset", dependencies=[Depends(_require_loopback)])
|
||||
async def stats_reset():
|
||||
"""Reset in-memory proxy stats for local test/debug isolation."""
|
||||
await proxy.metrics.reset_runtime()
|
||||
if proxy.cost_tracker:
|
||||
proxy.cost_tracker.reset_runtime()
|
||||
initialize_rtk_session_baseline()
|
||||
async with _stats_snapshot_lock:
|
||||
_stats_snapshot["value"] = None
|
||||
_stats_snapshot["expires_at"] = 0.0
|
||||
return JSONResponse(status_code=200, content={"status": "reset"})
|
||||
|
||||
@app.get("/stats-history")
|
||||
async def stats_history(
|
||||
format: Literal["json", "csv"] = "json",
|
||||
|
|
@ -2105,6 +2110,10 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
store = get_compression_store()
|
||||
|
||||
if query:
|
||||
if not store.exists(hash_key, clean_expired=True):
|
||||
raise HTTPException(
|
||||
status_code=404, detail="Entry not found or expired (TTL: 5 minutes)"
|
||||
)
|
||||
# Search within cached content
|
||||
results = store.search(hash_key, query)
|
||||
return {
|
||||
|
|
@ -2415,6 +2424,8 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
store = get_compression_store()
|
||||
|
||||
if query:
|
||||
if not store.exists(hash_key, clean_expired=True):
|
||||
raise HTTPException(status_code=404, detail="Entry not found or expired")
|
||||
results = store.search(hash_key, query)
|
||||
return {
|
||||
"hash": hash_key,
|
||||
|
|
@ -2911,13 +2922,6 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--log-file", help="Log file path")
|
||||
parser.add_argument("--log-messages", action="store_true", help="Log full messages")
|
||||
|
||||
# Smart routing (content-aware compression)
|
||||
parser.add_argument(
|
||||
"--no-smart-routing",
|
||||
action="store_true",
|
||||
help="Disable smart routing (use legacy sequential pipeline)",
|
||||
)
|
||||
|
||||
# Code-aware compression
|
||||
parser.add_argument(
|
||||
"--code-aware",
|
||||
|
|
@ -2934,7 +2938,6 @@ if __name__ == "__main__":
|
|||
|
||||
# Environment variable defaults (HEADROOM_* prefix)
|
||||
# CLI args override env vars, env vars override ProxyConfig defaults
|
||||
env_smart_routing = _get_env_bool("HEADROOM_SMART_ROUTING", True)
|
||||
env_code_aware = _get_env_bool("HEADROOM_CODE_AWARE_ENABLED", True)
|
||||
env_optimize = _get_env_bool("HEADROOM_OPTIMIZE", True)
|
||||
env_cache = _get_env_bool("HEADROOM_CACHE_ENABLED", True)
|
||||
|
|
@ -2942,7 +2945,6 @@ if __name__ == "__main__":
|
|||
|
||||
# Determine settings: CLI flags override env vars
|
||||
# --no-X explicitly disables, --X explicitly enables, neither uses env var
|
||||
smart_routing = env_smart_routing if not args.no_smart_routing else False
|
||||
code_aware_enabled = (
|
||||
env_code_aware
|
||||
if not (args.code_aware or args.no_code_aware)
|
||||
|
|
@ -2983,7 +2985,7 @@ if __name__ == "__main__":
|
|||
if args.log_file
|
||||
else os.environ.get("HEADROOM_LOG_FILE"),
|
||||
log_full_messages=args.log_messages or _get_env_bool("HEADROOM_LOG_MESSAGES", False),
|
||||
smart_routing=smart_routing,
|
||||
smart_routing=True,
|
||||
code_aware_enabled=code_aware_enabled,
|
||||
# Connection pool settings
|
||||
max_connections=_get_env_int("HEADROOM_MAX_CONNECTIONS", args.max_connections),
|
||||
|
|
|
|||
176
headroom/transforms/compression_units.py
Normal file
176
headroom/transforms/compression_units.py
Normal file
|
|
@ -0,0 +1,176 @@
|
|||
"""Provider-neutral compression units.
|
||||
|
||||
Provider adapters own request-envelope details and cache/live-zone decisions.
|
||||
They should extract only safe, mutable text ranges into ``CompressionUnit``
|
||||
objects, ask ContentRouter to compress each unit, then splice accepted
|
||||
replacements back into their native request shape.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
from .content_router import CompressionStrategy, ContentRouter, RouterCompressionResult
|
||||
|
||||
|
||||
class TokenCounterLike(Protocol):
|
||||
def count_text(self, text: str) -> int: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CompressionUnit:
|
||||
"""One provider-extracted, cache-safe text slot."""
|
||||
|
||||
text: str
|
||||
provider: str
|
||||
endpoint: str
|
||||
role: str
|
||||
item_type: str
|
||||
cache_zone: str = "live"
|
||||
mutable: bool = True
|
||||
context: str = ""
|
||||
question: str | None = None
|
||||
bias: float = 1.0
|
||||
min_bytes: int = 512
|
||||
metadata: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UnitCompressionResult:
|
||||
original: str
|
||||
compressed: str
|
||||
modified: bool
|
||||
tokens_before: int
|
||||
tokens_after: int
|
||||
tokens_saved: int
|
||||
transforms_applied: list[str]
|
||||
strategy: str
|
||||
reason: str | None = None
|
||||
router_result: RouterCompressionResult | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RoutedCompressionUnit:
|
||||
"""A unit paired with its provider-owned slot reference."""
|
||||
|
||||
unit: CompressionUnit
|
||||
slot: object
|
||||
|
||||
|
||||
def find_content_router(transforms: object) -> ContentRouter | None:
|
||||
"""Return the first ContentRouter in a pipeline or iterable."""
|
||||
|
||||
candidates = getattr(transforms, "transforms", transforms)
|
||||
try:
|
||||
iterator = iter(candidates) # type: ignore[arg-type]
|
||||
except TypeError:
|
||||
return None
|
||||
for transform in iterator:
|
||||
if isinstance(transform, ContentRouter):
|
||||
return transform
|
||||
return None
|
||||
|
||||
|
||||
def compress_unit_with_router(
|
||||
unit: CompressionUnit,
|
||||
*,
|
||||
router: ContentRouter,
|
||||
tokenizer: TokenCounterLike,
|
||||
) -> UnitCompressionResult:
|
||||
"""Compress one safe text unit through ContentRouter.
|
||||
|
||||
The final accept/reject gate uses the provider/model tokenizer, not the
|
||||
router's internal word-count estimates.
|
||||
"""
|
||||
|
||||
tokens_before = tokenizer.count_text(unit.text)
|
||||
base = {
|
||||
"original": unit.text,
|
||||
"compressed": unit.text,
|
||||
"modified": False,
|
||||
"tokens_before": tokens_before,
|
||||
"tokens_after": tokens_before,
|
||||
"tokens_saved": 0,
|
||||
"transforms_applied": [],
|
||||
"strategy": CompressionStrategy.PASSTHROUGH.value,
|
||||
"router_result": None,
|
||||
}
|
||||
|
||||
if not unit.mutable:
|
||||
return UnitCompressionResult(**base, reason="immutable")
|
||||
if unit.role == "user":
|
||||
return UnitCompressionResult(**base, reason="protected_user_message")
|
||||
if unit.role in {"system", "developer"}:
|
||||
return UnitCompressionResult(**base, reason="protected_system_message")
|
||||
if unit.role == "assistant" and unit.metadata.get("compress_assistant") != "true":
|
||||
return UnitCompressionResult(**base, reason="protected_assistant_message")
|
||||
if unit.cache_zone != "live":
|
||||
return UnitCompressionResult(**base, reason=f"cache_zone_{unit.cache_zone}")
|
||||
if len(unit.text) < unit.min_bytes:
|
||||
return UnitCompressionResult(**base, reason="below_unit_floor")
|
||||
if "Retrieve more: hash=" in unit.text or "Retrieve original: hash=" in unit.text:
|
||||
return UnitCompressionResult(**base, reason="already_compressed")
|
||||
|
||||
router_result = router.compress(
|
||||
unit.text,
|
||||
context=unit.context,
|
||||
question=unit.question,
|
||||
bias=unit.bias,
|
||||
)
|
||||
replacement = router_result.compressed
|
||||
strategy = router_result.strategy_used.value
|
||||
if replacement == unit.text:
|
||||
return UnitCompressionResult(
|
||||
**{**base, "strategy": strategy, "router_result": router_result},
|
||||
reason="router_no_change",
|
||||
)
|
||||
|
||||
tokens_after = tokenizer.count_text(replacement)
|
||||
if tokens_after >= tokens_before:
|
||||
return UnitCompressionResult(
|
||||
**{
|
||||
**base,
|
||||
"compressed": replacement,
|
||||
"tokens_after": tokens_after,
|
||||
"strategy": strategy,
|
||||
"router_result": router_result,
|
||||
},
|
||||
reason="rejected_not_smaller",
|
||||
)
|
||||
|
||||
return UnitCompressionResult(
|
||||
original=unit.text,
|
||||
compressed=replacement,
|
||||
modified=True,
|
||||
tokens_before=tokens_before,
|
||||
tokens_after=tokens_after,
|
||||
tokens_saved=tokens_before - tokens_after,
|
||||
transforms_applied=[
|
||||
f"router:{unit.provider}:{unit.endpoint}:{unit.item_type}:{strategy}",
|
||||
strategy,
|
||||
],
|
||||
strategy=strategy,
|
||||
reason=None,
|
||||
router_result=router_result,
|
||||
)
|
||||
|
||||
|
||||
def compress_units_with_router(
|
||||
units: Iterable[RoutedCompressionUnit],
|
||||
*,
|
||||
router: ContentRouter,
|
||||
tokenizer: TokenCounterLike,
|
||||
) -> list[tuple[object, UnitCompressionResult]]:
|
||||
"""Compress provider-extracted units and preserve provider slot refs.
|
||||
|
||||
Provider adapters use this when they have many candidate text slots in one
|
||||
request envelope. The slot object is intentionally opaque here; only the
|
||||
provider adapter knows how to splice the result back into its native shape.
|
||||
"""
|
||||
|
||||
return [
|
||||
(routed.slot, compress_unit_with_router(routed.unit, router=router, tokenizer=tokenizer))
|
||||
for routed in units
|
||||
]
|
||||
|
|
@ -1520,7 +1520,7 @@ class ContentRouter(Transform):
|
|||
skip_user = (
|
||||
kwargs.get("compress_user_messages") is not True and self.config.skip_user_messages
|
||||
)
|
||||
skip_system = kwargs.get("compress_system_messages") is False
|
||||
skip_system = kwargs.get("compress_system_messages") is not True
|
||||
protect_recent = kwargs.get("protect_recent", self.config.protect_recent_code)
|
||||
protect_analysis = kwargs.get(
|
||||
"protect_analysis_context", self.config.protect_analysis_context
|
||||
|
|
@ -1715,10 +1715,11 @@ class ContentRouter(Transform):
|
|||
route_counts["user_msg"] += 1
|
||||
continue
|
||||
|
||||
# Protection 1b: Never compress system messages (when disabled)
|
||||
if skip_system and role == "system":
|
||||
# Protection 1b: Never compress system/developer messages unless
|
||||
# explicitly opted in. These are cache-hot instruction bytes.
|
||||
if skip_system and role in {"system", "developer"}:
|
||||
result_slots[i] = message
|
||||
transforms_applied.append("router:protected:system_message")
|
||||
transforms_applied.append(f"router:protected:{role}_message")
|
||||
route_counts.setdefault("system_msg", 0)
|
||||
route_counts["system_msg"] += 1
|
||||
continue
|
||||
|
|
@ -2004,7 +2005,7 @@ class ContentRouter(Transform):
|
|||
# Role-based gate for `text` blocks. Tool/function roles are tool
|
||||
# outputs and compress freely; assistant defaults to skip (cache
|
||||
# safety) with explicit opt-in; unknown roles default to skip.
|
||||
if (skip_user and role == "user") or (skip_system and role == "system"):
|
||||
if (skip_user and role == "user") or (skip_system and role in {"system", "developer"}):
|
||||
protect_text_blocks = True
|
||||
elif role == "assistant" and not compress_assistant_text_blocks:
|
||||
protect_text_blocks = True
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ from ..utils import deep_copy_messages
|
|||
from .base import Transform
|
||||
from .cache_aligner import CacheAligner
|
||||
from .content_router import ContentRouter
|
||||
from .smart_crusher import SmartCrusher
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..providers.base import Provider
|
||||
|
|
@ -37,7 +36,6 @@ class TransformPipeline:
|
|||
1. Cache Aligner - normalize prefix for cache hits
|
||||
2. Content Router - intelligent content-aware compression (routes to appropriate
|
||||
compressor: Kompress for text, SmartCrusher for JSON, CodeCompressor for code, etc.)
|
||||
3. SmartCrusher - fallback if ContentRouter disabled
|
||||
|
||||
Phase B PR-B1 retired the IntelligentContextManager / RollingWindow
|
||||
"drop messages from history" stage. Live-zone-only compression is the
|
||||
|
|
@ -100,26 +98,13 @@ class TransformPipeline:
|
|||
# - Logs -> LogCompressor
|
||||
# - Search results -> SearchCompressor
|
||||
# - HTML -> HTMLExtractor
|
||||
if self.config.content_router_enabled:
|
||||
transforms.append(ContentRouter())
|
||||
logger.info("Pipeline using ContentRouter for intelligent content-aware compression")
|
||||
elif self.config.smart_crusher.enabled:
|
||||
# Fallback: SmartCrusher only handles JSON arrays
|
||||
from .smart_crusher import SmartCrusherConfig as SCConfig
|
||||
|
||||
smart_config = SCConfig(
|
||||
enabled=True,
|
||||
min_items_to_analyze=self.config.smart_crusher.min_items_to_analyze,
|
||||
min_tokens_to_crush=self.config.smart_crusher.min_tokens_to_crush,
|
||||
variance_threshold=self.config.smart_crusher.variance_threshold,
|
||||
uniqueness_threshold=self.config.smart_crusher.uniqueness_threshold,
|
||||
similarity_threshold=self.config.smart_crusher.similarity_threshold,
|
||||
max_items_after_crush=self.config.smart_crusher.max_items_after_crush,
|
||||
preserve_change_points=self.config.smart_crusher.preserve_change_points,
|
||||
factor_out_constants=self.config.smart_crusher.factor_out_constants,
|
||||
include_summaries=self.config.smart_crusher.include_summaries,
|
||||
if not self.config.content_router_enabled:
|
||||
logger.warning(
|
||||
"HeadroomConfig.content_router_enabled=False is deprecated and ignored; "
|
||||
"ContentRouter is always present in the default pipeline"
|
||||
)
|
||||
transforms.append(SmartCrusher(smart_config))
|
||||
transforms.append(ContentRouter())
|
||||
logger.info("Pipeline using ContentRouter for intelligent content-aware compression")
|
||||
|
||||
return transforms
|
||||
|
||||
|
|
|
|||
|
|
@ -15,6 +15,8 @@ from headroom.pipeline import (
|
|||
summarize_routing_markers,
|
||||
)
|
||||
from headroom.providers.base import Provider, TokenCounter
|
||||
from headroom.transforms import ContentRouter, TransformPipeline
|
||||
from headroom.transforms.content_router import CompressionStrategy
|
||||
|
||||
|
||||
class RecordingExtension:
|
||||
|
|
@ -146,6 +148,49 @@ def test_pipeline_extension_manager_uses_canonical_stage_contract():
|
|||
assert event.messages == [{"role": "user", "content": "mutated"}]
|
||||
|
||||
|
||||
def test_default_transform_pipeline_always_uses_content_router() -> None:
|
||||
config = HeadroomConfig(content_router_enabled=False)
|
||||
|
||||
pipeline = TransformPipeline(config)
|
||||
|
||||
assert any(isinstance(transform, ContentRouter) for transform in pipeline.transforms)
|
||||
assert not any(type(transform).__name__ == "SmartCrusher" for transform in pipeline.transforms)
|
||||
|
||||
|
||||
def test_content_router_protects_instruction_roles_but_compresses_tool_outputs() -> None:
|
||||
class Tokenizer:
|
||||
def count_text(self, text: str) -> int:
|
||||
return max(1, len(text.split()))
|
||||
|
||||
router = ContentRouter()
|
||||
calls: list[str] = []
|
||||
|
||||
def fake_compress(text: str, **kwargs: Any) -> SimpleNamespace:
|
||||
calls.append(text)
|
||||
return SimpleNamespace(
|
||||
compressed="COMPRESSED",
|
||||
compression_ratio=0.1,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
router.compress = fake_compress # type: ignore[method-assign]
|
||||
tool_text = "tool output " * 120
|
||||
messages = [
|
||||
{"role": "system", "content": "system instructions " * 120},
|
||||
{"role": "developer", "content": "developer instructions " * 120},
|
||||
{"role": "user", "content": "user prompt " * 120},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": tool_text},
|
||||
]
|
||||
|
||||
result = router.apply(messages, Tokenizer())
|
||||
|
||||
assert result.messages[0]["content"] == messages[0]["content"]
|
||||
assert result.messages[1]["content"] == messages[1]["content"]
|
||||
assert result.messages[2]["content"] == messages[2]["content"]
|
||||
assert result.messages[3]["content"] == "COMPRESSED"
|
||||
assert calls == [tool_text]
|
||||
|
||||
|
||||
def test_pipeline_extension_manager_replaces_events_and_ignores_failures(caplog):
|
||||
recorder = RecordingExtension()
|
||||
manager = PipelineExtensionManager(
|
||||
|
|
|
|||
128
tests/test_cli/test_unwrap_claude.py
Normal file
128
tests/test_cli/test_unwrap_claude.py
Normal file
|
|
@ -0,0 +1,128 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from click.testing import CliRunner
|
||||
|
||||
from headroom.cli import wrap as wrap_cli
|
||||
from headroom.cli.main import main
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runner() -> CliRunner:
|
||||
return CliRunner()
|
||||
|
||||
|
||||
def test_remove_claude_rtk_hooks_preserves_unrelated_hooks(tmp_path: Path) -> None:
|
||||
settings = tmp_path / "settings.json"
|
||||
settings.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"model": "opus",
|
||||
"hooks": {
|
||||
"PreToolUse": [
|
||||
{
|
||||
"matcher": "Bash",
|
||||
"hooks": [
|
||||
{
|
||||
"type": "command",
|
||||
"command": "/Users/test/.claude/hooks/rtk-rewrite.sh",
|
||||
},
|
||||
{"type": "command", "command": "echo keep"},
|
||||
],
|
||||
}
|
||||
],
|
||||
"SessionStart": [
|
||||
{"matcher": "startup", "hooks": [{"type": "command", "command": "keep"}]}
|
||||
],
|
||||
},
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert wrap_cli._remove_claude_rtk_hooks(settings) is True
|
||||
|
||||
payload = json.loads(settings.read_text(encoding="utf-8"))
|
||||
pre_tool_hooks = payload["hooks"]["PreToolUse"][0]["hooks"]
|
||||
assert pre_tool_hooks == [{"type": "command", "command": "echo keep"}]
|
||||
assert payload["hooks"]["SessionStart"][0]["hooks"][0]["command"] == "keep"
|
||||
|
||||
|
||||
def test_unwrap_claude_removes_mcp_rtk_and_stops_proxy(
|
||||
runner: CliRunner,
|
||||
tmp_path: Path,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
home = str(tmp_path)
|
||||
monkeypatch.setenv("HOME", home)
|
||||
monkeypatch.setenv("USERPROFILE", home)
|
||||
claude_dir = tmp_path / ".claude"
|
||||
claude_dir.mkdir()
|
||||
settings = claude_dir / "settings.json"
|
||||
settings.write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"hooks": {
|
||||
"PreToolUse": [
|
||||
{
|
||||
"matcher": "Bash",
|
||||
"hooks": [
|
||||
{"type": "command", "command": str(claude_dir / "rtk-rewrite.sh")}
|
||||
],
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
)
|
||||
+ "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
stopped: list[int] = []
|
||||
unregistered: list[str] = []
|
||||
|
||||
class Registrar:
|
||||
def detect(self) -> bool:
|
||||
return True
|
||||
|
||||
def unregister_server(self, server_name: str) -> bool:
|
||||
unregistered.append(server_name)
|
||||
return True
|
||||
|
||||
with (
|
||||
patch("headroom.mcp_registry.ClaudeRegistrar", return_value=Registrar()),
|
||||
patch(
|
||||
"headroom.cli.wrap._stop_local_proxy_for_unwrap",
|
||||
side_effect=lambda port: stopped.append(port) or "stopped",
|
||||
),
|
||||
):
|
||||
result = runner.invoke(main, ["unwrap", "claude", "--port", "9999"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert unregistered == ["headroom", "codebase-memory-mcp"]
|
||||
assert stopped == [9999]
|
||||
assert "Stopped local Headroom proxy on port 9999" in result.output
|
||||
assert "hooks" not in json.loads(settings.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def test_unwrap_claude_keep_flags_skip_cleanup(
|
||||
runner: CliRunner,
|
||||
) -> None:
|
||||
with (
|
||||
patch("headroom.mcp_registry.ClaudeRegistrar") as registrar,
|
||||
patch("headroom.cli.wrap._remove_claude_rtk_hooks") as remove_rtk,
|
||||
patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy,
|
||||
):
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "claude", "--keep-mcp", "--keep-rtk", "--no-stop-proxy"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
registrar.assert_not_called()
|
||||
remove_rtk.assert_not_called()
|
||||
stop_proxy.assert_not_called()
|
||||
|
|
@ -80,11 +80,30 @@ class TestStripCodexHeadroomBlocks:
|
|||
def test_removes_stray_top_level_model_provider_line(self) -> None:
|
||||
# Old wrap versions left `model_provider = "headroom"` outside markers.
|
||||
content = 'foo = 1\nmodel_provider = "headroom"\nbar = 2\n'
|
||||
cleaned = wrap_mod._strip_codex_headroom_blocks(content)
|
||||
cleaned = wrap_mod._strip_codex_headroom_blocks(content, remove_mcp=True)
|
||||
assert 'model_provider = "headroom"' not in cleaned
|
||||
assert "foo = 1" in cleaned
|
||||
assert "bar = 2" in cleaned
|
||||
|
||||
def test_removes_codex_mcp_blocks(self) -> None:
|
||||
content = (
|
||||
'[profiles.default]\nmodel = "gpt-4o"\n\n'
|
||||
f"{wrap_mod._CODEX_MCP_MARKER}\n"
|
||||
"[mcp_servers.headroom]\n"
|
||||
'command = "headroom"\n'
|
||||
f"{wrap_mod._CODEX_MCP_END}\n\n"
|
||||
f"{wrap_mod._MEMORY_MCP_MARKER}\n"
|
||||
"[mcp_servers.headroom_memory]\n"
|
||||
'command = "python"\n'
|
||||
f"{wrap_mod._MEMORY_MCP_END}\n"
|
||||
)
|
||||
|
||||
cleaned = wrap_mod._strip_codex_headroom_blocks(content, remove_mcp=True)
|
||||
|
||||
assert "[mcp_servers.headroom]" not in cleaned
|
||||
assert "[mcp_servers.headroom_memory]" not in cleaned
|
||||
assert 'model = "gpt-4o"' in cleaned
|
||||
|
||||
|
||||
class TestSnapshotCodexConfig:
|
||||
"""Tests for ``_snapshot_codex_config_if_unwrapped``."""
|
||||
|
|
@ -240,6 +259,35 @@ class TestInjectAndRestoreRoundTrip:
|
|||
assert 'model_provider = "headroom"' not in cleaned
|
||||
assert 'model = "gpt-4o"' in cleaned
|
||||
|
||||
def test_unwrap_without_backup_removes_provider_and_mcp_blocks(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
config_dir = tmp_path / ".codex"
|
||||
config_dir.mkdir()
|
||||
config_file = config_dir / "config.toml"
|
||||
config_file.write_text(
|
||||
'[profiles.default]\nmodel = "gpt-4o"\n\n'
|
||||
f"{wrap_mod._CODEX_TOP_LEVEL_MARKER}\n"
|
||||
'model_provider = "headroom"\n'
|
||||
f"{wrap_mod._CODEX_END_MARKER}\n\n"
|
||||
f"{wrap_mod._CODEX_MCP_MARKER}\n"
|
||||
"[mcp_servers.headroom]\n"
|
||||
'command = "headroom"\n'
|
||||
f"{wrap_mod._CODEX_MCP_END}\n\n"
|
||||
f"{wrap_mod._MEMORY_MCP_MARKER}\n"
|
||||
"[mcp_servers.headroom_memory]\n"
|
||||
'command = "python"\n'
|
||||
f"{wrap_mod._MEMORY_MCP_END}\n"
|
||||
)
|
||||
|
||||
status, _ = wrap_mod._restore_codex_provider_config()
|
||||
|
||||
assert status == "cleaned"
|
||||
cleaned = config_file.read_text()
|
||||
assert 'model = "gpt-4o"' in cleaned
|
||||
assert "headroom" not in cleaned
|
||||
|
||||
def test_unwrap_handles_malformed_prior_config(
|
||||
self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -358,6 +406,62 @@ def test_wrap_codex_prepare_only_creates_backup_and_config(
|
|||
assert backup.read_text() == original
|
||||
|
||||
|
||||
def test_start_proxy_uses_separate_session_for_signal_isolation(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
"""Proxy child should not receive Ctrl-C intended for the wrapped CLI."""
|
||||
popen_kwargs: dict[str, object] = {}
|
||||
|
||||
class FakeProc:
|
||||
returncode = None
|
||||
|
||||
def poll(self) -> None:
|
||||
return None
|
||||
|
||||
def fake_popen(*args: object, **kwargs: object) -> FakeProc:
|
||||
popen_kwargs.update(kwargs)
|
||||
return FakeProc()
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_get_log_path", lambda: tmp_path / "proxy.log")
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda port: True)
|
||||
monkeypatch.setattr(wrap_mod.subprocess, "Popen", fake_popen)
|
||||
|
||||
proc = wrap_mod._start_proxy(8787, agent_type="codex")
|
||||
|
||||
assert isinstance(proc, FakeProc)
|
||||
assert popen_kwargs["start_new_session"] == (wrap_mod.os.name == "posix")
|
||||
|
||||
|
||||
def test_launch_tool_ignores_sigint_in_wrapper(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Ctrl-C should be handled by the child CLI, not kill the proxy from wrapper."""
|
||||
signal_handlers: dict[object, object] = {}
|
||||
|
||||
class FakeCompleted:
|
||||
returncode = 0
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_ensure_proxy", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
wrap_mod.signal, "signal", lambda sig, fn: signal_handlers.setdefault(sig, fn)
|
||||
)
|
||||
monkeypatch.setattr(wrap_mod.subprocess, "run", lambda *args, **kwargs: FakeCompleted())
|
||||
|
||||
with pytest.raises(SystemExit) as exc:
|
||||
wrap_mod._launch_tool(
|
||||
binary="codex",
|
||||
args=(),
|
||||
env={},
|
||||
port=8787,
|
||||
no_proxy=True,
|
||||
tool_label="CODEX",
|
||||
env_vars_display=[],
|
||||
)
|
||||
|
||||
assert exc.value.code == 0
|
||||
assert signal_handlers[wrap_mod.signal.SIGINT] is wrap_mod._ignore_child_sigint
|
||||
|
||||
|
||||
def test_wrap_codex_prepare_only_updates_stale_mcp_proxy_url(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
|
|
@ -407,7 +511,13 @@ def test_unwrap_codex_restores_prior_config_end_to_end(
|
|||
assert wrap_result.exit_code == 0, wrap_result.output
|
||||
assert 'model_provider = "headroom"' in config_file.read_text()
|
||||
|
||||
unwrap_result = runner.invoke(main, ["unwrap", "codex"])
|
||||
stopped: list[int] = []
|
||||
|
||||
with patch(
|
||||
"headroom.cli.wrap._stop_local_proxy_for_unwrap",
|
||||
side_effect=lambda port: stopped.append(port) or "stopped",
|
||||
):
|
||||
unwrap_result = runner.invoke(main, ["unwrap", "codex", "--port", "9999"])
|
||||
assert unwrap_result.exit_code == 0, unwrap_result.output
|
||||
|
||||
# Config must be byte-for-byte what the user had before wrap, and the
|
||||
|
|
@ -416,6 +526,49 @@ def test_unwrap_codex_restores_prior_config_end_to_end(
|
|||
assert config_file.read_text() == original
|
||||
assert 'model_provider = "headroom"' not in config_file.read_text()
|
||||
assert not (tmp_path / ".codex" / "config.toml.headroom-backup").exists()
|
||||
assert stopped == [9999]
|
||||
assert "Stopped local Headroom proxy on port 9999" in unwrap_result.output
|
||||
|
||||
|
||||
def test_unwrap_codex_no_stop_proxy_leaves_proxy_alone(
|
||||
runner: CliRunner, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
_set_test_home(monkeypatch, tmp_path)
|
||||
|
||||
with patch("headroom.cli.wrap._stop_local_proxy_for_unwrap") as stop_proxy:
|
||||
result = runner.invoke(main, ["unwrap", "codex", "--no-stop-proxy"])
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
stop_proxy.assert_not_called()
|
||||
|
||||
|
||||
def test_stop_local_proxy_for_unwrap_kills_identified_headroom_proxy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
killed: list[tuple[int, int]] = []
|
||||
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda port: True)
|
||||
monkeypatch.setattr(wrap_mod, "_query_proxy_config", lambda port: {"pid": "12345"})
|
||||
monkeypatch.setattr(
|
||||
wrap_mod,
|
||||
"_kill_proxy_by_pid",
|
||||
lambda pid, port: killed.append((pid, port)) or True,
|
||||
)
|
||||
|
||||
assert wrap_mod._stop_local_proxy_for_unwrap(8787) == "stopped"
|
||||
assert killed == [(12345, 8787)]
|
||||
|
||||
|
||||
def test_stop_local_proxy_for_unwrap_refuses_unidentified_listener(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(wrap_mod, "_check_proxy", lambda port: True)
|
||||
monkeypatch.setattr(wrap_mod, "_query_proxy_config", lambda port: None)
|
||||
|
||||
with patch("headroom.cli.wrap._kill_proxy_by_pid") as kill_proxy:
|
||||
assert wrap_mod._stop_local_proxy_for_unwrap(8787) == "unidentified"
|
||||
|
||||
kill_proxy.assert_not_called()
|
||||
|
||||
|
||||
def test_unwrap_codex_is_safe_noop_with_no_prior_wrap(
|
||||
|
|
|
|||
|
|
@ -407,6 +407,28 @@ def test_build_openclaw_unwrap_entry_preserves_top_level_metadata() -> None:
|
|||
assert entry["config"] == {"customFlag": True}
|
||||
|
||||
|
||||
def test_unwrap_openclaw_stops_proxy_by_default(runner: CliRunner) -> None:
|
||||
calls: list[dict] = []
|
||||
|
||||
def which(name: str) -> str | None:
|
||||
return "openclaw" if name == "openclaw" else None
|
||||
|
||||
with patch("headroom.cli.wrap.shutil.which", side_effect=which):
|
||||
with patch("headroom.cli.wrap.subprocess.run", side_effect=_make_successful_run(calls)):
|
||||
with patch(
|
||||
"headroom.cli.wrap._stop_local_proxy_for_unwrap",
|
||||
return_value="stopped",
|
||||
) as stop_proxy:
|
||||
result = runner.invoke(
|
||||
main,
|
||||
["unwrap", "openclaw", "--proxy-port", "9999", "--no-restart"],
|
||||
)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
stop_proxy.assert_called_once_with(9999)
|
||||
assert "Stopped local Headroom proxy on port 9999" in result.output
|
||||
|
||||
|
||||
def test_wrap_openclaw_no_auto_start_does_not_default_python_path(
|
||||
runner: CliRunner, plugin_dir: Path
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -803,8 +803,40 @@ class TestCompressionStoreSearch:
|
|||
results = store.search(hash_key, "query")
|
||||
assert results == []
|
||||
|
||||
def test_search_plain_text_returns_matching_chunks(self, store: CompressionStore):
|
||||
"""search() can find content in Kompress-style plain-text originals."""
|
||||
original = (
|
||||
"The OpenAI handler contains def _compress_openai_responses_payload "
|
||||
"for Responses API live-zone compression. Other text is irrelevant."
|
||||
)
|
||||
hash_key = store.store(original=original, compressed="compressed")
|
||||
|
||||
results = store.search(hash_key, "def _compress_openai_responses_payload")
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["type"] == "text"
|
||||
assert "_compress_openai_responses_payload" in results[0]["text"]
|
||||
|
||||
def test_search_json_object_returns_matching_leaf(self, store: CompressionStore):
|
||||
"""search() can find values inside JSON objects, not only arrays."""
|
||||
original = json.dumps(
|
||||
{
|
||||
"module": {
|
||||
"name": "openai",
|
||||
"function": "_compress_openai_responses_payload",
|
||||
}
|
||||
}
|
||||
)
|
||||
hash_key = store.store(original=original, compressed="{}")
|
||||
|
||||
results = store.search(hash_key, "_compress_openai_responses_payload")
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0]["path"] == "module.function"
|
||||
assert results[0]["value"] == "_compress_openai_responses_payload"
|
||||
|
||||
def test_search_non_array_returns_empty(self, store: CompressionStore):
|
||||
"""search() returns empty for non-array content."""
|
||||
"""search() returns empty for JSON objects without matching leaves."""
|
||||
hash_key = store.store(original=json.dumps({"key": "value"}), compressed="{}")
|
||||
results = store.search(hash_key, "query")
|
||||
assert results == []
|
||||
|
|
|
|||
160
tests/test_compression_units.py
Normal file
160
tests/test_compression_units.py
Normal file
|
|
@ -0,0 +1,160 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from headroom.transforms.compression_units import (
|
||||
CompressionUnit,
|
||||
RoutedCompressionUnit,
|
||||
compress_unit_with_router,
|
||||
compress_units_with_router,
|
||||
)
|
||||
from headroom.transforms.content_router import (
|
||||
CompressionStrategy,
|
||||
RouterCompressionResult,
|
||||
)
|
||||
|
||||
|
||||
class TokenCounter:
|
||||
def count_text(self, text: str) -> int:
|
||||
return len(text.split())
|
||||
|
||||
|
||||
class Router:
|
||||
def __init__(self, compressed: str):
|
||||
self.compressed = compressed
|
||||
|
||||
def compress(self, content: str, **_kwargs):
|
||||
return RouterCompressionResult(
|
||||
compressed=self.compressed,
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
|
||||
def test_compression_unit_accepts_token_shrinking_replacement():
|
||||
result = compress_unit_with_router(
|
||||
CompressionUnit(
|
||||
text="alpha beta gamma delta epsilon",
|
||||
provider="openai",
|
||||
endpoint="responses",
|
||||
role="tool",
|
||||
item_type="local_shell_call_output",
|
||||
min_bytes=1,
|
||||
),
|
||||
router=Router("alpha beta"),
|
||||
tokenizer=TokenCounter(),
|
||||
)
|
||||
|
||||
assert result.modified is True
|
||||
assert result.tokens_saved == 3
|
||||
assert result.compressed == "alpha beta"
|
||||
assert "router:openai:responses:local_shell_call_output:kompress" in result.transforms_applied
|
||||
|
||||
|
||||
def test_compression_unit_rejects_non_shrinking_replacement():
|
||||
result = compress_unit_with_router(
|
||||
CompressionUnit(
|
||||
text="alpha beta",
|
||||
provider="anthropic",
|
||||
endpoint="messages",
|
||||
role="tool",
|
||||
item_type="tool_result",
|
||||
min_bytes=1,
|
||||
),
|
||||
router=Router("alpha beta gamma"),
|
||||
tokenizer=TokenCounter(),
|
||||
)
|
||||
|
||||
assert result.modified is False
|
||||
assert result.reason == "rejected_not_smaller"
|
||||
assert result.original == "alpha beta"
|
||||
|
||||
|
||||
def test_compression_unit_respects_cache_zone_and_floor():
|
||||
frozen = compress_unit_with_router(
|
||||
CompressionUnit(
|
||||
text="alpha beta gamma delta",
|
||||
provider="anthropic",
|
||||
endpoint="messages",
|
||||
role="tool",
|
||||
item_type="tool_result",
|
||||
cache_zone="frozen",
|
||||
min_bytes=1,
|
||||
),
|
||||
router=Router("alpha"),
|
||||
tokenizer=TokenCounter(),
|
||||
)
|
||||
small = compress_unit_with_router(
|
||||
CompressionUnit(
|
||||
text="small text",
|
||||
provider="openai",
|
||||
endpoint="responses",
|
||||
role="tool",
|
||||
item_type="function_call_output",
|
||||
min_bytes=500,
|
||||
),
|
||||
router=Router("small"),
|
||||
tokenizer=TokenCounter(),
|
||||
)
|
||||
|
||||
assert frozen.modified is False
|
||||
assert frozen.reason == "cache_zone_frozen"
|
||||
assert small.modified is False
|
||||
assert small.reason == "below_unit_floor"
|
||||
|
||||
|
||||
def test_batch_compression_preserves_provider_slot_references():
|
||||
routed = [
|
||||
RoutedCompressionUnit(
|
||||
unit=CompressionUnit(
|
||||
text="alpha beta gamma",
|
||||
provider="openai",
|
||||
endpoint="responses",
|
||||
role="tool",
|
||||
item_type="function_call_output",
|
||||
min_bytes=1,
|
||||
),
|
||||
slot=("input", 3, "output"),
|
||||
),
|
||||
RoutedCompressionUnit(
|
||||
unit=CompressionUnit(
|
||||
text="one two three",
|
||||
provider="gemini",
|
||||
endpoint="generateContent",
|
||||
role="user",
|
||||
item_type="part.text",
|
||||
min_bytes=1,
|
||||
),
|
||||
slot={"path": ["contents", 0, "parts", 0, "text"]},
|
||||
),
|
||||
]
|
||||
|
||||
results = compress_units_with_router(
|
||||
routed,
|
||||
router=Router("short"),
|
||||
tokenizer=TokenCounter(),
|
||||
)
|
||||
|
||||
assert results[0][0] == ("input", 3, "output")
|
||||
assert results[1][0] == {"path": ["contents", 0, "parts", 0, "text"]}
|
||||
assert [result.modified for _slot, result in results] == [True, False]
|
||||
|
||||
|
||||
def test_compress_unit_protects_prompt_roles() -> None:
|
||||
for role, reason in [
|
||||
("user", "protected_user_message"),
|
||||
("developer", "protected_system_message"),
|
||||
("system", "protected_system_message"),
|
||||
("assistant", "protected_assistant_message"),
|
||||
]:
|
||||
unit = CompressionUnit(
|
||||
text="alpha beta gamma delta",
|
||||
provider="openai",
|
||||
endpoint="responses",
|
||||
role=role,
|
||||
item_type="message",
|
||||
min_bytes=1,
|
||||
)
|
||||
|
||||
result = compress_unit_with_router(unit, router=Router("alpha"), tokenizer=TokenCounter())
|
||||
|
||||
assert result.modified is False
|
||||
assert result.reason == reason
|
||||
|
|
@ -246,6 +246,26 @@ def test_handle_openai_responses_routes_chatgpt_auth_to_backend_api(monkeypatch)
|
|||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_handle_openai_responses_routes_api_key_auth_direct_to_openai(monkeypatch):
|
||||
request = _build_request(
|
||||
{"model": "gpt-4o-mini", "input": "hello"},
|
||||
{"Authorization": "Bearer sk-test"},
|
||||
)
|
||||
handler = _DummyOpenAIHandler()
|
||||
|
||||
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())
|
||||
|
||||
response = anyio.run(handler.handle_openai_responses, request)
|
||||
|
||||
assert handler.captured_request is not None
|
||||
method, url, headers, body = handler.captured_request
|
||||
assert method == "POST"
|
||||
assert url == "https://api.openai.com/v1/responses"
|
||||
assert headers.get("ChatGPT-Account-ID") is None
|
||||
assert body["input"] == "hello"
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
def test_handle_openai_responses_stream_skips_python_compression(monkeypatch):
|
||||
"""PR-C5: Python no longer compresses /v1/responses (Rust handles it
|
||||
natively). The streaming forward path must still fire — only the
|
||||
|
|
|
|||
|
|
@ -288,6 +288,7 @@ async def test_ws_session_metrics_include_response_completed_usage():
|
|||
assert recorded["input_tokens"] == 100
|
||||
assert recorded["output_tokens"] == 12
|
||||
assert recorded["cache_read_tokens"] == 75
|
||||
assert recorded["cache_write_tokens"] == 25
|
||||
assert recorded["uncached_input_tokens"] == 25
|
||||
|
||||
|
||||
|
|
|
|||
194
tests/test_openai_responses_compression_units.py
Normal file
194
tests/test_openai_responses_compression_units.py
Normal file
|
|
@ -0,0 +1,194 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from types import MethodType, SimpleNamespace
|
||||
|
||||
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
|
||||
from headroom.transforms.content_router import (
|
||||
CompressionStrategy,
|
||||
ContentRouter,
|
||||
RouterCompressionResult,
|
||||
)
|
||||
|
||||
|
||||
class TokenCounter:
|
||||
def count_text(self, text: str) -> int:
|
||||
return len(text.split())
|
||||
|
||||
|
||||
def _handler_with_router(router: ContentRouter) -> OpenAIHandlerMixin:
|
||||
handler = OpenAIHandlerMixin()
|
||||
handler.openai_pipeline = SimpleNamespace(transforms=[router])
|
||||
handler.openai_provider = SimpleNamespace(
|
||||
get_token_counter=lambda _model: TokenCounter(),
|
||||
)
|
||||
return handler
|
||||
|
||||
|
||||
def test_openai_responses_adapter_compresses_only_live_text_slots():
|
||||
router = ContentRouter()
|
||||
|
||||
def compress(self, content: str, **_kwargs):
|
||||
return RouterCompressionResult(
|
||||
compressed="kept words",
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
router.compress = MethodType(compress, router)
|
||||
handler = _handler_with_router(router)
|
||||
long_text = " ".join(f"word{i}" for i in range(180))
|
||||
payload = {
|
||||
"model": "gpt-5",
|
||||
"input": [
|
||||
{"type": "reasoning", "encrypted_content": long_text},
|
||||
{"type": "function_call", "arguments": long_text},
|
||||
{"type": "local_shell_call_output", "call_id": "c1", "output": long_text},
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": long_text}],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
new_payload, modified, saved, transforms = (
|
||||
handler._compress_openai_responses_live_text_units_with_router(
|
||||
payload,
|
||||
model="gpt-5",
|
||||
request_id="req_test",
|
||||
)
|
||||
)
|
||||
|
||||
assert modified is True
|
||||
assert saved > 0
|
||||
assert new_payload["input"][0]["encrypted_content"] == long_text
|
||||
assert new_payload["input"][1]["arguments"] == long_text
|
||||
assert new_payload["input"][2]["output"] == "kept words"
|
||||
assert new_payload["input"][3]["content"][0]["text"] == long_text
|
||||
assert any(t.startswith("router:openai:responses:") for t in transforms)
|
||||
|
||||
|
||||
def test_openai_responses_adapter_preserves_headroom_retrieve_outputs():
|
||||
router = ContentRouter()
|
||||
|
||||
def compress(self, content: str, **_kwargs):
|
||||
return RouterCompressionResult(
|
||||
compressed="compressed retrieve output",
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
router.compress = MethodType(compress, router)
|
||||
handler = _handler_with_router(router)
|
||||
retrieved = " ".join(f"retrieved{i}" for i in range(180))
|
||||
payload = {
|
||||
"model": "gpt-5",
|
||||
"input": [
|
||||
{
|
||||
"type": "function_call",
|
||||
"call_id": "call_retrieve",
|
||||
"name": "mcp__headroom__headroom_retrieve",
|
||||
"arguments": "{}",
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_retrieve",
|
||||
"output": retrieved,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
new_payload, modified, saved, transforms = (
|
||||
handler._compress_openai_responses_live_text_units_with_router(
|
||||
payload,
|
||||
model="gpt-5",
|
||||
request_id="req_test",
|
||||
)
|
||||
)
|
||||
|
||||
assert modified is False
|
||||
assert saved == 0
|
||||
assert transforms == []
|
||||
assert new_payload == payload
|
||||
|
||||
|
||||
def test_openai_responses_adapter_keeps_small_and_opaque_items():
|
||||
router = ContentRouter()
|
||||
|
||||
def compress(self, content: str, **_kwargs):
|
||||
return RouterCompressionResult(
|
||||
compressed="short",
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
router.compress = MethodType(compress, router)
|
||||
handler = _handler_with_router(router)
|
||||
payload = {
|
||||
"model": "gpt-5",
|
||||
"input": [
|
||||
{"type": "local_shell_call_output", "call_id": "c1", "output": "too small"},
|
||||
{"type": "compaction", "encrypted_content": " ".join(["secret"] * 200)},
|
||||
],
|
||||
}
|
||||
|
||||
new_payload, modified, saved, transforms = (
|
||||
handler._compress_openai_responses_live_text_units_with_router(
|
||||
payload,
|
||||
model="gpt-5",
|
||||
request_id="req_test",
|
||||
)
|
||||
)
|
||||
|
||||
assert modified is False
|
||||
assert saved == 0
|
||||
assert transforms == []
|
||||
assert new_payload == payload
|
||||
|
||||
|
||||
def test_openai_responses_payload_routes_through_content_router_without_rust(
|
||||
monkeypatch,
|
||||
):
|
||||
router = ContentRouter()
|
||||
|
||||
def compress(self, content: str, **_kwargs):
|
||||
return RouterCompressionResult(
|
||||
compressed="compressed fallback",
|
||||
original=content,
|
||||
strategy_used=CompressionStrategy.KOMPRESS,
|
||||
)
|
||||
|
||||
router.compress = MethodType(compress, router)
|
||||
handler = _handler_with_router(router)
|
||||
|
||||
import headroom._core as core
|
||||
|
||||
def rust_must_not_run(*_args, **_kwargs):
|
||||
raise AssertionError("Responses payload compression should route through ContentRouter")
|
||||
|
||||
monkeypatch.setattr(core, "compress_openai_responses_live_zone", rust_must_not_run)
|
||||
|
||||
payload = {
|
||||
"model": "gpt-5",
|
||||
"input": [
|
||||
{
|
||||
"type": "local_shell_call_output",
|
||||
"call_id": "c1",
|
||||
"output": " ".join(f"word{i}" for i in range(180)),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
new_payload, modified, saved, transforms, reason, _, _ = (
|
||||
handler._compress_openai_responses_payload(
|
||||
payload,
|
||||
model="gpt-5",
|
||||
request_id="req_router",
|
||||
)
|
||||
)
|
||||
|
||||
assert modified is True
|
||||
assert saved > 0
|
||||
assert reason is None
|
||||
assert new_payload["input"][0]["output"] == "compressed fallback"
|
||||
assert any(t.startswith("router:openai:responses:") for t in transforms)
|
||||
|
|
@ -118,6 +118,35 @@ class TestCCRRetrieveEndpoint:
|
|||
assert "results" in data
|
||||
assert data["count"] >= 1
|
||||
|
||||
def test_retrieve_with_search_plain_text_original(self, client):
|
||||
"""Query retrieval searches plain-text originals stored by Kompress."""
|
||||
store = get_compression_store()
|
||||
original = (
|
||||
"Codex WS compression stores plain text originals. "
|
||||
"The target symbol is _compress_openai_responses_payload."
|
||||
)
|
||||
hash_key = store.store(original=original, compressed="compressed")
|
||||
|
||||
response = client.post(
|
||||
"/v1/retrieve",
|
||||
json={"hash": hash_key, "query": "_compress_openai_responses_payload"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
data = response.json()
|
||||
assert data["hash"] == hash_key
|
||||
assert data["count"] == 1
|
||||
assert data["results"][0]["type"] == "text"
|
||||
assert "_compress_openai_responses_payload" in data["results"][0]["text"]
|
||||
|
||||
def test_retrieve_with_search_nonexistent_hash_returns_404(self, client):
|
||||
"""Query mode should not mask a missing hash as an empty search."""
|
||||
response = client.post(
|
||||
"/v1/retrieve",
|
||||
json={"hash": "nonexistent123", "query": "anything"},
|
||||
)
|
||||
assert response.status_code == 404
|
||||
|
||||
def test_retrieve_increments_count(self, client):
|
||||
"""Each retrieval increments the retrieval count."""
|
||||
store = get_compression_store()
|
||||
|
|
@ -185,6 +214,21 @@ class TestCCRRetrieveGetEndpoint:
|
|||
# Results should be a list (may be empty if BM25 threshold not met)
|
||||
assert isinstance(data["results"], list)
|
||||
|
||||
def test_get_retrieve_with_query_plain_text_original(self, client):
|
||||
"""GET query retrieval searches plain-text originals."""
|
||||
store = get_compression_store()
|
||||
hash_key = store.store(
|
||||
original="plain text contains _compress_openai_responses_payload",
|
||||
compressed="plain text",
|
||||
)
|
||||
|
||||
response = client.get(f"/v1/retrieve/{hash_key}?query=_compress_openai_responses_payload")
|
||||
assert response.status_code == 200
|
||||
|
||||
data = response.json()
|
||||
assert data["count"] == 1
|
||||
assert data["results"][0]["type"] == "text"
|
||||
|
||||
def test_get_retrieve_nonexistent(self, client):
|
||||
"""GET with nonexistent hash returns 404."""
|
||||
response = client.get("/v1/retrieve/nonexistent123")
|
||||
|
|
|
|||
|
|
@ -30,17 +30,25 @@ class _ToinStub:
|
|||
@pytest.fixture(autouse=True)
|
||||
def _reset_rtk_stats_cache() -> None:
|
||||
proxy_helpers._rtk_stats_cache.update({"expires_at": 0.0, "has_value": False, "value": None})
|
||||
proxy_helpers._rtk_session_baseline.update(
|
||||
{"initialized": False, "total_commands": 0, "tokens_saved": 0}
|
||||
)
|
||||
|
||||
|
||||
def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
now = {"value": 100.0}
|
||||
calls = {"run": 0}
|
||||
totals = [
|
||||
{"total_commands": 7, "total_saved": 1234},
|
||||
{"total_commands": 9, "total_saved": 1500},
|
||||
]
|
||||
|
||||
def _fake_run(*args, **kwargs):
|
||||
calls["run"] += 1
|
||||
summary = totals[min(calls["run"] - 1, len(totals) - 1)]
|
||||
return SimpleNamespace(
|
||||
returncode=0,
|
||||
stdout=json.dumps({"summary": {"total_commands": 7, "total_saved": 1234}}),
|
||||
stdout=json.dumps({"summary": summary}),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(proxy_helpers.time, "monotonic", lambda: now["value"])
|
||||
|
|
@ -53,8 +61,8 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch
|
|||
assert first == second
|
||||
assert first == {
|
||||
"installed": True,
|
||||
"total_commands": 7,
|
||||
"tokens_saved": 1234,
|
||||
"total_commands": 0,
|
||||
"tokens_saved": 0,
|
||||
"avg_savings_pct": 0.0,
|
||||
}
|
||||
assert calls["run"] == 1
|
||||
|
|
@ -62,7 +70,12 @@ def test_get_rtk_stats_memoizes_subprocess_calls(monkeypatch: pytest.MonkeyPatch
|
|||
now["value"] += proxy_helpers.RTK_STATS_CACHE_TTL_SECONDS + 0.1
|
||||
third = proxy_helpers._get_rtk_stats()
|
||||
|
||||
assert third == first
|
||||
assert third == {
|
||||
"installed": True,
|
||||
"total_commands": 2,
|
||||
"tokens_saved": 266,
|
||||
"avg_savings_pct": 0.0,
|
||||
}
|
||||
assert calls["run"] == 2
|
||||
|
||||
|
||||
|
|
@ -135,8 +148,71 @@ def test_stats_cached_query_reuses_short_ttl_snapshot(monkeypatch: pytest.Monkey
|
|||
assert first.json()["tokens"]["saved"] == 5
|
||||
assert first.json()["tokens"]["proxy_compression_saved"] == 0
|
||||
assert first.json()["tokens"]["rtk_saved"] == 5
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["tokens"] == 5
|
||||
assert first.json()["tokens"]["all_layers_saved"] == 5
|
||||
assert (
|
||||
first.json()["tokens"]["savings_percent"]
|
||||
== first.json()["tokens"]["all_layers_savings_percent"]
|
||||
)
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["tokens"] == 0
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["rtk_tokens"] == 5
|
||||
assert first.json()["savings"]["by_layer"]["compression"]["all_layers_tokens"] == 5
|
||||
|
||||
|
||||
def test_stats_reset_clears_runtime_proxy_counters(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
pytest.importorskip("fastapi")
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
import headroom.proxy.server as server
|
||||
from headroom.proxy.loopback_guard import require_loopback
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_store",
|
||||
lambda: _StatsStub({"store": 0}, "store", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_telemetry_collector",
|
||||
lambda: _StatsStub({"telemetry": 0}, "telemetry", {}),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"get_compression_feedback",
|
||||
lambda: _StatsStub({"feedback": 0}, "feedback", {}),
|
||||
)
|
||||
monkeypatch.setattr(server, "_get_rtk_stats", lambda: None)
|
||||
monkeypatch.setattr(server, "get_toin", lambda: _ToinStub())
|
||||
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=False,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
)
|
||||
)
|
||||
app.dependency_overrides[require_loopback] = lambda: None
|
||||
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
proxy.metrics.tokens_saved_total = 123
|
||||
proxy.metrics.tokens_input_total = 456
|
||||
proxy.metrics.requests_total = 2
|
||||
|
||||
before = client.get("/stats").json()
|
||||
reset = client.post("/stats/reset")
|
||||
after = client.get("/stats").json()
|
||||
|
||||
assert before["tokens"]["proxy_compression_saved"] == 123
|
||||
assert reset.status_code == 200
|
||||
assert after["tokens"]["proxy_compression_saved"] == 0
|
||||
assert after["tokens"]["input"] == 0
|
||||
assert after["requests"]["total"] == 0
|
||||
|
||||
|
||||
def test_dashboard_uses_cached_stats_and_lazy_history_feed_polling() -> None:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from unittest.mock import patch
|
|||
|
||||
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
|
||||
from headroom.proxy.handlers.openai import OpenAIHandlerMixin, _decode_openai_bearer_payload
|
||||
from headroom.proxy.helpers import _headroom_bypass_enabled
|
||||
|
||||
|
||||
def _jwt(payload: object) -> str:
|
||||
|
|
@ -79,6 +80,17 @@ def test_openai_handler_prefix_helpers_cover_edge_cases() -> None:
|
|||
assert changed == 1
|
||||
|
||||
|
||||
def test_headroom_bypass_helper_is_transport_neutral() -> None:
|
||||
assert _headroom_bypass_enabled({"x-headroom-bypass": "true"}) is True
|
||||
assert _headroom_bypass_enabled({"x-headroom-bypass": " TRUE "}) is True
|
||||
assert _headroom_bypass_enabled({"x-headroom-mode": "passthrough"}) is True
|
||||
assert _headroom_bypass_enabled({"x-headroom-mode": " PASSTHROUGH "}) is True
|
||||
assert _headroom_bypass_enabled({"x-headroom-bypass": "false"}) is False
|
||||
assert _headroom_bypass_enabled({}) is False
|
||||
assert _headroom_bypass_enabled(None) is False
|
||||
assert OpenAIHandlerMixin._headroom_bypass_enabled({"x-headroom-bypass": "true"}) is True
|
||||
|
||||
|
||||
def test_anthropic_tool_sort_and_context_append_helpers() -> None:
|
||||
tools = [
|
||||
{"type": "function", "function": {"name": "beta"}},
|
||||
|
|
|
|||
109
tests/test_proxy_openai_responses_bypass.py
Normal file
109
tests/test_proxy_openai_responses_bypass.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
|
||||
from headroom.proxy.loopback_guard import require_loopback # noqa: E402
|
||||
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
|
||||
|
||||
|
||||
class _MemoryHandler:
|
||||
def __init__(self) -> None:
|
||||
self.search_calls = 0
|
||||
self.tool_calls = 0
|
||||
self.config = SimpleNamespace(inject_context=True, inject_tools=True)
|
||||
|
||||
async def search_and_format_context(self, user_id: str, messages: list[dict[str, Any]]) -> str:
|
||||
self.search_calls += 1
|
||||
return "memory context that must not be injected"
|
||||
|
||||
def compute_memory_tool_definitions(self, provider: str) -> list[dict[str, Any]]:
|
||||
self.tool_calls += 1
|
||||
return [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "memory_search",
|
||||
"description": "search memory",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
def has_memory_tool_calls(self, response: dict[str, Any], provider: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def test_responses_bypass_skips_memory_and_compression_mutation() -> None:
|
||||
app = create_app(
|
||||
ProxyConfig(
|
||||
optimize=True,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
)
|
||||
)
|
||||
app.dependency_overrides[require_loopback] = lambda: None
|
||||
|
||||
original_input = [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": "summarize"}],
|
||||
},
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": "large tool output " * 200,
|
||||
},
|
||||
]
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
with TestClient(app) as client:
|
||||
proxy = client.app.state.proxy
|
||||
memory_handler = _MemoryHandler()
|
||||
proxy.memory_handler = memory_handler
|
||||
|
||||
async def _fake_retry(
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: dict[str, Any],
|
||||
stream: bool = False,
|
||||
**kwargs: Any,
|
||||
) -> httpx.Response:
|
||||
captured["body"] = body
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "resp_1",
|
||||
"output": [],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
|
||||
proxy._retry_request = _fake_retry
|
||||
|
||||
response = client.post(
|
||||
"/v1/responses",
|
||||
headers={
|
||||
"authorization": "Bearer test-key",
|
||||
"x-headroom-bypass": "true",
|
||||
"x-headroom-user-id": "user-1",
|
||||
},
|
||||
json={"model": "gpt-4o-mini", "input": original_input},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["body"]["input"] == original_input
|
||||
assert "tools" not in captured["body"]
|
||||
assert memory_handler.search_calls == 0
|
||||
assert memory_handler.tool_calls == 0
|
||||
|
|
@ -22,6 +22,7 @@ pytest.importorskip("httpx")
|
|||
|
||||
from fastapi.testclient import TestClient # noqa: E402
|
||||
|
||||
from headroom.proxy.loopback_guard import require_loopback # noqa: E402
|
||||
from headroom.proxy.server import ProxyConfig, create_app # noqa: E402
|
||||
|
||||
|
||||
|
|
@ -35,6 +36,7 @@ def openai_responses_client():
|
|||
cost_tracking_enabled=False,
|
||||
)
|
||||
app = create_app(config)
|
||||
app.dependency_overrides[require_loopback] = lambda: None
|
||||
with TestClient(app) as client:
|
||||
yield client
|
||||
|
||||
|
|
@ -295,8 +297,9 @@ class TestOpenAIResponsesCompression:
|
|||
assert response.status_code == 200
|
||||
|
||||
stats = openai_responses_client.get("/stats").json()
|
||||
# With bypass, no tokens should be saved
|
||||
assert stats["tokens"]["saved"] == 0
|
||||
# With bypass, proxy compression should not save tokens. The headline
|
||||
# saved count may include RTK CLI savings from the developer shell.
|
||||
assert stats["tokens"]["proxy_compression_saved"] == 0
|
||||
|
||||
|
||||
class TestOpenAIResponsesStats:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ The non-streaming Anthropic path and the Bedrock streaming path were the
|
|||
only ones that called `self.logger.log(...)`.
|
||||
"""
|
||||
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
|
|
@ -49,6 +50,35 @@ def _stream_state(output_tokens: int = 42) -> dict:
|
|||
}
|
||||
|
||||
|
||||
def test_parse_openai_responses_completed_usage_from_sse_buffer():
|
||||
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
||||
completed = {
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_1",
|
||||
"usage": {
|
||||
"input_tokens": 844_000,
|
||||
"input_tokens_details": {"cached_tokens": 657_400},
|
||||
"output_tokens": 6_635,
|
||||
},
|
||||
},
|
||||
}
|
||||
state = {
|
||||
"sse_buffer": bytearray(
|
||||
f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode()
|
||||
)
|
||||
}
|
||||
|
||||
usage = proxy._parse_sse_usage_from_buffer(state, "openai")
|
||||
|
||||
assert usage == {
|
||||
"input_tokens": 844_000,
|
||||
"output_tokens": 6_635,
|
||||
"cache_read_input_tokens": 657_400,
|
||||
}
|
||||
assert state["sse_buffer"] == bytearray()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_stream_response_logs_request_for_feed():
|
||||
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
||||
|
|
@ -153,6 +183,51 @@ async def test_finalize_stream_response_handles_zero_original_tokens():
|
|||
assert entries[0]["savings_percent"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_openai_responses_stream_uses_provider_usage_for_dashboard():
|
||||
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
||||
state = _stream_state(output_tokens=6_635)
|
||||
state["input_tokens"] = 844_000
|
||||
state["cache_read_input_tokens"] = 657_400
|
||||
|
||||
await proxy._finalize_stream_response(
|
||||
body={"model": "gpt-5.5", "input": [{"type": "message", "role": "user"}]},
|
||||
provider="openai",
|
||||
model="gpt-5.5",
|
||||
request_id="req-openai-responses-stream",
|
||||
original_tokens=0,
|
||||
optimized_tokens=0,
|
||||
tokens_saved=663_000,
|
||||
transforms_applied=["openai_responses_live_zone"],
|
||||
optimization_latency=26.0,
|
||||
stream_state=state,
|
||||
start_time=0.0,
|
||||
)
|
||||
|
||||
entries = proxy.logger.get_recent(10)
|
||||
assert len(entries) == 1
|
||||
entry = entries[0]
|
||||
assert entry["input_tokens_optimized"] == 844_000
|
||||
assert entry["input_tokens_original"] == 1_507_000
|
||||
assert entry["tokens_saved"] == 663_000
|
||||
assert entry["savings_percent"] == pytest.approx(663_000 / 1_507_000 * 100)
|
||||
assert entry["output_tokens"] == 6_635
|
||||
|
||||
proxy.metrics.record_request.assert_awaited_once()
|
||||
metrics_kwargs = proxy.metrics.record_request.await_args.kwargs
|
||||
assert metrics_kwargs["input_tokens"] == 844_000
|
||||
assert metrics_kwargs["output_tokens"] == 6_635
|
||||
assert metrics_kwargs["tokens_saved"] == 663_000
|
||||
assert metrics_kwargs["cache_read_tokens"] == 657_400
|
||||
assert metrics_kwargs["uncached_input_tokens"] == 186_600
|
||||
|
||||
proxy.cost_tracker.record_tokens.assert_called_once()
|
||||
cost_args, cost_kwargs = proxy.cost_tracker.record_tokens.call_args
|
||||
assert cost_args[:3] == ("gpt-5.5", 663_000, 844_000)
|
||||
assert cost_kwargs["cache_read_tokens"] == 657_400
|
||||
assert cost_kwargs["uncached_tokens"] == 186_600
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalize_stream_response_no_op_when_logger_disabled():
|
||||
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
||||
|
|
|
|||
|
|
@ -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, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(body, "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(body, auth_mode, "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(body, "payg", "")
|
||||
out, modified, _saved, _transforms, _reason = compress(body, "payg", "")
|
||||
assert isinstance(out, bytes)
|
||||
assert modified is False
|
||||
|
||||
|
|
@ -118,13 +118,15 @@ class TestNoExceptionsLeak:
|
|||
|
||||
def test_garbage_bytes_no_raise(self):
|
||||
compress = _ensure_binding()
|
||||
out, modified, _saved, _transforms = compress(b"\xff\xfe\x00\xff", "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = 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, _saved, _transforms = compress(b"", "payg", "gpt-4o-mini")
|
||||
out, modified, _saved, _transforms, _reason = compress(b"", "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
assert out == b""
|
||||
|
||||
|
|
@ -142,11 +144,12 @@ class TestTelemetryFields:
|
|||
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")
|
||||
out, modified, saved, transforms, reason = compress(body, "payg", "gpt-4o-mini")
|
||||
assert modified is False
|
||||
assert out == body
|
||||
assert saved == 0
|
||||
assert transforms == []
|
||||
assert reason == "no_eligible_items"
|
||||
|
||||
def test_field_types(self):
|
||||
"""Pin the wire shape so downstream callers don't break."""
|
||||
|
|
@ -154,13 +157,14 @@ class TestTelemetryFields:
|
|||
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 len(result) == 5
|
||||
out, modified, saved, transforms, reason = 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)
|
||||
assert reason is None or isinstance(reason, str)
|
||||
|
||||
def test_large_local_shell_output_compresses_with_telemetry(self):
|
||||
"""End-to-end check: a payload large enough to clear the
|
||||
|
|
@ -186,10 +190,11 @@ class TestTelemetryFields:
|
|||
],
|
||||
}
|
||||
).encode()
|
||||
out, modified, saved, transforms = compress(body, "payg", "gpt-4o")
|
||||
out, modified, saved, transforms, reason = compress(body, "payg", "gpt-4o")
|
||||
assert modified is True
|
||||
assert saved > 0
|
||||
assert transforms, "expected at least one strategy in transforms"
|
||||
assert reason is None
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ def _ensure_binding():
|
|||
def _ws_compress_first_frame(
|
||||
first_msg_raw: str,
|
||||
auth_mode_value: str = "payg",
|
||||
bypass: bool = False,
|
||||
) -> tuple[str, bool]:
|
||||
"""Replicates the WS-handler compression block as a pure function.
|
||||
|
||||
|
|
@ -45,6 +46,9 @@ def _ws_compress_first_frame(
|
|||
up a full WebSocket fixture. If you change the handler's
|
||||
compression block, mirror it here so the tests catch the drift.
|
||||
"""
|
||||
if bypass:
|
||||
return first_msg_raw, False
|
||||
|
||||
compress = _ensure_binding()
|
||||
|
||||
try:
|
||||
|
|
@ -60,7 +64,9 @@ 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, _saved, _transforms = compress(inner_bytes, auth_mode_value, model)
|
||||
new_bytes, modified, _saved, _transforms, _reason = compress(
|
||||
inner_bytes, auth_mode_value, model
|
||||
)
|
||||
if not modified:
|
||||
return first_msg_raw, False
|
||||
|
||||
|
|
@ -111,6 +117,37 @@ class TestWrappedEnvelopeShape:
|
|||
assert modified is False
|
||||
assert json.loads(out) == json.loads(first_msg)
|
||||
|
||||
def test_bypass_header_short_circuits_first_frame(self):
|
||||
first_msg = json.dumps(
|
||||
{
|
||||
"type": "response.create",
|
||||
"response": {
|
||||
"model": "gpt-5",
|
||||
"input": [
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": "call_1",
|
||||
"output": json.dumps(
|
||||
[
|
||||
{
|
||||
"id": i,
|
||||
"name": f"Item {i}",
|
||||
"desc": "large repeated payload " * 20,
|
||||
}
|
||||
for i in range(100)
|
||||
]
|
||||
),
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
out, modified = _ws_compress_first_frame(first_msg, bypass=True)
|
||||
|
||||
assert modified is False
|
||||
assert out == first_msg
|
||||
|
||||
|
||||
class TestUnwrappedShape:
|
||||
"""Older Codex versions (and some test fixtures) send the Responses
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue