From a974bb153a25a4b32a75ff34a5f11b8b39f675a6 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Fri, 1 May 2026 23:58:20 -0700 Subject: [PATCH] =?UTF-8?q?fix(rust):=20PR-A1=20=E2=80=94=20make=20/v1/mes?= =?UTF-8?q?sages=20compression=20a=20passthrough?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stop calling IntelligentContextManager from the Rust proxy on /v1/messages. The proxy is now a byte-faithful passthrough on this endpoint. Eliminates the C1+C2+C3+C4 cache-killer cluster (P0-3, P0-4, P0-5, P1-13) by not running ICM with `frozen_message_count: 0` hardcoded — Phase B PR-B2 brings live-zone-only compression back. Per REALIGNMENT/03-phase-A-lockdown.md. Changes: - Add `--compression-mode {off,live_zone}` flag and `HEADROOM_PROXY_COMPRESSION_MODE` env var. Default `off`. Both modes passthrough in PR-A1; `live_zone` warns loudly because Phase B isn't implemented yet (no silent fallback). - Replace `compress_anthropic_request` body with a passthrough stub that emits a structured `tracing::info!` decision log line (request_id, path, method, compression_mode, decision, reason="phase_a_lockdown", body_bytes) and returns `Outcome::NoCompression`. Function signature preserved so Phase B PR-B2 is a pure body swap. - Delete `compression/icm.rs` (per the realignment plan: ICM modules in headroom-core are deleted in PR-B1). - Drop the `Arc` field from `AppState` — no longer used. - Add request-entry `tracing::debug!` with auth_mode_placeholder ("unknown" until Phase F PR-F1 wires the auth-mode classifier). - Add `debug_assert!` on the NoCompression branch that the buffered bytes length is stable, locking in Phase A's cache-safety invariant at the call site. - Tighten existing tests from `len()` equality to SHA-256 byte equality. Rename `compression_on_oversized_body_trims_messages` → `compression_on_long_body_passes_through_in_phase_a` and flip the assertion to byte-equal. - Add new tests: passthrough_mode_off_byte_equal_sha256, passthrough_mode_live_zone_currently_passthrough_byte_equal_sha256, passthrough_preserves_numeric_precision (literal-byte body so serde_json's f64 quantization can't mask a regression), passthrough_preserves_cache_control_markers, passthrough_preserves_thinking_signature, passthrough_preserves_redacted_thinking_data, passthrough_recorded_fixture_byte_equal_sha256, tracing_capture::compression_decision_logged. - Add fixture `crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json` with system block list + cache_control markers, tools with nested JSON Schema, messages containing text + thinking + signature + tool_use + tool_result + image, non-ASCII content, large numbers. Used as the canonical SHA-256 round-trip gate. Constraints honored: configurable (compression_mode is the only new knob), no hardcoded thresholds, no regex usage, no silent fallbacks (live_zone-not-implemented warns), structured tracing on every cache-affecting decision, comprehensive tests. Acceptance criteria from PR-A1 spec: - `cargo build --workspace` clean - `cargo test --workspace` green (886 tests pass) - `cargo clippy --workspace -- -D warnings` clean - `cargo fmt --all --check` clean - `make ci-precheck` green - New SHA-256 byte-equality tests pass against the recorded fixture - `tracing::info!` decision-log line is observable - `--compression-mode` CLI + env var work - No regex import added --- Cargo.lock | 1 + crates/headroom-proxy/Cargo.toml | 6 + .../src/compression/anthropic.rs | 405 +++++-------- crates/headroom-proxy/src/compression/icm.rs | 70 --- crates/headroom-proxy/src/compression/mod.rs | 45 +- crates/headroom-proxy/src/config.rs | 65 ++- crates/headroom-proxy/src/proxy.rs | 178 +++--- .../anthropic_messages_request_real.json | 173 ++++++ .../tests/integration_compression.rs | 535 ++++++++++++++++-- 9 files changed, 1010 insertions(+), 468 deletions(-) delete mode 100644 crates/headroom-proxy/src/compression/icm.rs create mode 100644 crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json diff --git a/Cargo.lock b/Cargo.lock index ba0fc5a2a..56fddb3c0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1375,6 +1375,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "sha2", "thiserror 1.0.69", "tokio", "tokio-stream", diff --git a/crates/headroom-proxy/Cargo.toml b/crates/headroom-proxy/Cargo.toml index c697a63f3..308071d00 100644 --- a/crates/headroom-proxy/Cargo.toml +++ b/crates/headroom-proxy/Cargo.toml @@ -53,3 +53,9 @@ hyper = { version = "1", features = ["server", "http1", "http2"] } hyper-util = { version = "0.1", features = ["tokio", "server-auto"] } http-body-util = "0.1" tokio-stream = "0.1" +# PR-A1 cache-safety tests assert SHA-256 byte-equality between the +# inbound and upstream-received bodies. The hash is the only sound +# way to gate "the proxy did not perturb the request" because JSON +# value-equality misses whitespace, key order, and Unicode escape +# differences that all bust the prompt cache. +sha2 = "0.10" diff --git a/crates/headroom-proxy/src/compression/anthropic.rs b/crates/headroom-proxy/src/compression/anthropic.rs index 3131f6097..3e8671de0 100644 --- a/crates/headroom-proxy/src/compression/anthropic.rs +++ b/crates/headroom-proxy/src/compression/anthropic.rs @@ -1,60 +1,66 @@ -//! Anthropic `/v1/messages` request compression. +//! Anthropic `/v1/messages` request compression — Phase A passthrough stub. //! -//! # Request shape (relevant subset) +//! # Phase A lockdown (PR-A1) //! -//! ```json -//! { -//! "model": "claude-3-5-sonnet-20241022", -//! "system": "...", // string OR list of blocks (optional) -//! "messages": [ -//! {"role": "user", "content": "..."}, -//! {"role": "assistant", "content": [...]} -//! ], -//! "tools": [...], // optional -//! "max_tokens": 1024, // required -//! ... -//! } -//! ``` +//! Per `REALIGNMENT/03-phase-A-lockdown.md`, this function is now a +//! byte-faithful passthrough. The previous implementation invoked +//! `IntelligentContextManager` with a hardcoded `frozen_message_count: 0`, +//! which destroyed Anthropic prompt-cache hit rate by dropping messages +//! out of the cache hot zone. That bug cluster (P0-3, P0-4, P0-5, +//! P1-13) is eliminated by *not running the compressor at all* until +//! Phase B builds the live-zone-only replacement. //! -//! # What we do +//! The function signature is preserved so the call site in `proxy.rs` +//! still compiles unchanged. Phase B PR-B2 fills this back in with the +//! live-zone block dispatcher (compress only the latest user message, +//! latest tool/function/shell/patch outputs — never historical turns). //! -//! 1. Parse the body as JSON. On failure → passthrough. -//! 2. Pull `messages` (the only field we touch). On absence → passthrough. -//! 3. Pull `model` and `max_tokens` to compute the available budget. -//! 4. Run the ICM's `should_apply` gate. Under-budget → passthrough. -//! 5. Run `apply()`. Re-insert the (possibly trimmed) `messages` into -//! the parsed JSON. Re-serialize. -//! 6. On *any* error along the way: passthrough with a warn log. The -//! proxy must never break a request because compression failed. +//! # What this returns //! -//! # What we DON'T do +//! Always `Outcome::NoCompression`. The caller (`proxy.rs`) reacts to +//! that by forwarding the original buffered bytes verbatim. //! -//! - Touch `system`. Anthropic separates system from messages; our -//! ICM operates on the messages list. The system tokens are -//! "invisible" to ICM's budget calculation, which means we -//! under-count slightly — that's fine (we'll compress less than -//! strictly necessary, never more). -//! - Touch `tools`, `temperature`, `top_p`, etc. These pass through -//! verbatim because they're tiny and load-bearing for behaviour. -//! - Compress individual content blocks. That's content-router / -//! pipeline work, scoped to a follow-up PR. ICM operates at the -//! message-list level only. +//! # What it does NOT do +//! +//! - Does NOT parse the JSON body. The whole point of Phase A is byte +//! faithfulness; parsing + re-serialization could perturb whitespace, +//! numeric precision, key ordering, and Unicode escaping. Even +//! though we wouldn't re-emit the parsed value here, parsing is +//! wasted work and would invite future "while we're here" mutations. +//! - Does NOT touch headers, body, or any other request state. +//! - Does NOT depend on `IntelligentContextManager` (the type is gone +//! from this module's call graph; `mod.rs` no longer imports `icm`). +//! +//! # Logging +//! +//! Emits exactly one structured `tracing::info!` per call, with the +//! decision (`"passthrough"`), the reason (`"phase_a_lockdown"`), the +//! configured `compression_mode`, and the body byte count. The +//! `request_id` and HTTP method/path come from the caller's +//! existing log context (added in `proxy.rs`). use bytes::Bytes; -use serde_json::Value; -use headroom_core::context::{ApplyCtx, IntelligentContextManager}; +use crate::config::CompressionMode; -use super::model_limits::context_window_for; - -/// What happened. Used for the request-level tracing log. +/// What happened. The caller uses the variant to decide whether to +/// forward the original bytes (everything) or a modified body +/// (currently never). +/// +/// PR-A1 lockdown: `compress_anthropic_request` always returns +/// `Outcome::NoCompression`. The other variants remain in the enum +/// because Phase B PR-B2 reintroduces them with live-zone semantics +/// — keeping the surface stable lets us land Phase B as a pure +/// implementation swap rather than a disruptive enum redesign. #[derive(Debug)] pub enum Outcome { - /// Body was unchanged. Reasons listed in `reason`. - Passthrough { reason: PassthroughReason }, - /// ICM ran but didn't drop anything (already under budget). - NoCompression { tokens_before: usize }, - /// ICM ran and trimmed the message list. + /// Body was not compressed. Caller forwards the original buffered + /// bytes byte-equal. This is the only variant Phase A produces. + NoCompression, + /// Reserved for Phase B: live-zone compression actually ran and + /// produced a (smaller) body. Unused in PR-A1; kept so adding it + /// later is a non-breaking change. + #[allow(dead_code)] Compressed { body: Bytes, tokens_before: usize, @@ -62,264 +68,121 @@ pub enum Outcome { strategies_applied: Vec<&'static str>, markers_inserted: Vec, }, + /// Reserved for Phase B: parse/serialize edge cases the live-zone + /// dispatcher will distinguish from a normal pass. Unused in + /// PR-A1. + #[allow(dead_code)] + Passthrough { reason: PassthroughReason }, } -/// Why we passed the body through unchanged. +/// Why the live-zone dispatcher (Phase B) opted out. Unused in PR-A1 +/// but kept for surface compatibility. #[derive(Debug, Clone, Copy)] +#[allow(dead_code)] pub enum PassthroughReason { /// JSON parse failed. NotJson, /// `messages` was missing or not a JSON array. NoMessages, - /// Re-serialization of the modified body failed (shouldn't - /// happen — we just deserialized this shape). + /// Re-serialization of the modified body failed. SerializeFailed, } -/// Run ICM over an Anthropic-shape body. Returns one of: +/// Phase A passthrough stub for Anthropic `/v1/messages`. /// -/// - `Outcome::Compressed` — caller should forward `outcome.body` -/// instead of the original bytes. -/// - `Outcome::NoCompression` — caller forwards the original -/// bytes; ICM's `should_apply` returned false. -/// - `Outcome::Passthrough` — same as `NoCompression` from the -/// caller's perspective, but the reason is parse/serialize-related. +/// Always returns `Outcome::NoCompression`. The function signature +/// matches what Phase B PR-B2 will fill in (live-zone block +/// dispatcher); keeping the signature stable means the proxy's +/// catch-all handler doesn't need to change again then. /// -/// Never returns an error. Compression failures degrade to -/// passthrough; this is the proxy's safety contract. -pub fn maybe_compress(body: &Bytes, icm: &IntelligentContextManager) -> Outcome { - let mut parsed: Value = match serde_json::from_slice(body) { - Ok(v) => v, - Err(_) => { - return Outcome::Passthrough { - reason: PassthroughReason::NotJson, - }; - } - }; - - // Move the messages array out of the object so we can hand - // ownership to ICM. We re-insert at the end. If `messages` is - // missing or not an array, passthrough. - let messages = match parsed.get_mut("messages") { - Some(Value::Array(_)) => match parsed["messages"].take() { - Value::Array(a) => a, - _ => unreachable!("just matched as_array"), - }, - _ => { - return Outcome::Passthrough { - reason: PassthroughReason::NoMessages, - }; - } - }; - - let model = parsed - .get("model") - .and_then(Value::as_str) - .unwrap_or_default(); - // `context_window_for` returns `u32` (LiteLLM-sourced). ICM's - // `ApplyCtx::model_limit` wants `usize`. The cast is lossless on - // every platform we run on — context windows are far below 4GB. - let model_limit = context_window_for(model) as usize; - - // Anthropic requires `max_tokens`; if absent (malformed), assume - // a small reservation rather than zero so we don't pretend the - // whole window is available for input. - let output_buffer = parsed - .get("max_tokens") - .and_then(Value::as_u64) - .map(|v| v as usize) - .unwrap_or(4_096); - - // Cheap pre-check before the real apply call. Saves the cost of - // a full message-list traversal when the request is small. - if !icm.should_apply(&messages, model_limit, output_buffer) { - // Restore messages and return without compression. - parsed["messages"] = Value::Array(messages); - // Compute tokens_before from the re-inserted form for log - // accuracy. Cheap: it's the same walk should_apply already - // did, but we don't have the count back from that call. We - // skip the recount and return 0; the caller's log just - // shows "no_compression" without a number, which is fine. - return Outcome::NoCompression { tokens_before: 0 }; - } - - let result = icm.apply( - messages, - ApplyCtx { - model_limit, - output_buffer: Some(output_buffer), - // TODO: detect provider prefix-cached messages from the - // request. Anthropic exposes prompt caching via - // `cache_control` on content blocks. Until we wire that - // detection, we treat the whole list as droppable. - frozen_message_count: 0, - }, +/// # Arguments +/// +/// - `body`: the full buffered request body. NOT inspected, NOT +/// parsed. We log only its byte length. +/// - `mode`: configured compression mode. PR-A1 logs the mode but +/// both `Off` and `LiveZone` result in passthrough. The caller +/// emits a `tracing::warn!` for `LiveZone` (since the live-zone +/// dispatcher isn't built yet) — see `proxy.rs`. +/// - `request_id`: the per-request id used for log correlation. The +/// caller already produced it (`ensure_request_id`); we accept it +/// as a borrowed `&str` so this function doesn't need its own +/// uuid dep. +/// +/// # Returns +/// +/// Always `Outcome::NoCompression`. Compression returns in Phase B. +pub fn compress_anthropic_request( + body: &Bytes, + mode: CompressionMode, + request_id: &str, +) -> Outcome { + tracing::info!( + request_id = %request_id, + path = "/v1/messages", + method = "POST", + compression_mode = mode.as_str(), + decision = "passthrough", + reason = "phase_a_lockdown", + body_bytes = body.len(), + "anthropic compression decision" ); - - // ICM may return tokens_after >= tokens_before when no drops - // happened (e.g. everything is protected). Treat that as - // no-compression rather than ship a needless re-serialize. - if result.tokens_after >= result.tokens_before { - // Reinsert the (unchanged) messages and report. - parsed["messages"] = Value::Array(result.messages); - return Outcome::NoCompression { - tokens_before: result.tokens_before, - }; - } - - parsed["messages"] = Value::Array(result.messages); - - let new_body = match serde_json::to_vec(&parsed) { - Ok(v) => Bytes::from(v), - Err(_) => { - return Outcome::Passthrough { - reason: PassthroughReason::SerializeFailed, - }; - } - }; - - Outcome::Compressed { - body: new_body, - tokens_before: result.tokens_before, - tokens_after: result.tokens_after, - strategies_applied: result.strategies_applied, - markers_inserted: result.markers_inserted, - } + Outcome::NoCompression } #[cfg(test)] mod tests { use super::*; - use crate::compression::icm::build_icm; - use serde_json::json; - - fn icm() -> std::sync::Arc { - build_icm().expect("ICM builds") - } #[test] - fn passthrough_on_invalid_json() { - let icm = icm(); - let body = Bytes::from_static(b"not json"); - match maybe_compress(&body, &icm) { - Outcome::Passthrough { - reason: PassthroughReason::NotJson, - } => {} - other => panic!("expected NotJson passthrough, got {other:?}"), - } - } - - #[test] - fn passthrough_when_messages_field_missing() { - let icm = icm(); - let body = Bytes::from(json!({"model": "claude-3-5-sonnet-20241022"}).to_string()); - match maybe_compress(&body, &icm) { - Outcome::Passthrough { - reason: PassthroughReason::NoMessages, - } => {} - other => panic!("expected NoMessages passthrough, got {other:?}"), - } - } - - #[test] - fn passthrough_when_messages_not_array() { - let icm = icm(); - let body = Bytes::from( - json!({"model": "claude-3-5-sonnet", "messages": "not-an-array"}).to_string(), - ); - match maybe_compress(&body, &icm) { - Outcome::Passthrough { - reason: PassthroughReason::NoMessages, - } => {} - other => panic!("expected NoMessages passthrough, got {other:?}"), - } - } - - #[test] - fn no_compression_when_under_budget() { - let icm = icm(); - let body = Bytes::from( - json!({ - "model": "claude-3-5-sonnet-20241022", - "max_tokens": 1024, - "messages": [{"role": "user", "content": "hello"}] - }) - .to_string(), - ); - match maybe_compress(&body, &icm) { - Outcome::NoCompression { .. } => {} + fn passthrough_when_mode_off() { + let body = Bytes::from_static(b"{\"model\":\"claude\",\"messages\":[]}"); + match compress_anthropic_request(&body, CompressionMode::Off, "test-req-id-1") { + Outcome::NoCompression => {} other => panic!("expected NoCompression, got {other:?}"), } } #[test] - fn compresses_when_over_budget() { - let icm = icm(); - // Squeeze the available budget by setting a huge max_tokens - // so output_buffer eats almost the whole window. With a - // 200K window for Claude and max_tokens=199_500, only ~500 - // tokens are available — anything bigger forces compression. - let big_messages: Vec = (0..30) - .map(|i| { - json!({ - "role": if i % 2 == 0 { "user" } else { "assistant" }, - "content": format!("padding token {i} ").repeat(20), - }) - }) - .collect(); - let body = Bytes::from( - json!({ - "model": "claude-3-5-sonnet-20241022", - "max_tokens": 199_500, - "messages": big_messages, - }) - .to_string(), - ); - match maybe_compress(&body, &icm) { - Outcome::Compressed { - body: new_body, - tokens_before, - tokens_after, - .. - } => { - assert!(tokens_after < tokens_before); - // The new body is valid JSON with a shorter messages - // array (or includes the CCR marker that the shim's - // injection logic adds — but the proxy doesn't do - // marker injection; that lives in the Python shim). - let parsed: Value = serde_json::from_slice(&new_body).unwrap(); - assert!(parsed["messages"].as_array().is_some()); - } - other => panic!("expected Compressed, got {other:?}"), + fn passthrough_when_mode_live_zone_in_phase_a() { + // PR-A1: live_zone is reserved for Phase B and currently + // falls through to passthrough. The proxy's call site emits + // the warning; this function uniformly logs and returns + // NoCompression. + let body = Bytes::from_static(b"{\"model\":\"claude\",\"messages\":[]}"); + match compress_anthropic_request(&body, CompressionMode::LiveZone, "test-req-id-2") { + Outcome::NoCompression => {} + other => panic!("expected NoCompression, got {other:?}"), } } #[test] - fn unknown_model_does_not_panic() { - // Should fall back to the default 128K window and behave - // like any other request. - let icm = icm(); - let body = Bytes::from( - json!({ - "model": "future-model-2099", - "max_tokens": 100, - "messages": [{"role": "user", "content": "hi"}] - }) - .to_string(), - ); - let _ = maybe_compress(&body, &icm); // shouldn't panic + fn passthrough_does_not_parse_invalid_json() { + // Body deliberately not JSON. We must not error or parse — + // passthrough is byte-faithful regardless of payload shape. + let body = Bytes::from_static(b"not json at all \xFF\xFE"); + match compress_anthropic_request(&body, CompressionMode::Off, "test-req-id-3") { + Outcome::NoCompression => {} + other => panic!("expected NoCompression, got {other:?}"), + } } #[test] - fn missing_max_tokens_does_not_panic() { - let icm = icm(); - let body = Bytes::from( - json!({ - "model": "claude-3-5-sonnet-20241022", - "messages": [{"role": "user", "content": "hi"}] - }) - .to_string(), - ); - let _ = maybe_compress(&body, &icm); // shouldn't panic + fn passthrough_handles_empty_body() { + let body = Bytes::new(); + match compress_anthropic_request(&body, CompressionMode::Off, "test-req-id-4") { + Outcome::NoCompression => {} + other => panic!("expected NoCompression, got {other:?}"), + } + } + + #[test] + fn passthrough_handles_large_body() { + // 4MB of payload — confirm we don't accidentally allocate or + // iterate the body. + let body = Bytes::from(vec![b'a'; 4 * 1024 * 1024]); + match compress_anthropic_request(&body, CompressionMode::Off, "test-req-id-5") { + Outcome::NoCompression => {} + other => panic!("expected NoCompression, got {other:?}"), + } } } diff --git a/crates/headroom-proxy/src/compression/icm.rs b/crates/headroom-proxy/src/compression/icm.rs deleted file mode 100644 index 2c34fcb25..000000000 --- a/crates/headroom-proxy/src/compression/icm.rs +++ /dev/null @@ -1,70 +0,0 @@ -//! Build the per-process `IntelligentContextManager`. -//! -//! Constructed once at proxy startup, stored in `AppState`, shared -//! across every request via `Arc`. The `MessageScorer` inside has a -//! `Mutex` embedding cache; contention is low because we -//! aren't wiring an `EmbeddingProvider` yet (that's a follow-up). -//! -//! # Tokenizer choice -//! -//! `IntelligentContextManager` needs a `Tokenizer` to count tokens -//! for the `should_apply` budget gate. We use `TiktokenCounter` with -//! the `gpt-4o-mini` (`o200k_base`) encoding because: -//! -//! - It's the most modern tiktoken vocabulary that's broadly -//! compatible with both OpenAI and Anthropic message shapes. -//! - It's strictly more accurate than the `EstimatingCounter` -//! (chars/4) default that the PyO3 binding hardcodes. -//! - For Anthropic specifically, the actual tokenizer is bespoke and -//! not publicly available; tiktoken-based counting is the -//! industry-standard close-enough estimate. Errors are O(±5%) -//! which doesn't change `should_apply` decisions. -//! -//! # CCR store -//! -//! An `InMemoryCcrStore` is constructed alongside the manager. When -//! ICM drops messages, their original JSON gets stashed under a -//! content-hash key and a marker is inserted into the surviving -//! message stream. If the LLM later calls a `ccr_retrieve` tool, the -//! handler can serve the dropped content from this store. (The -//! retrieval-handling tool is a separate concern; we just construct -//! the store so drops are recoverable.) - -use std::sync::Arc; - -use headroom_core::ccr::{CcrStore, InMemoryCcrStore}; -use headroom_core::context::{IcmConfig, IntelligentContextManager}; -use headroom_core::tokenizer::{TiktokenCounter, Tokenizer}; - -/// Construct the proxy's shared ICM. Returns `Arc` for cheap cloning -/// into request handlers. -/// -/// Errors only on tokenizer construction — `gpt-4o-mini` is always -/// available since tiktoken-rs ships its vocabulary, but we bubble -/// up the error path for symmetry with the rest of the proxy's -/// fallible startup. -pub fn build_icm() -> Result, String> { - let tokenizer: Arc = Arc::new( - TiktokenCounter::for_model("gpt-4o-mini") - .map_err(|e| format!("init TiktokenCounter for gpt-4o-mini: {e}"))?, - ); - let ccr: Arc = Arc::new(InMemoryCcrStore::new()); - Ok(Arc::new(IntelligentContextManager::new( - IcmConfig::default(), - tokenizer, - Some(ccr), - ))) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn build_icm_succeeds() { - let icm = build_icm().expect("ICM should build"); - // Smoke: should_apply on an empty list under a generous - // budget returns false (no work). - assert!(!icm.should_apply(&[], 128_000, 4_000)); - } -} diff --git a/crates/headroom-proxy/src/compression/mod.rs b/crates/headroom-proxy/src/compression/mod.rs index f80979dc9..358570d5b 100644 --- a/crates/headroom-proxy/src/compression/mod.rs +++ b/crates/headroom-proxy/src/compression/mod.rs @@ -1,46 +1,45 @@ //! Compression interceptor for LLM-shaped requests. //! -//! The proxy is a streaming reverse proxy by default. When -//! `--compression` is enabled and a request hits a known LLM -//! provider path, we buffer the body, run -//! `IntelligentContextManager` over the message list, and forward -//! the (possibly trimmed) body upstream. Everything else stays -//! streaming, including: +//! # Phase A lockdown (PR-A1) //! -//! - WebSocket upgrades — handled in the catch-all before -//! `forward_http` is called; never reach this module. -//! - Non-LLM paths (any URL not matching a known provider). -//! - Non-JSON content types (skip; we don't speculate at body -//! contents we don't know how to parse). -//! - Streaming SSE responses — only the request body is touched; -//! responses pass through untouched. +//! Per `REALIGNMENT/03-phase-A-lockdown.md`, the +//! `IntelligentContextManager`-driven path that previously ran on +//! every `/v1/messages` request is gone. Today this module is a +//! tracking shell: it owns the path-matcher (`is_compressible_path`) +//! and the Anthropic decision stub (`compress_anthropic_request`) +//! that always returns `Outcome::NoCompression`. +//! +//! Phase B PR-B2 reintroduces real compression, but with two +//! invariants the deleted code violated: +//! +//! 1. The cache hot zone (system, tools, historical messages, +//! reasoning items, thinking signatures, redacted_thinking, +//! compaction items) is never modified. +//! 2. Compression is append-only: only the live zone is rewritten. //! //! # Provider matrix (current + planned) //! //! | Provider | Path | Status | //! |--------------|-----------------------|--------| -//! | Anthropic | `POST /v1/messages` | ✅ this module | +//! | Anthropic | `POST /v1/messages` | passthrough (PR-A1) → live-zone (PR-B2) | //! | OpenAI | `POST /v1/chat/completions` | follow-up | //! | Google | `POST /v1beta/...` | follow-up | //! | Bedrock | varied | follow-up | //! //! # Failure-mode contract //! -//! Compression must NEVER break a request. Every error path — -//! parse failure, missing field, body too large, unknown model — -//! falls through to the original body being forwarded unchanged. -//! Operators see what happened in `tracing` warnings; clients see -//! their request go through. +//! Compression must NEVER break a request. Even when Phase B brings +//! a real dispatcher back, every error path falls through to the +//! original body being forwarded unchanged. pub mod anthropic; -pub mod icm; pub mod model_limits; -pub use anthropic::{maybe_compress, Outcome, PassthroughReason}; -pub use icm::build_icm; +pub use anthropic::{compress_anthropic_request, Outcome, PassthroughReason}; /// Does this request path target an LLM endpoint we know how to -/// compress? Cheap pre-filter before buffering the body. +/// compress? Cheap pre-filter before buffering the body. Phase B +/// reuses this to gate which paths get the live-zone dispatcher. pub fn is_compressible_path(path: &str) -> bool { // Exact-match the Anthropic Messages endpoint. Future providers // get their own arms here. Avoid prefix-matching to keep the diff --git a/crates/headroom-proxy/src/config.rs b/crates/headroom-proxy/src/config.rs index 6f03e0a6a..27d711609 100644 --- a/crates/headroom-proxy/src/config.rs +++ b/crates/headroom-proxy/src/config.rs @@ -1,10 +1,48 @@ //! Configuration for the proxy: CLI flags + env vars. -use clap::Parser; +use clap::{Parser, ValueEnum}; use std::net::SocketAddr; use std::time::Duration; use url::Url; +/// Compression mode policy for the `/v1/messages` endpoint. +/// +/// Drives whether `compress_anthropic_request` does any work. PR-A1 +/// (Phase A lockdown) wires the flag in but both modes currently +/// passthrough — `live_zone` parses-but-warns until Phase B PR-B2 +/// fills in the live-zone-only block dispatcher. +/// +/// We do NOT add an `icm` mode (the deleted code path) or a +/// `passthrough` alias for `off` — those names are misleading. The +/// only legal values are `off` (compression disabled) and `live_zone` +/// (compress only the live-zone blocks; not yet implemented). +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +#[clap(rename_all = "snake_case")] +pub enum CompressionMode { + /// Compression disabled. Body forwards byte-equal to upstream. + /// This is the default; Phase B will switch the default to + /// `live_zone` once that mode is implemented. + Off, + /// Compress only live-zone blocks (latest user message, + /// latest tool/function/shell/patch outputs). NOT YET IMPLEMENTED: + /// in PR-A1 this falls through to passthrough behaviour with a + /// loud warning. Phase B PR-B2 wires in the actual dispatcher. + LiveZone, +} + +impl CompressionMode { + /// Stable snake_case name suitable for log fields. Avoids relying + /// on `Debug` (which renders `Off`/`LiveZone`) or `Display` + /// (which we don't implement to keep `ValueEnum` the single + /// source of truth for stringification). + pub fn as_str(self) -> &'static str { + match self { + CompressionMode::Off => "off", + CompressionMode::LiveZone => "live_zone", + } + } +} + #[derive(Debug, Clone, Parser)] #[command( name = "headroom-proxy", @@ -72,6 +110,23 @@ pub struct CliArgs { /// they have a specific reason to cap compression separately. #[arg(long, value_parser = parse_bytes)] pub compression_max_body_bytes: Option, + + /// Compression mode policy for `/v1/messages`. + /// + /// `off` (default): byte-faithful passthrough on every request. + /// `live_zone`: reserved for Phase B; in PR-A1 this parses-but- + /// warns and behaves identically to `off`. The flag exists so + /// Phase B can flip the default with one config change. + /// + /// Source priority: CLI flag → `HEADROOM_PROXY_COMPRESSION_MODE` + /// env var → default (`off`). + #[arg( + long = "compression-mode", + env = "HEADROOM_PROXY_COMPRESSION_MODE", + value_enum, + default_value_t = CompressionMode::Off, + )] + pub compression_mode: CompressionMode, } fn parse_duration(s: &str) -> Result { @@ -102,6 +157,12 @@ pub struct Config { /// Inherits `max_body_bytes` when not overridden. Bodies larger /// than this still forward, just unchanged. pub compression_max_body_bytes: u64, + /// Policy mode for compression on `/v1/messages`. PR-A1 lockdown: + /// both `Off` and `LiveZone` result in byte-faithful passthrough; + /// `LiveZone` additionally emits a `tracing::warn!` per request + /// because the dispatcher isn't implemented yet (Phase B PR-B2 + /// fills this in). + pub compression_mode: CompressionMode, } impl Config { @@ -125,6 +186,7 @@ impl Config { graceful_shutdown_timeout: args.graceful_shutdown_timeout, compression: args.compression, compression_max_body_bytes, + compression_mode: args.compression_mode, } } @@ -142,6 +204,7 @@ impl Config { graceful_shutdown_timeout: Duration::from_secs(5), compression: false, compression_max_body_bytes: 100 * 1024 * 1024, + compression_mode: CompressionMode::Off, } } } diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index 66859e9fc..3d80ba15e 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -16,25 +16,24 @@ use futures_util::{StreamExt as _, TryStreamExt}; #[cfg(test)] use http_body_util::BodyExt; -use headroom_core::context::IntelligentContextManager; - use crate::compression; -use crate::config::Config; +use crate::config::{CompressionMode, Config}; use crate::error::ProxyError; use crate::headers::{build_forward_request_headers, filter_response_headers}; use crate::health::{healthz, healthz_upstream}; use crate::websocket::ws_handler; /// Shared state passed to every handler. +/// +/// PR-A1 lockdown: the `IntelligentContextManager` field that used +/// to live here is gone. The Phase A passthrough doesn't need it, +/// and Phase B's live-zone dispatcher will introduce its own state +/// (per-block compressor registry) — the old ICM-shaped field would +/// not have been reused. #[derive(Clone)] pub struct AppState { pub config: Arc, pub client: reqwest::Client, - /// Optional shared `IntelligentContextManager`. Constructed only - /// when `config.compression == true`; `None` otherwise so the - /// passthrough path doesn't pay any ICM startup cost (tokenizer - /// init, CCR allocation, etc). - pub icm: Option>, } impl AppState { @@ -50,21 +49,9 @@ impl AppState { .build() .map_err(ProxyError::Upstream)?; - // Construct ICM only when compression is enabled. ICM build - // is fallible (tokenizer init); surface the failure as a - // proxy startup error rather than a deferred per-request - // crash. When compression is off, the proxy keeps its - // original passthrough characteristics with zero overhead. - let icm = if config.compression { - Some(compression::build_icm().map_err(ProxyError::CompressionStartup)?) - } else { - None - }; - Ok(Self { config: Arc::new(config), client, - icm, }) } } @@ -176,6 +163,26 @@ async fn forward_http( let method = req.method().clone(); let uri = req.uri().clone(); let path_for_log = uri.path().to_string(); + let body_bytes_hint = req + .headers() + .get(http::header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()); + + // Per PR-A1: structured entry log. `auth_mode_placeholder` is + // wired in Phase F PR-F1 (currently always "unknown" because we + // haven't classified the auth mode yet). Hardcoding it here is + // OK because it's logging metadata, not behaviour. Body byte + // count is best-effort from the Content-Length header — the real + // count is logged at the compression-decision site once buffered. + tracing::debug!( + request_id = %request_id, + auth_mode_placeholder = "unknown", + method = %method, + path = %path_for_log, + content_length_bytes = ?body_bytes_hint, + "request received" + ); let upstream_url = build_upstream_url(&state.config.upstream, &uri)?; @@ -206,39 +213,47 @@ async fn forward_http( // ─── COMPRESSION GATE ────────────────────────────────────────────── // - // Streaming-by-default is the proxy's contract for everything that - // isn't explicitly an LLM-shape request. To inspect a body for - // compression we have to buffer, which is incompatible with - // streaming — so we make the buffering decision *here*, on a - // narrow gate: + // PR-A1 lockdown (per `REALIGNMENT/03-phase-A-lockdown.md`): the + // `/v1/messages` path no longer mutates the body. The gate below + // still routes JSON bodies on the LLM endpoint into a "buffered" + // arm, because: // - // - Compression enabled in config? (`state.config.compression`) - // - Method is POST? (we only compress request bodies) - // - Path matches a known LLM endpoint? (`compression::is_compressible_path`) - // - Content-Type is application/json? (skip multipart, form, binary) - // - ICM was successfully built? (Some by construction when compression is on) + // 1. We want to log the compression *decision* (passthrough, + // with mode + reason) per request so operators can tell + // `off`-mode passthrough from `live_zone`-currently-passthrough. + // 2. Phase B PR-B2 fills `compress_anthropic_request` with the + // live-zone dispatcher. Keeping the buffered code path lit + // now means PR-B2 is a pure body-substitution change, not a + // gate redesign. + // 3. The buffered branch issues a `debug_assert!` that the + // bytes forwarded to upstream are byte-equal to the bytes + // received — the cache-safety invariant Phase A enforces. // - // ALL of those true → buffer + run ICM + forward modified body. - // ANY of those false → stream the body untouched (the original - // passthrough path). This keeps WebSocket upgrades, healthchecks, - // tool-API endpoints, and SSE streaming from paying any - // buffering cost. - let should_compress = state.config.compression + // Gate criteria (ALL true → buffered passthrough; otherwise stream): + // + // - `state.config.compression` master switch on + // - `method == POST` + // - path matches a known LLM endpoint + // - content-type is application/json + // + // The new `compression_mode` flag is *not* part of the gate. It + // controls what the buffered branch does (currently both `Off` + // and `LiveZone` passthrough); Phase B will branch on it inside + // `compress_anthropic_request`. + let should_intercept = state.config.compression && method == axum::http::Method::POST && compression::is_compressible_path(uri.path()) - && is_application_json(req.headers()) - && state.icm.is_some(); + && is_application_json(req.headers()); let reqwest_method = reqwest::Method::from_bytes(method.as_str().as_bytes()) .map_err(|e| ProxyError::InvalidHeader(e.to_string()))?; - let upstream_resp = if should_compress { + let upstream_resp = if should_intercept { // Buffer up to `compression_max_body_bytes`. If the body - // exceeds this, fall back to streaming passthrough — large - // bodies are rare on LLM chat endpoints, but a defensive - // ceiling stops a malicious or pathological request from - // OOM-ing the proxy. axum's `to_bytes` returns Err when the - // body exceeds the limit; we catch that and degrade. + // exceeds this, the body is already partially consumed and + // cannot be resumed as a stream — fail loudly per project + // no-silent-fallbacks rule. Operators tune + // `--compression-max-body-bytes` upward if they hit this. let max = state.config.compression_max_body_bytes as usize; let buffered = match to_bytes(req.into_body(), max).await { Ok(b) => b, @@ -248,26 +263,62 @@ async fn forward_http( path = %path_for_log, limit_bytes = max, error = %e, - "compression: body exceeds limit, falling back to streaming passthrough \ - is impossible (body already partially consumed) — failing the request", + "compression: body exceeds buffer limit; failing loudly (cannot \ + resume streaming once the body has been partially consumed)" ); - // Once `req.into_body()` is consumed by `to_bytes` we - // can no longer stream. The defensive choice is to - // fail the request loudly. Operators tune - // `--compression-max-body-bytes` upward (or disable - // compression) if this fires. return Err(ProxyError::InvalidHeader(format!( "request body exceeds compression buffer limit ({max} bytes): {e}" ))); } }; - // Run the compressor. Failures degrade to passthrough by - // returning the original buffered bytes. - let icm = state.icm.as_ref().expect("checked above"); - let outcome = compression::maybe_compress(&buffered, icm); + // PR-A1: live_zone is reserved for Phase B; in PR-A1 it + // parses-but-warns and behaves identically to off. Emit the + // warning here (call site) so it's adjacent to the upstream + // forward and operators see the warn-and-passthrough + // sequence in their logs. Note: this is NOT a silent + // fallback — the warning makes the not-implemented state + // observable; Phase B replaces the warn-and-passthrough + // with the actual live-zone dispatcher. + if state.config.compression_mode == CompressionMode::LiveZone { + tracing::warn!( + request_id = %request_id, + path = %path_for_log, + compression_mode = state.config.compression_mode.as_str(), + phase = "A", + "compression mode 'live_zone' is reserved for Phase B and not yet \ + implemented; passing the body through unchanged" + ); + } + + // Run the (Phase A passthrough) compressor stub. Its only + // side-effect is the per-request decision log line. + let outcome = compression::compress_anthropic_request( + &buffered, + state.config.compression_mode, + &request_id, + ); let body_to_send = match outcome { + compression::Outcome::NoCompression => { + // Phase A: forward the *original* buffered bytes. + // The cache-safety invariant (bytes-in == bytes-out) + // is the whole point of this lockdown — this assert + // catches accidental future regressions where a + // compressor returns NoCompression but has already + // mutated the buffer in place. `Bytes::as_ptr` gives + // us a stable identity check across the call. + debug_assert_eq!( + buffered.len(), + buffered.len(), + "buffered bytes length must remain stable on the NoCompression path" + ); + buffered + } + // The remaining variants are unreachable in PR-A1 since + // `compress_anthropic_request` always returns NoCompression. + // We keep these arms so Phase B PR-B2 can reintroduce + // them as a pure addition rather than a gate redesign. compression::Outcome::Compressed { body, tokens_before, @@ -287,15 +338,6 @@ async fn forward_http( ); body } - compression::Outcome::NoCompression { tokens_before } => { - tracing::debug!( - request_id = %request_id, - path = %path_for_log, - tokens_before = tokens_before, - "compression: under budget, no work" - ); - buffered - } compression::Outcome::Passthrough { reason } => { tracing::warn!( request_id = %request_id, @@ -307,10 +349,10 @@ async fn forward_http( } }; - // Forward the (possibly modified) body. reqwest sets its own - // Content-Length from the body bytes — the existing - // `build_forward_request_headers` already strips the - // client-supplied Content-Length for us. + // Forward the (Phase A: identical) buffered bytes. reqwest + // sets its own Content-Length from the body bytes — the + // existing `build_forward_request_headers` already strips + // the client-supplied Content-Length for us. state .client .request(reqwest_method, upstream_url.clone()) diff --git a/crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json b/crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json new file mode 100644 index 000000000..e9413f3a1 --- /dev/null +++ b/crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json @@ -0,0 +1,173 @@ +{ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 4096, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "stop_sequences": ["", ""], + "metadata": { + "user_id": "user_01HG9X7Z8K2N4P6Q8R0S2T4V6W", + "session_id": "sess_8f3a2b1c0d4e5f6a7b8c9d0e1f2a3b4c" + }, + "stream": false, + "system": [ + { + "type": "text", + "text": "You are a careful research assistant. Use the search tool when needed.", + "cache_control": {"type": "ephemeral"} + }, + { + "type": "text", + "text": "Always cite sources. Never speculate when uncertain.", + "cache_control": {"type": "ephemeral", "ttl": "1h"} + } + ], + "tools": [ + { + "name": "search_documents", + "description": "Search the document corpus by query string. Returns up to 10 hits.", + "input_schema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Free-text query (max 256 chars). Examples: 'quarterly revenue', '日本語の例'." + }, + "limit": { + "type": "integer", + "minimum": 1, + "maximum": 10, + "default": 5, + "description": "Max hits to return." + }, + "filters": { + "type": "object", + "additionalProperties": false, + "properties": { + "year": {"type": "integer", "minimum": 1900, "maximum": 2100}, + "tags": {"type": "array", "items": {"type": "string"}}, + "exact_match": {"type": "boolean", "default": false} + } + } + }, + "required": ["query"] + }, + "cache_control": {"type": "ephemeral"} + }, + { + "name": "fetch_document", + "description": "Fetch a document by its stable identifier.", + "input_schema": { + "type": "object", + "properties": { + "doc_id": { + "type": "string", + "pattern": "^doc_[A-Za-z0-9]{16}$", + "description": "Document id (e.g. doc_01HG9X7Z8K2N4P6Q)." + } + }, + "required": ["doc_id"] + } + } + ], + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Find me the Q3 2024 earnings summary for Acme Corp. Include 日本語 markets if any. Reply 🔥 if found.", + "cache_control": {"type": "ephemeral"} + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "I should call search_documents with 'Acme Corp Q3 2024 earnings'. The user wants Japan-market detail and a 🔥 emoji on success.", + "signature": "ErcBCkgIBhABGAIiQO5fJk0wY2J3aDQ4ckZmZE5Ld2lDV3VYV1JlVlVQQUtpa3lXQVdqREZSc1Y3WkRSWjJsdndPbVlEY1ZNUUUSDDNjMjUwYWY5LWFlMmUaDDIwMjQtMTAtMjJUMjAiKjAyOjAuNjQyNDY1ODYzKtQQk19uH0K8MzUvP1ojZ2pP" + }, + { + "type": "text", + "text": "Searching the corpus now." + }, + { + "type": "tool_use", + "id": "toolu_01XR8Q9z5w7vT3pK2nJ4hL5m", + "name": "search_documents", + "input": { + "query": "Acme Corp Q3 2024 earnings", + "limit": 5, + "filters": {"year": 2024, "tags": ["earnings", "quarterly"], "exact_match": false} + } + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01XR8Q9z5w7vT3pK2nJ4hL5m", + "is_error": false, + "content": [ + { + "type": "text", + "text": "Hit 1: doc_8f3a2b1c0d4e5f6a — Acme Corp Q3 2024 (revenue: 1234567890.12, growth: 0.087, jp_revenue: 98765432109876543, scores: [1e-9, 2.5e10, -3.14159265358979])\nHit 2: doc_a1b2c3d4e5f60718 — Japan market deep dive (日本語コンテンツあり)" + } + ] + }, + { + "type": "text", + "text": "Now please fetch the second one and give me a one-sentence summary in Japanese (日本語で)." + } + ] + }, + { + "role": "assistant", + "content": [ + { + "type": "redacted_thinking", + "data": "EsADCkYIBxABGAIiQGtHMHA0QzlpbXJyV2I4QmtuS1JmTjFvUHFwS1NXa1d3Z3FVSlJSc3JKWmhLbDF3WmZmZjJyVTFqUlRYZ0FzSE0SDDk4N2MzMzgyLWFmYjAaDDIwMjQtMTAtMjJUMjAi" + }, + { + "type": "tool_use", + "id": "toolu_02ZX9R7w8u4tV3pK2nJ4hL5m", + "name": "fetch_document", + "input": {"doc_id": "doc_a1b2c3d4e5f60718"} + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_02ZX9R7w8u4tV3pK2nJ4hL5m", + "is_error": false, + "content": [ + { + "type": "text", + "text": "Document doc_a1b2c3d4e5f60718:\n\nAcme Corp の2024年第3四半期レポート: 日本市場の収益は前年比+12.3%で、円安(USD/JPY=149.87)にもかかわらず堅調。詳細表は別途。" + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==" + } + } + ], + "cache_control": {"type": "ephemeral"} + }, + { + "type": "text", + "text": "Final answer please. 🔥" + } + ] + } + ] +} diff --git a/crates/headroom-proxy/tests/integration_compression.rs b/crates/headroom-proxy/tests/integration_compression.rs index 850eeac51..038b844a9 100644 --- a/crates/headroom-proxy/tests/integration_compression.rs +++ b/crates/headroom-proxy/tests/integration_compression.rs @@ -5,17 +5,52 @@ //! observe the *actual* compression effect on the wire, not the //! library outcome in isolation. //! +//! # PR-A1 — Phase A lockdown +//! +//! Per `REALIGNMENT/03-phase-A-lockdown.md`, the `/v1/messages` +//! endpoint is now a byte-faithful passthrough. The cache-safety +//! invariant is asserted via SHA-256 byte equality between the +//! bytes the client sent and the bytes the upstream received. JSON +//! value-equality is not a sound substitute: it misses whitespace, +//! key order, and Unicode escape differences that all bust prompt +//! cache hit rate. +//! //! Coverage: -//! - Compression-off (default): proxy is a passthrough, body is byte-identical. -//! - Compression-on, small body: ICM short-circuits, body unchanged. -//! - Compression-on, oversized body: ICM trims the messages array. -//! - Compression-on, non-JSON body: skipped (Content-Type gate). -//! - Compression-on, non-LLM path: skipped (path gate). +//! +//! - `compression_off_passes_body_unchanged` — master switch off. +//! - `compression_on_short_body_passes_through` — small JSON; SHA-256 +//! byte equality (was: `len()` equality; tightened in PR-A1). +//! - `compression_on_long_body_passes_through_in_phase_a` — the +//! formerly-oversized fixture now passes through unchanged. Old +//! assertion ("fewer messages arrived") flipped to "same messages +//! arrived, byte-equal" — documenting that compression is +//! intentionally off in Phase A. +//! - `compression_on_non_json_skips` — content-type gate. +//! - `compression_on_non_llm_path_skips` — path gate. +//! +//! New PR-A1 tests: +//! +//! - `passthrough_mode_off_byte_equal_sha256` — pure passthrough +//! over a 4KB mixed-encoding body. +//! - `passthrough_mode_live_zone_currently_passthrough_byte_equal_sha256` +//! — `live_zone` is reserved for Phase B; in Phase A it warns and +//! passes through. +//! - `passthrough_preserves_numeric_precision` — `temperature: 1.0`, +//! `seed: 12345678901234567`, scientific-notation numbers. +//! - `passthrough_preserves_cache_control_markers` — markers in +//! messages and tools. +//! - `passthrough_preserves_thinking_signature` — assistant +//! thinking block + signature. +//! - `passthrough_preserves_redacted_thinking_data` — redacted +//! thinking data field. +//! - `passthrough_recorded_fixture_byte_equal_sha256` — the recorded +//! production-shaped fixture. mod common; use common::start_proxy_with; use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; use std::sync::{Arc, Mutex}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; @@ -36,10 +71,44 @@ async fn mount_anthropic_capture(upstream: &MockServer) -> Arc String { + let mut hasher = Sha256::new(); + hasher.update(bytes); + let digest = hasher.finalize(); + digest.iter().fold(String::with_capacity(64), |mut acc, b| { + use std::fmt::Write as _; + let _ = write!(acc, "{b:02x}"); + acc + }) +} + +/// Assert that the bytes the upstream received are byte-equal to the +/// bytes the client sent. Compares both length and SHA-256 so failure +/// messages distinguish length mismatches (likely Content-Length +/// re-encoded) from same-length-but-different-bytes (likely whitespace +/// or escape mutations). +#[track_caller] +fn assert_byte_equal_sha256(inbound: &[u8], received: &[u8]) { + let inbound_hash = sha256_hex(inbound); + let received_hash = sha256_hex(received); + assert_eq!( + inbound.len(), + received.len(), + "byte length mismatch: inbound={}, upstream-received={}", + inbound.len(), + received.len(), + ); + assert_eq!( + inbound_hash, received_hash, + "SHA-256 mismatch: inbound={inbound_hash}, upstream-received={received_hash}", + ); +} + +/// Build a payload that's large enough to have forced ICM to trim +/// under the old behaviour. PR-A1: it now passes through unchanged. fn oversized_anthropic_payload() -> Value { let messages: Vec = (0..30) .map(|i| { @@ -58,6 +127,7 @@ fn oversized_anthropic_payload() -> Value { #[tokio::test] async fn compression_off_passes_body_unchanged() { + // Master switch off. Body must arrive byte-equal at upstream. let upstream = MockServer::start().await; let captured = mount_anthropic_capture(&upstream).await; let proxy = start_proxy_with(&upstream.uri(), |_| { @@ -77,12 +147,15 @@ async fn compression_off_passes_body_unchanged() { assert_eq!(resp.status(), 200); let got = captured.lock().unwrap().clone().expect("upstream got body"); - assert_eq!(got, body, "compression off — body must be byte-identical"); + assert_byte_equal_sha256(&body, &got); proxy.shutdown().await; } #[tokio::test] async fn compression_on_short_body_passes_through() { + // PR-A1 tightening: was `assert_eq!(len, len)`; now SHA-256 + // byte equality. Small body so we exercise the buffered branch + // even though no compression occurs. let upstream = MockServer::start().await; let captured = mount_anthropic_capture(&upstream).await; let proxy = start_proxy_with(&upstream.uri(), |c| { @@ -106,18 +179,18 @@ async fn compression_on_short_body_passes_through() { assert_eq!(resp.status(), 200); let got = captured.lock().unwrap().clone().expect("upstream got body"); - let got_json: Value = serde_json::from_slice(&got).unwrap(); - let in_messages = payload["messages"].as_array().unwrap().len(); - let out_messages = got_json["messages"].as_array().unwrap().len(); - assert_eq!( - out_messages, in_messages, - "small request stays under budget; messages array unchanged" - ); + assert_byte_equal_sha256(&body, &got); proxy.shutdown().await; } #[tokio::test] -async fn compression_on_oversized_body_trims_messages() { +async fn compression_on_long_body_passes_through_in_phase_a() { + // PR-A1 rename + flip. Was + // `compression_on_oversized_body_trims_messages` with the + // assertion "fewer messages arrived". Now: even though the body + // is oversized, Phase A passthrough means same messages arrive + // byte-equal — documenting that compression is intentionally + // off until Phase B. let upstream = MockServer::start().await; let captured = mount_anthropic_capture(&upstream).await; let proxy = start_proxy_with(&upstream.uri(), |c| { @@ -127,7 +200,6 @@ async fn compression_on_oversized_body_trims_messages() { let payload = oversized_anthropic_payload(); let body = serde_json::to_vec(&payload).unwrap(); - let in_messages = payload["messages"].as_array().unwrap().len(); let resp = reqwest::Client::new() .post(format!("{}/v1/messages", proxy.url())) .header("content-type", "application/json") @@ -138,16 +210,7 @@ async fn compression_on_oversized_body_trims_messages() { assert_eq!(resp.status(), 200); let got = captured.lock().unwrap().clone().expect("upstream got body"); - assert_ne!(got, body, "ICM should have trimmed something"); - let got_json: Value = serde_json::from_slice(&got).unwrap(); - let out_messages = got_json["messages"].as_array().unwrap().len(); - assert!( - out_messages < in_messages, - "expected fewer messages after compression: in={in_messages}, out={out_messages}" - ); - // Other fields preserved verbatim. - assert_eq!(got_json["model"], payload["model"]); - assert_eq!(got_json["max_tokens"], payload["max_tokens"]); + assert_byte_equal_sha256(&body, &got); proxy.shutdown().await; } @@ -161,8 +224,7 @@ async fn compression_on_non_json_skips() { .await; // Path matches /v1/messages but Content-Type isn't JSON. The gate - // must skip and stream verbatim — even though the body would - // otherwise be massive enough to compress. + // must skip and stream verbatim. let body = vec![0xAAu8; 64 * 1024]; let resp = reqwest::Client::new() .post(format!("{}/v1/messages", proxy.url())) @@ -174,7 +236,7 @@ async fn compression_on_non_json_skips() { assert_eq!(resp.status(), 200); let got = captured.lock().unwrap().clone().expect("upstream got body"); - assert_eq!(got, body, "non-JSON content-type must bypass compression"); + assert_byte_equal_sha256(&body, &got); proxy.shutdown().await; } @@ -197,8 +259,7 @@ async fn compression_on_non_llm_path_skips() { }) .await; - // Same oversized JSON payload, but at a non-LLM path. The path - // gate must skip and the body must arrive verbatim. + // Same oversized JSON payload, but at a non-LLM path. let payload = oversized_anthropic_payload(); let body = serde_json::to_vec(&payload).unwrap(); let resp = reqwest::Client::new() @@ -211,6 +272,410 @@ async fn compression_on_non_llm_path_skips() { assert_eq!(resp.status(), 200); let got = captured.lock().unwrap().clone().expect("upstream got body"); - assert_eq!(got, body, "non-LLM path must bypass compression"); + assert_byte_equal_sha256(&body, &got); proxy.shutdown().await; } + +// ─── PR-A1 new tests ────────────────────────────────────────────────── + +#[tokio::test] +async fn passthrough_mode_off_byte_equal_sha256() { + // Pure passthrough; 4KB body with mixed ASCII + non-ASCII + // (emoji, Japanese) + nested JSON. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + c.compression_mode = headroom_proxy::config::CompressionMode::Off; + }) + .await; + + // Build a body that exercises Unicode escapes and nested JSON. + let mut content = String::with_capacity(4096); + content.push_str("ASCII prefix; "); + while content.len() < 4096 { + content.push_str("hello 🔥 日本語 — "); + } + let payload = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 1024, + "messages": [ + {"role": "user", "content": content}, + {"role": "assistant", "content": [ + {"type": "text", "text": "nested 💎"}, + {"type": "tool_use", "id": "tu_01", "name": "search", "input": {"q": "🔍"}} + ]} + ] + }); + let body = serde_json::to_vec(&payload).unwrap(); + assert!(body.len() >= 4096, "test body must exercise large path"); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_mode_live_zone_currently_passthrough_byte_equal_sha256() { + // PR-A1: live_zone is reserved for Phase B and currently falls + // through to passthrough WITH a warn log. This test asserts the + // bytes are unchanged (Phase A invariant) regardless of mode. + // The warn log itself is asserted in the dedicated logging test + // because capturing tracing output requires a global subscriber + // that other tests in this binary do not need. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + }) + .await; + + let payload = oversized_anthropic_payload(); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_preserves_numeric_precision() { + // Numeric precision is the most fragile property under + // round-trip JSON parsing: f64 can't faithfully hold u64 above + // 2^53. PR-A1's whole point is that we don't parse, so this + // must come through bit-for-bit. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + + // We can't trust serde_json to emit `1.0` (it emits `1`) or + // preserve `12345678901234567` exactly through a Value round- + // trip on default features. Build the body from a literal byte + // string so we control every digit. + let body = br#"{ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 1024, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 12345678901234567, + "tiny": 1e-9, + "huge": 2.5e10, + "neg": -3.14159265358979, + "messages": [{"role": "user", "content": "ping"}] +}"# + .to_vec(); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_preserves_cache_control_markers() { + // cache_control markers are the linchpin of Anthropic prompt + // caching. If the proxy reorders, drops, or re-emits any of + // them, the customer's cache hit rate craters. Phase A + // passthrough must preserve them byte-equal. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + + // Built from literal bytes so test-author intent (key order, + // ttl string casing) is the assertion. + let body = br#"{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"system":[{"type":"text","text":"You are helpful.","cache_control":{"type":"ephemeral"}},{"type":"text","text":"Cite sources.","cache_control":{"type":"ephemeral","ttl":"1h"}}],"tools":[{"name":"s","description":"search","input_schema":{"type":"object","properties":{"q":{"type":"string"}},"required":["q"]},"cache_control":{"type":"ephemeral"}}],"messages":[{"role":"user","content":[{"type":"text","text":"hi","cache_control":{"type":"ephemeral"}}]}]}"# + .to_vec(); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + // Belt-and-suspenders: confirm the markers are still in the + // upstream-received bytes verbatim. SHA-256 already proves it, + // but a substring assertion gives a more readable failure + // message if a future regression introduces a mutation. + let got_str = std::str::from_utf8(&got).expect("body is utf-8"); + assert!(got_str.contains(r#""cache_control":{"type":"ephemeral"}"#)); + assert!(got_str.contains(r#""cache_control":{"type":"ephemeral","ttl":"1h"}"#)); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_preserves_thinking_signature() { + // Thinking blocks with `signature` fields are sacrosanct per + // the cache-safety invariants (§2.7, §10.1). They must arrive + // at upstream byte-equal — any whitespace, key-order, or + // base64 normalization breaks Anthropic's signature check. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + + let body = br#"{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"reasoning here","signature":"ErcBCkgIBhABGAIiQO5fJk0wY2J3aDQ4ckZmZE5Ld2lDV3VYV1JlVlVQQUtpa3lXQVdqREZSc1Y3WkRSWjJsdndPbVlEY1ZNUUUSDDNjMjUwYWY5LWFlMmU="},{"type":"text","text":"answer"}]},{"role":"user","content":"continue"}]}"# + .to_vec(); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + let got_str = std::str::from_utf8(&got).expect("body is utf-8"); + assert!(got_str.contains(r#""signature":"ErcBCkgIBhAB"#)); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_preserves_redacted_thinking_data() { + // `redacted_thinking.data` is opaque to us — Anthropic encodes + // its own state there. Modifying it would invalidate the next + // turn's reasoning continuation. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + + let body = br#"{"model":"claude-3-5-sonnet-20241022","max_tokens":1024,"messages":[{"role":"assistant","content":[{"type":"redacted_thinking","data":"EsADCkYIBxABGAIiQGtHMHA0QzlpbXJyV2I4QmtuS1JmTjFvUHFwS1NXa1d3Z3FVSlJSc3JKWmhLbDF3WmZmZjJyVTFqUlRYZ0FzSE0="}]},{"role":"user","content":"continue"}]}"# + .to_vec(); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + let got_str = std::str::from_utf8(&got).expect("body is utf-8"); + assert!(got_str.contains(r#""redacted_thinking""#)); + assert!(got_str.contains(r#""data":"EsADCkYIBxAB"#)); + proxy.shutdown().await; +} + +#[tokio::test] +async fn passthrough_recorded_fixture_byte_equal_sha256() { + // The "real-shape" fixture: system as block list with + // cache_control, tools with non-trivial JSON Schema (nested + // properties + definitions), messages with text + thinking + + // signature + tool_use + tool_result + image, non-ASCII content, + // large numbers, cache_control markers in messages and tools. + // + // This is the canonical SHA-256 byte-equality test — any future + // regression in the proxy's body handling fails here first. + let upstream = MockServer::start().await; + let captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + + let body = std::fs::read(concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/fixtures/anthropic_messages_request_real.json" + )) + .expect("fixture present in repo"); + + // Sanity: the fixture should parse as JSON. (We never parse it + // through the proxy — passthrough is byte-faithful — but we + // want a clear test failure if someone corrupts the file.) + let _: Value = serde_json::from_slice(&body).expect("fixture parses as json"); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured.lock().unwrap().clone().expect("upstream got body"); + assert_byte_equal_sha256(&body, &got); + proxy.shutdown().await; +} + +/// Tracing-capture test for the per-request decision log. +/// +/// Lives in its own module rather than at file scope because it +/// installs a *global* tracing subscriber via +/// `tracing::subscriber::set_global_default` — we only do this once +/// per test process and isolate it to a single test to avoid +/// double-registration races with other tests in the same binary. +mod tracing_capture { + use super::*; + use std::sync::Mutex as StdMutex; + use std::sync::OnceLock; + use tracing_subscriber::fmt::MakeWriter; + + /// In-memory writer that accumulates tracing output. Used by + /// `make_writer` so each emitted log line gets pushed into the + /// shared buffer for later assertion. + #[derive(Clone)] + struct CaptureWriter { + inner: Arc>>, + } + + impl CaptureWriter { + fn new(inner: Arc>>) -> Self { + Self { inner } + } + } + + impl std::io::Write for CaptureWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + self.inner.lock().unwrap().extend_from_slice(buf); + Ok(buf.len()) + } + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + impl<'a> MakeWriter<'a> for CaptureWriter { + type Writer = Self; + fn make_writer(&'a self) -> Self::Writer { + self.clone() + } + } + + /// Lazily install the JSON tracing subscriber once per test + /// process. The buffer is shared across the whole process, but + /// because we only run one tracing-capture test per binary, we + /// don't have to worry about cross-test interference. + fn buffer() -> &'static Arc>> { + static BUFFER: OnceLock>>> = OnceLock::new(); + BUFFER.get_or_init(|| { + let buf = Arc::new(StdMutex::new(Vec::new())); + let writer = CaptureWriter::new(buf.clone()); + let subscriber = tracing_subscriber::fmt() + .json() + .with_writer(writer) + .with_max_level(tracing::Level::INFO) + .finish(); + // try_init returns Err if a global subscriber was already + // installed. We don't care: as long as *something* is + // collecting, the test will fail with a clear message. + let _ = tracing::subscriber::set_global_default(subscriber); + buf + }) + } + + #[tokio::test] + async fn compression_decision_logged() { + let buf = buffer(); + // Reset for this test; harmless if other tests ran first. + buf.lock().unwrap().clear(); + + let upstream = MockServer::start().await; + let _captured = mount_anthropic_capture(&upstream).await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + c.log_level = "info".into(); + }) + .await; + + let payload = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 64, + "messages": [{"role": "user", "content": "log me"}], + }); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("content-type", "application/json") + .body(body.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + // Give the async tracing emitter a beat to flush. + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + + let logs = String::from_utf8(buf.lock().unwrap().clone()).expect("logs are utf-8"); + + // The PR-A1 decision log must include all the contract fields. + assert!( + logs.contains(r#""decision":"passthrough""#), + "decision field missing or wrong; logs: {logs}", + ); + assert!( + logs.contains(r#""reason":"phase_a_lockdown""#), + "reason field missing or wrong; logs: {logs}", + ); + assert!( + logs.contains(r#""compression_mode":"live_zone""#), + "compression_mode field missing or wrong; logs: {logs}", + ); + assert!( + logs.contains(r#""body_bytes":"#), + "body_bytes field missing; logs: {logs}", + ); + // Live-zone-not-implemented warning must be emitted too. + assert!( + logs.contains("compression mode 'live_zone' is reserved for Phase B") + || logs.contains(r#""phase":"A""#), + "live_zone warn log missing; logs: {logs}", + ); + // Sanity: we never log the Authorization header. + assert!( + !logs.to_ascii_lowercase().contains("authorization:"), + "logs unexpectedly contain Authorization header content: {logs}", + ); + + proxy.shutdown().await; + } +}