akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(rust): PR-A1 — make /v1/messages compression a passthrough

Browse source 
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<IntelligentContextManager>` 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
This commit is contained in:
chopratejas 2026-05-01 23:58:20 -07:00
parent 0ce2243dfb
commit a974bb153a
9 changed files with 1010 additions and 468 deletions
Download patch file Download diff file Expand all files Collapse all files

1
Cargo.lock generated
View file

@ -1375,6 +1375,7 @@ dependencies = [
"reqwest",
"serde",
"serde_json",
"sha2",
"thiserror 1.0.69",
"tokio",
"tokio-stream",

6
crates/headroom-proxy/Cargo.toml
View file

@ -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"

405
crates/headroom-proxy/src/compression/anthropic.rs
View file

@ -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<String>,
},
/// 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<IntelligentContextManager> {
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<Value> = (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:?}"),
}
}
}

70
crates/headroom-proxy/src/compression/icm.rs
View file

@ -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<HashMap>` 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<Arc<IntelligentContextManager>, String> {
let tokenizer: Arc<dyn Tokenizer> = Arc::new(
TiktokenCounter::for_model("gpt-4o-mini")
.map_err(|e| format!("init TiktokenCounter for gpt-4o-mini: {e}"))?,
);
let ccr: Arc<dyn CcrStore> = 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));
}
}

45
crates/headroom-proxy/src/compression/mod.rs
View file

@ -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

65
crates/headroom-proxy/src/config.rs
View file

@ -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<u64>,
/// 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<Duration, String> {
@ -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,
}
}
}

178
crates/headroom-proxy/src/proxy.rs
View file

@ -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<Config>,
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<Arc<IntelligentContextManager>>,
}
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::<u64>().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())

173
crates/headroom-proxy/tests/fixtures/anthropic_messages_request_real.json vendored Normal file
View file

@ -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": ["</search>", "</done>"],
"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. 🔥"
}
]
}
]
}

535
crates/headroom-proxy/tests/integration_compression.rs
View file

@ -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<Mutex<Option<Vec<
captured
}
/// Build a payload that's large enough to force ICM to trim. Uses the
/// same pattern as the `compresses_when_over_budget` unit test in the
/// anthropic module: huge max_tokens eats the budget, leaving very few
/// tokens for input and forcing drops.
/// Compute the lowercase hex SHA-256 of a byte slice. Used to gate
/// "the proxy did not perturb the request body" — the only sound way
/// to assert byte-faithfulness.
fn sha256_hex(bytes: &[u8]) -> 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<Value> = (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<StdMutex<Vec<u8>>>,
}
impl CaptureWriter {
fn new(inner: Arc<StdMutex<Vec<u8>>>) -> Self {
Self { inner }
}
}
impl std::io::Write for CaptureWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
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<StdMutex<Vec<u8>>> {
static BUFFER: OnceLock<Arc<StdMutex<Vec<u8>>>> = 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;
}
}