From 4a3b76bcc879b3ba4a6a29096b5b50cb87d34f40 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Mon, 4 May 2026 15:27:54 -0700 Subject: [PATCH] fix: PR-E1 tool array deterministic sort (Phase E) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sort `tools[]` alphabetically by name on the way out so cache hits no longer depend on the customer-side iteration order (commonly hash- randomized via `set()` / `dict`). Mutates request bytes only when: 1. Auth mode is PAYG (`headroom_core::auth_mode::classify`). 2. No tool already carries a `cache_control` marker (reordering would shift cache scope and silently void customer intent). Every gate skip emits a structured `e1_skipped` event with `reason = auth_mode | marker_present` so dashboards can see policy adoption. Wired into all three live-zone walkers — Anthropic `/v1/messages`, OpenAI `/v1/chat/completions`, OpenAI `/v1/responses` — plus the Bedrock invoke + invoke-streaming entry points. Each passes `auth_mode` (already pre-classified by Phase F PR-F1 middleware) into the dispatcher so the gate evaluates without re-classifying. Sort key uses `tool["name"]` (Anthropic) or `tool["function"]["name"]` (OpenAI). Unnamed tools (rare; malformed inputs only) fall back to MD5 of canonical-JSON serialization for a stable in-process key — collision odds are astronomically small and `Vec::sort_by` is stable. Tests: unit tests for sort + marker detection + idempotency + the permutation property; integration tests boot the real proxy in front of a wiremock upstream and assert PAYG -> sorted, OAuth/Subscription/ marker -> byte-equal passthrough (SHA-256). --- Cargo.lock | 1 + crates/headroom-proxy/Cargo.toml | 7 + crates/headroom-proxy/src/bedrock/invoke.rs | 9 +- .../src/bedrock/invoke_streaming.rs | 9 +- .../src/cache_stabilization/mod.rs | 33 +- .../cache_stabilization/tool_def_normalize.rs | 297 +++++++++++ .../src/compression/live_zone_anthropic.rs | 494 +++++++++++++++--- .../src/compression/live_zone_openai.rs | 208 +++++++- .../src/compression/live_zone_responses.rs | 214 +++++++- crates/headroom-proxy/src/proxy.rs | 2 + .../tests/integration_tool_sort.rs | 257 +++++++++ 11 files changed, 1441 insertions(+), 90 deletions(-) create mode 100644 crates/headroom-proxy/src/cache_stabilization/tool_def_normalize.rs create mode 100644 crates/headroom-proxy/tests/integration_tool_sort.rs diff --git a/Cargo.lock b/Cargo.lock index d871bf240..c908ff8b2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1879,6 +1879,7 @@ dependencies = [ "hyper", "hyper-util", "lru", + "md-5", "pin-project-lite", "prometheus", "proptest", diff --git a/crates/headroom-proxy/Cargo.toml b/crates/headroom-proxy/Cargo.toml index 6a38c03d6..88690ce50 100644 --- a/crates/headroom-proxy/Cargo.toml +++ b/crates/headroom-proxy/Cargo.toml @@ -77,6 +77,13 @@ lru = "0.12" # Cargo.toml for rationale. gcp_auth = { workspace = true } async-trait = "0.1" +# Phase E PR-E1: tool array deterministic sort uses MD5 of canonical +# JSON as a fallback sort key for unnamed tools. MD5 is sufficient +# because the value is opaque and only used for stable in-process +# ordering — never persisted, never compared cross-host. Same crate +# the core uses for the CCR cache_key, so no additional hash backend +# enters the dep tree. +md-5 = "0.10" [dev-dependencies] tower = { workspace = true, features = ["util"] } diff --git a/crates/headroom-proxy/src/bedrock/invoke.rs b/crates/headroom-proxy/src/bedrock/invoke.rs index c503de4e8..84c93a6cb 100644 --- a/crates/headroom-proxy/src/bedrock/invoke.rs +++ b/crates/headroom-proxy/src/bedrock/invoke.rs @@ -157,7 +157,7 @@ pub async fn handle_invoke( let is_anthropic = model_id.starts_with(ANTHROPIC_VENDOR_PREFIX); let outbound_body: Bytes = if is_anthropic { - run_anthropic_compression(&body, &state, &request_id) + run_anthropic_compression(&body, &state, auth_mode, &request_id) } else { tracing::info!( event = "bedrock_compression_skipped", @@ -357,7 +357,12 @@ pub async fn handle_invoke( /// re-emission step (`ensure_anthropic_version_first`) almost always /// no-ops. We still call it as a defence-in-depth assertion that /// the byte order is correct before signing. -fn run_anthropic_compression(body: &Bytes, state: &AppState, request_id: &str) -> Bytes { +fn run_anthropic_compression( + body: &Bytes, + state: &AppState, + _auth_mode: AuthMode, + request_id: &str, +) -> Bytes { // Validate envelope shape. If the body isn't a valid Bedrock // envelope we still forward verbatim — the compressor would have // refused too — but log loudly. diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index d816cc84a..ea52bee58 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -154,7 +154,7 @@ pub async fn handle_invoke_streaming( // 1. Live-zone compression for Anthropic-shape bodies (same as D1). let is_anthropic = model_id.starts_with(ANTHROPIC_VENDOR_PREFIX); let outbound_body: Bytes = if is_anthropic { - run_anthropic_compression(&body, &state, &request_id) + run_anthropic_compression(&body, &state, auth_mode, &request_id) } else { tracing::info!( event = "bedrock_compression_skipped", @@ -828,7 +828,12 @@ fn error_response(status: StatusCode, event: &str, msg: &str) -> Response { /// flow (no body buffering required at the caller, the handler always /// owns the bytes). When PR-D3 merges, both arms can converge into a /// single helper. -fn run_anthropic_compression(body: &Bytes, state: &AppState, request_id: &str) -> Bytes { +fn run_anthropic_compression( + body: &Bytes, + state: &AppState, + _auth_mode: AuthMode, + request_id: &str, +) -> Bytes { use crate::bedrock::envelope::BedrockEnvelope; if let Err(e) = BedrockEnvelope::parse(body) { diff --git a/crates/headroom-proxy/src/cache_stabilization/mod.rs b/crates/headroom-proxy/src/cache_stabilization/mod.rs index f32cad858..99994a794 100644 --- a/crates/headroom-proxy/src/cache_stabilization/mod.rs +++ b/crates/headroom-proxy/src/cache_stabilization/mod.rs @@ -3,13 +3,18 @@ //! The realignment plan (`REALIGNMENT/07-phase-E-cache-stabilization.md`) //! groups every cache-stabilization mechanism behind one module so //! operators searching for "what does Headroom do to keep prompt -//! caches warm" land in one place. Phase E PRs in this module sit -//! *next to* the request path — either as observers -//! (volatile_detector, drift_detector) or as PAYG-gated mutators -//! (openai_cache_key, anthropic cache_control). The Phase A -//! "passthrough is sacred" invariant still holds: mutators MUST -//! gate on `AuthMode::Payg` at their call sites before invoking -//! any function that mutates the body. Observers never mutate. +//! caches warm" land in one place. Phase E PRs in this module either: +//! +//! - **Observe** inbound bodies and emit structured warnings so +//! customers can see why their prompt-cache hit rate is degrading +//! ([`volatile_detector`], PR-E5; [`drift_detector`], PR-E6). These +//! never mutate request bytes. +//! - **Normalize** request bytes to make cache hits deterministic +//! under PAYG mode ([`tool_def_normalize`], PR-E1 / PR-E2; +//! [`anthropic_cache_control`], PR-E3; [`openai_cache_key`], PR-E4). +//! These mutate bytes only when the auth-mode gate and per-policy +//! preconditions (e.g. no customer `cache_control` marker) all clear; +//! OAuth and Subscription always passthrough. //! //! Currently shipped: //! @@ -22,6 +27,10 @@ //! `cache_drift_first_request` on first sight and //! `cache_drift_observed` when consecutive requests on the same //! session disagree on any of the three dimensions. +//! - [`tool_def_normalize`] — PR-E1 / PR-E2: sorts `tools[]` +//! alphabetically by name; recursively sorts JSON Schema object +//! keys inside each tool's `input_schema`. PAYG-only; skipped when +//! any tool already carries `cache_control`. //! - [`anthropic_cache_control`] — PR-E3: on PAYG-classified //! requests where the customer hasn't placed any `cache_control` //! marker, auto-inserts one ephemeral marker on the last tool @@ -36,13 +45,13 @@ //! cache lookup to a tenant-stable identity. **Mutates the body** //! (only on PAYG) — see its docs for the gating contract. //! -//! Future PRs (E1 — tool-array sort, E2 — JSON Schema key sort, E3 — -//! `cache_control` auto-placement) hang sibling submodules off this -//! same `mod.rs`. Conflict resolution between parallel Phase E PRs -//! is intentionally trivial: each detector lives in its own file, -//! the only shared surface is this `mod.rs`'s `pub mod` list. +//! Sibling PRs hang additional submodules off this `mod.rs`. Conflict +//! resolution between parallel Phase E PRs is intentionally trivial: +//! each lives in its own file, the only shared surface is this +//! `mod.rs`'s `pub mod` list. pub mod anthropic_cache_control; pub mod drift_detector; pub mod openai_cache_key; +pub mod tool_def_normalize; pub mod volatile_detector; diff --git a/crates/headroom-proxy/src/cache_stabilization/tool_def_normalize.rs b/crates/headroom-proxy/src/cache_stabilization/tool_def_normalize.rs new file mode 100644 index 000000000..fb3ac0031 --- /dev/null +++ b/crates/headroom-proxy/src/cache_stabilization/tool_def_normalize.rs @@ -0,0 +1,297 @@ +//! PR-E1: tool array deterministic sort. +//! +//! Many client SDKs accumulate `tools[]` from a Python `set()` or a +//! `dict` whose iteration order is hash-randomized between processes. +//! The proxy sees a different tool order on every restart even though +//! the customer's source code never changed. Each shuffle busts every +//! prompt-cache hit on the cached prefix that contains the tools +//! definition, because cache hits require byte-identical bytes. +//! +//! This module provides a single mutation: sort `tools[]` alphabetically +//! by `tool["name"]`. Idempotent — re-sorting an already-sorted array +//! is a no-op (we still pay for the JSON walk but produce the same +//! bytes). +//! +//! # Cache-safety contract +//! +//! Mutating the request body is only safe under three preconditions +//! (the caller checks all three): +//! +//! 1. **PAYG auth mode.** OAuth and Subscription clients are +//! passthrough-prefer; reordering bytes for a subscription client +//! can look like cache-evasion to the upstream and trigger +//! revocation. The caller gates with [`AuthMode::Payg`]. +//! 2. **No `cache_control` marker on any tool.** When the customer has +//! explicitly placed a marker on a tool object, reordering the +//! array shifts what is "before" their marker and silently changes +//! cache scope. Their intentional layout wins. See +//! [`any_tool_has_cache_control`]. +//! 3. **Idempotency.** Re-running on already-sorted input must yield +//! byte-identical bytes; the walker uses a stable sort and rebuilds +//! objects with `serde_json::Map` (which preserves insertion order +//! via the workspace `preserve_order` feature, so the second sort +//! sees the same input as the first). +//! +//! # Why no regex +//! +//! Per `feedback_realignment_build_constraints.md` (Realignment build +//! policy): no regex for parsing. Marker detection here is a structured +//! key lookup (`tool.get("cache_control")`), not a pattern match against +//! serialized JSON. + +use md5::{Digest, Md5}; +use serde_json::Value; + +/// Sort `tools[]` deterministically by name, in place. +/// +/// Sort key: `tool["name"]` as a string. For tools missing a name (rare; +/// the API requires it but malformed inputs do reach the proxy), the +/// fallback key is the MD5 hex digest of the canonical-JSON serialization +/// of the tool object. MD5 is sufficient — the value is opaque, used +/// only for in-process ordering, never persisted, never compared across +/// hosts. +/// +/// Returns `true` if the sort changed the order, `false` if the array +/// was already sorted (idempotent signal). The caller emits a structured +/// event using this signal so dashboards can see how often the policy +/// fires. +/// +/// # Stability +/// +/// Uses `Vec::sort_by` (a stable sort: equal keys preserve original +/// order). Two unnamed tools that happen to MD5-collide will keep their +/// original relative order — collision is astronomically rare for any +/// realistic input but the contract still holds. +pub fn sort_tools_deterministically(tools: &mut Vec) -> bool { + // Capture the pre-sort key sequence so the return-value contract + // (`true` iff anything moved) is exact. We compare keys, not full + // values, because the sort is by key — equal-key swaps would not + // affect cache bytes. + let before: Vec = tools.iter().map(sort_key).collect(); + tools.sort_by(|a, b| sort_key(a).cmp(&sort_key(b))); + let after: Vec = tools.iter().map(sort_key).collect(); + before != after +} + +/// Build the deterministic sort key for a tool. Public only inside +/// this module; the public API is [`sort_tools_deterministically`]. +/// +/// Looks for the name at two known locations: +/// +/// 1. `tool["name"]` — Anthropic shape (`{"name": "...", +/// "input_schema": ...}`). +/// 2. `tool["function"]["name"]` — OpenAI Chat Completions shape +/// (`{"type": "function", "function": {"name": "...", +/// "parameters": ...}}`). +/// +/// Both providers carry the tool name in exactly one of these +/// positions; tools that match neither are rare malformed inputs that +/// fall back to the MD5-of-canonical-JSON fallback. +fn sort_key(tool: &Value) -> String { + if let Some(name) = tool.get("name").and_then(Value::as_str) { + return name.to_string(); + } + if let Some(name) = tool + .get("function") + .and_then(|f| f.get("name")) + .and_then(Value::as_str) + { + return name.to_string(); + } + // Fallback for unnamed tools: MD5 of canonical-JSON + // serialization. `serde_json::to_vec` is deterministic for a + // given `Value` because the `preserve_order` workspace feature + // pins object key order to insertion order. + let serialized = serde_json::to_vec(tool).unwrap_or_default(); + let mut hasher = Md5::new(); + hasher.update(&serialized); + let digest = hasher.finalize(); + // Hex-encode by hand to keep the dep surface tiny — `format!` + // with `{:02x}` produces the same lowercase hex `hex::encode` + // would. + let mut out = String::with_capacity(32); + for byte in digest { + out.push_str(&format!("{byte:02x}")); + } + out +} + +/// Return `true` if any tool object carries a `cache_control` field at +/// its top level. +/// +/// The Anthropic API places `cache_control` on the tool object itself +/// (e.g. `{"name": "x", "cache_control": {"type": "ephemeral"}, ...}`). +/// The customer uses this to mark a cache breakpoint that depends on +/// the tool's *position* in the array — reordering tools would shift +/// what's "before" the marker and silently change cache scope, voiding +/// their intent. So when any tool has the marker, we skip the sort. +/// +/// This function only checks the top-level field. Markers nested inside +/// `input_schema` (none of the public APIs put one there) would not be +/// caught — but they would also not be position-dependent, so the +/// safety contract still holds. +pub fn any_tool_has_cache_control(tools: &[Value]) -> bool { + tools + .iter() + .any(|tool| tool.get("cache_control").is_some()) +} + +#[cfg(test)] +mod tests { + use super::*; + use proptest::prelude::*; + use serde_json::json; + + // ─── E1: sort_tools_deterministically ───────────────────────── + + #[test] + fn sort_alphabetic_by_name() { + let mut tools = vec![ + json!({"name": "B"}), + json!({"name": "A"}), + json!({"name": "C"}), + ]; + let changed = sort_tools_deterministically(&mut tools); + assert!(changed, "out-of-order input should report a reorder"); + let names: Vec<&str> = tools + .iter() + .map(|t| t.get("name").and_then(Value::as_str).unwrap()) + .collect(); + assert_eq!(names, vec!["A", "B", "C"]); + } + + #[test] + fn idempotent_resort_no_change() { + let mut tools = vec![ + json!({"name": "A"}), + json!({"name": "B"}), + json!({"name": "C"}), + ]; + let changed = sort_tools_deterministically(&mut tools); + assert!(!changed, "already-sorted input must report no reorder"); + let names: Vec<&str> = tools + .iter() + .map(|t| t.get("name").and_then(Value::as_str).unwrap()) + .collect(); + assert_eq!(names, vec!["A", "B", "C"]); + } + + #[test] + fn byte_stable_across_runs() { + // Two independently-shuffled inputs produce byte-identical + // serialized output after sort. This is the core invariant: + // upstream sees the same bytes regardless of upstream client + // tool-collection order. + let mut input_a = vec![ + json!({"name": "search", "description": "x"}), + json!({"name": "fetch", "description": "y"}), + json!({"name": "edit", "description": "z"}), + ]; + let mut input_b = vec![ + json!({"name": "edit", "description": "z"}), + json!({"name": "search", "description": "x"}), + json!({"name": "fetch", "description": "y"}), + ]; + sort_tools_deterministically(&mut input_a); + sort_tools_deterministically(&mut input_b); + let a_bytes = serde_json::to_vec(&input_a).unwrap(); + let b_bytes = serde_json::to_vec(&input_b).unwrap(); + assert_eq!( + a_bytes, b_bytes, + "different inputs with same tool set must serialize identically after sort" + ); + } + + #[test] + fn sort_alphabetic_by_openai_function_name() { + // OpenAI Chat shape: name lives at `tool.function.name`. + let mut tools = vec![ + json!({"type": "function", "function": {"name": "Z_tool"}}), + json!({"type": "function", "function": {"name": "A_tool"}}), + json!({"type": "function", "function": {"name": "M_tool"}}), + ]; + let changed = sort_tools_deterministically(&mut tools); + assert!(changed); + let names: Vec<&str> = tools + .iter() + .map(|t| { + t.get("function") + .and_then(|f| f.get("name")) + .and_then(Value::as_str) + .unwrap() + }) + .collect(); + assert_eq!(names, vec!["A_tool", "M_tool", "Z_tool"]); + } + + #[test] + fn unnamed_tool_uses_md5_fallback() { + // Two unnamed tools — the MD5 of canonical JSON breaks ties + // deterministically. The serialized output must be stable + // across runs. + let mut tools = vec![ + json!({"description": "second"}), + json!({"description": "first"}), + ]; + let _ = sort_tools_deterministically(&mut tools); + let bytes_run1 = serde_json::to_vec(&tools).unwrap(); + + let mut tools2 = vec![ + json!({"description": "first"}), + json!({"description": "second"}), + ]; + let _ = sort_tools_deterministically(&mut tools2); + let bytes_run2 = serde_json::to_vec(&tools2).unwrap(); + + assert_eq!( + bytes_run1, bytes_run2, + "unnamed-tool MD5 fallback must produce stable byte output" + ); + } + + #[test] + fn cache_control_detection_finds_marker() { + let with_marker = vec![ + json!({"name": "A"}), + json!({"name": "B", "cache_control": {"type": "ephemeral"}}), + json!({"name": "C"}), + ]; + assert!(any_tool_has_cache_control(&with_marker)); + + let without_marker = vec![ + json!({"name": "A"}), + json!({"name": "B"}), + json!({"name": "C"}), + ]; + assert!(!any_tool_has_cache_control(&without_marker)); + } + + #[test] + fn cache_control_detection_returns_false_on_empty_tools() { + let empty: Vec = Vec::new(); + assert!(!any_tool_has_cache_control(&empty)); + } + + proptest! { + /// Sort is a permutation: no tools added, no tools removed, + /// for any reasonable mix of named / unnamed tools. + #[test] + fn sort_is_permutation( + names in prop::collection::vec( + prop::option::of("[a-zA-Z][a-zA-Z0-9_]{0,15}"), + 0..16, + ) + ) { + let mut tools: Vec = names + .iter() + .map(|maybe_name| match maybe_name { + Some(n) => json!({"name": n, "description": "x"}), + None => json!({"description": "unnamed"}), + }) + .collect(); + let len_before = tools.len(); + sort_tools_deterministically(&mut tools); + prop_assert_eq!(tools.len(), len_before); + } + } +} diff --git a/crates/headroom-proxy/src/compression/live_zone_anthropic.rs b/crates/headroom-proxy/src/compression/live_zone_anthropic.rs index c33c41ce5..5bfe76b14 100644 --- a/crates/headroom-proxy/src/compression/live_zone_anthropic.rs +++ b/crates/headroom-proxy/src/compression/live_zone_anthropic.rs @@ -45,10 +45,14 @@ use headroom_core::transforms::{ compress_anthropic_live_zone, AuthMode, BlockAction, ExclusionReason, LiveZoneError, LiveZoneOutcome, }; +use serde_json::Value; use crate::cache_stabilization::anthropic_cache_control::{ auto_place_anthropic_cache_control, AutoPlaceOutcome, SkipReason, }; +use crate::cache_stabilization::tool_def_normalize::{ + any_tool_has_cache_control, sort_tools_deterministically, +}; use crate::compression::resolve_frozen_count; use crate::config::{CacheControlAutoFrozen, CompressionMode}; @@ -112,12 +116,13 @@ pub enum PassthroughReason { /// in the body. Disabled → floor=0 (everything is in the live /// zone). /// - `auth_mode`: F1's [`RequestAuthMode`] classification of the -/// inbound request. PR-E3 gates `cache_control` auto-placement -/// on `Payg` only — OAuth and Subscription modes pass through -/// byte-equal (mutating their bytes risks looking like -/// cache-evasion to the upstream). The live-zone dispatcher -/// itself still runs on every mode in PR-B/C; the auth-mode -/// gate is local to PR-E3. +/// inbound request. Gates every Phase E byte-mutating pass — +/// PR-E1 (tool-array sort), PR-E2 (JSON Schema key sort), and +/// PR-E3 (`cache_control` auto-placement) — on `Payg` only. +/// OAuth and Subscription modes pass through byte-equal because +/// mutating their bytes risks looking like cache-evasion to the +/// upstream. The live-zone dispatcher itself still runs on every +/// mode in PR-B/C; the auth-mode gate is local to Phase E. /// - `request_id`: per-request id used for log correlation. pub fn compress_anthropic_request( body: &Bytes, @@ -166,21 +171,37 @@ pub fn compress_anthropic_request( let frozen_count = resolve_frozen_count(&parsed, cache_control_policy, request_id); - // ── PR-E3: Anthropic cache_control auto-placement ───────────── + // ── Phase E byte-mutating passes ────────────────────────────── // - // Gate 1 (auth-mode): only PAYG. Mutating bytes on OAuth / - // subscription would look like cache-evasion to the upstream. - // Gate 2 (customer-placement-wins): handled inside - // `auto_place_anthropic_cache_control` — if any marker exists - // anywhere in the body we return Skipped { MarkerPresent }. + // Three PAYG-gated passes run on the same parsed body, in this + // order, before the live-zone dispatcher: // - // When Applied, we re-serialize the parsed body and use the - // new bytes for the rest of the pipeline. The live-zone - // dispatcher will re-parse internally — this costs one extra - // serialize on the (rare) Applied path; on the Skipped / - // non-PAYG paths we don't touch the bytes at all. - let mut e3_body_bytes: Option = None; + // 1. PR-E1 — sort `tools[]` alphabetically by name. + // Skipped if any tool carries a `cache_control` marker. + // 2. PR-E3 — auto-place a `cache_control` marker on the + // (now-sorted) last tool. Skipped if any marker is already + // present anywhere in the body. + // + // Why this order: E1 must run before E3 so E3 places its marker + // on the deterministic "last tool after sort". If E3 ran first, + // E1 would correctly skip on `marker_present` but the marker + // would be on a non-deterministic tool. + // + // OAuth and Subscription auth modes pass through byte-equal — + // mutating their bytes can look like cache-evasion to the + // upstream and trigger revocation. + // + // Each gate skip emits a structured `eN_skipped` event so + // dashboards can see how often each policy fires in production. + // Each apply emits `eN_applied` with diagnostic fields. + + // PR-E1: sort tools[] in-place on the parsed value. + let normalization_applied = + normalize_tool_definitions(&mut parsed, auth_mode, request_id); + + // PR-E3: auto-place anthropic cache_control on the last tool. let mut e3_locations: Vec = Vec::new(); + let mut e3_applied: bool = false; let e3_skipped: bool; if matches!(auth_mode, RequestAuthMode::Payg) { match auto_place_anthropic_cache_control(&mut parsed) { @@ -190,35 +211,16 @@ pub fn compress_anthropic_request( } => { e3_skipped = false; if placed_count > 0 { - match serde_json::to_vec(&parsed) { - Ok(v) => { - tracing::info!( - event = "e3_applied", - request_id = %request_id, - path = "/v1/messages", - placed_count = placed_count, - locations = ?locations, - "auto-placed anthropic cache_control marker(s)" - ); - e3_body_bytes = Some(Bytes::from(v)); - e3_locations = locations; - } - Err(err) => { - // We just parsed successfully; serialize - // failure is unreachable in practice. If - // it ever fires, fall back to the - // original body bytes — never poison the - // request. Loud log so operators notice. - tracing::error!( - event = "e3_serialize_failed", - request_id = %request_id, - path = "/v1/messages", - error = %err, - "auto-placement mutated parsed body but \ - serialize-back failed; forwarding original bytes" - ); - } - } + tracing::info!( + event = "e3_applied", + request_id = %request_id, + path = "/v1/messages", + placed_count = placed_count, + locations = ?locations, + "auto-placed anthropic cache_control marker(s)" + ); + e3_applied = true; + e3_locations = locations; } else { // Applied with placed_count = 0 means "ran but // nothing to do" (no tools array, empty array, @@ -270,9 +272,32 @@ pub fn compress_anthropic_request( // counts without re-deriving them. let _ = e3_skipped; - // For the rest of the pipeline, use the E3-modified bytes when - // E3 applied, else the original buffer. - let working_body: Bytes = e3_body_bytes.clone().unwrap_or_else(|| body.clone()); + // Re-serialize the parsed value once if any Phase E pass mutated + // it. The live-zone dispatcher will re-parse internally — this + // costs one extra serialize on the (rare) mutated path; on the + // all-skipped path we don't touch the bytes at all. + let dispatch_body: Bytes = if normalization_applied.any() || e3_applied { + match serde_json::to_vec(&parsed) { + Ok(v) => Bytes::from(v), + Err(err) => { + // We just parsed successfully; serialize failure is + // unreachable in practice. If it ever fires, fall + // back to the original body bytes — never poison the + // request. Loud log so operators notice. + tracing::error!( + event = "phase_e_serialize_failed", + request_id = %request_id, + path = "/v1/messages", + error = %err, + "Phase E pass(es) mutated parsed body but \ + serialize-back failed; forwarding original bytes" + ); + body.clone() + } + } + } else { + body.clone() + }; // PR-B4: extract `body["model"]` so the live-zone dispatcher can // route the tokenizer registry to the right backend for the @@ -295,7 +320,7 @@ pub fn compress_anthropic_request( // `NoChange` otherwise (live zone empty, every compressor // declined, or every compressor produced output whose token // count was not strictly less than the input's). - match compress_anthropic_live_zone(&working_body, frozen_count, AuthMode::Payg, model) { + match compress_anthropic_live_zone(&dispatch_body, frozen_count, AuthMode::Payg, model) { Ok(LiveZoneOutcome::NoChange { manifest }) => { let block_count = manifest.block_outcomes.len(); let blocks_excluded = manifest @@ -325,18 +350,22 @@ pub fn compress_anthropic_request( live_zone_blocks_excluded = blocks_excluded, "anthropic live-zone dispatch" ); - // If E3 applied, we still need to forward the modified - // bytes even though the live-zone dispatcher made no - // additional changes. Translate to `Compressed` so the - // proxy substitutes the body — `tokens_*` are zero - // because no token-bearing block was rewritten; - // `markers_inserted` carries the E3 placement locations. - if let Some(new_body) = e3_body_bytes { + // The live-zone dispatcher made no change — but if any + // Phase E pass (E1 sort, E3 cache_control auto-placement) + // rewrote bytes, the proxy must still forward the new + // bytes. Surface as `Compressed` with the union of + // strategies and markers so the outer log/metrics layer + // attributes the byte change correctly. + if normalization_applied.any() || e3_applied { + let mut strategies = normalization_applied.strategies(); + if e3_applied { + strategies.push("e3_anthropic_cache_control"); + } Outcome::Compressed { - body: new_body, + body: dispatch_body, tokens_before: 0, tokens_after: 0, - strategies_applied: vec!["e3_anthropic_cache_control"], + strategies_applied: strategies, markers_inserted: e3_locations, } } else { @@ -372,6 +401,22 @@ pub fn compress_anthropic_request( } } } + // Stitch in the PR-E1 / PR-E2 / PR-E3 strategy tags so + // downstream log/metrics layers attribute the + // normalization / auto-placement to its distinct + // cache-stabilization surface rather than to a live-zone + // compressor that didn't actually run. + for strategy in normalization_applied.strategies() { + if !strategies.contains(&strategy) { + strategies.push(strategy); + } + } + if e3_applied { + let s = "e3_anthropic_cache_control"; + if !strategies.contains(&s) { + strategies.push(s); + } + } let body_bytes_in = body.len(); let new_body_bytes = Bytes::copy_from_slice(new_body.get().as_bytes()); let body_bytes_out = new_body_bytes.len(); @@ -403,8 +448,9 @@ pub fn compress_anthropic_request( tokens_before: original_tokens_total, tokens_after: compressed_tokens_total, strategies_applied: strategies, - // PR-B7 wires CCR retrieval-marker injection. - markers_inserted: Vec::new(), + // PR-E3 surfaces tool-slot location(s); PR-B7 will + // append CCR retrieval markers when wired. + markers_inserted: e3_locations, } } Err(LiveZoneError::BodyNotJson(_)) => { @@ -439,6 +485,111 @@ pub fn compress_anthropic_request( } } +/// Tracks which Phase E normalization steps actually mutated the +/// dispatch body. Each `bool` is `true` only when the gate cleared AND +/// the operation produced a byte-different result. Used by the caller +/// to attribute strategies on the `Outcome::Compressed` payload. +/// +/// Currently carries a single field for PR-E1; PR-E2 lands in the +/// follow-up commit and adds `e2_schema_sort` here. Keeping the flag +/// set in a struct (rather than a bare `bool`) means the second commit +/// is a pure addition. +#[derive(Debug, Clone, Copy, Default)] +pub(super) struct NormalizationApplied { + pub e1_tool_sort: bool, +} + +impl NormalizationApplied { + pub(super) fn any(self) -> bool { + self.e1_tool_sort + } + + pub(super) fn strategies(self) -> Vec<&'static str> { + let mut out = Vec::new(); + if self.e1_tool_sort { + out.push("tool_array_sort"); + } + out + } +} + +/// Apply PR-E1 (tool-array sort) in-place on the parsed body when +/// the auth-mode + marker gates clear. +/// +/// The caller owns re-serialization (because PR-E3 may also mutate +/// the same parsed value before bytes are produced). Returns a flag +/// set indicating which Phase E normalization step actually ran. +/// +/// Every gate skip emits a structured `tracing::info!` event so +/// dashboards can see how often each policy fires in production. +pub(super) fn normalize_tool_definitions( + parsed: &mut Value, + auth_mode: RequestAuthMode, + request_id: &str, +) -> NormalizationApplied { + // Auth-mode gate first — PR-E1 mutates request bytes, which is + // only safe under PAYG. OAuth and Subscription clients pass + // through byte-equal so the proxy never looks like a cache- + // evasion intermediary to the upstream. + if !matches!(auth_mode, RequestAuthMode::Payg) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/messages", + reason = "auth_mode", + auth_mode = auth_mode.as_str(), + "tool-array sort skipped: non-PAYG auth mode passes through byte-equal" + ); + return NormalizationApplied::default(); + } + + // The body must carry a `tools` array for any normalization to + // be possible. Missing / non-array `tools` → no work; this is + // not a "skip" event because it is the customer's request shape, + // not a policy gate firing. + let Some(tools_in) = parsed.get("tools").and_then(Value::as_array) else { + return NormalizationApplied::default(); + }; + if tools_in.is_empty() { + return NormalizationApplied::default(); + } + + // PR-E1 marker check. Reordering tools when any tool already + // carries `cache_control` shifts what's "before" the marker and + // silently changes cache scope. Skip the sort and pass through. + if any_tool_has_cache_control(tools_in) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/messages", + reason = "marker_present", + tool_count = tools_in.len(), + "tool-array sort skipped: customer cache_control marker present \ + on at least one tool; preserving customer-intentional order" + ); + return NormalizationApplied::default(); + } + + let tools = parsed + .get_mut("tools") + .and_then(Value::as_array_mut) + .expect("tools array verified above"); + + let mut applied = NormalizationApplied::default(); + applied.e1_tool_sort = sort_tools_deterministically(tools); + if applied.e1_tool_sort { + tracing::info!( + event = "e1_applied", + request_id = %request_id, + path = "/v1/messages", + tool_count = tools.len(), + "tool-array sort applied: tools reordered alphabetically by name" + ); + } + + applied +} + #[cfg(test)] mod tests { use super::*; @@ -580,11 +731,14 @@ mod tests { } } + // ─── PR-E3 cache_control auto-placement: unit tests ────────── + #[test] fn pr_e3_payg_with_tools_and_no_markers_returns_compressed_with_marker() { // PR-E3 happy path: PAYG body with one tool and no markers // anywhere → dispatcher inserts a marker on the last tool - // and returns Compressed with the new bytes. + // and returns Compressed with the new bytes. With one tool, + // E1 sort is a no-op so the only mutation is E3. let original = serde_json::json!({ "model": "claude-3-5-sonnet-20241022", "tools": [ @@ -609,7 +763,10 @@ mod tests { markers_inserted, .. } => { - assert_eq!(strategies_applied, vec!["e3_anthropic_cache_control"]); + assert!( + strategies_applied.contains(&"e3_anthropic_cache_control"), + "expected e3_anthropic_cache_control strategy, got: {strategies_applied:?}", + ); assert_eq!(markers_inserted, vec!["tools[0]".to_string()]); let parsed: serde_json::Value = serde_json::from_slice(&new_body).expect("re-parse new body"); @@ -713,4 +870,213 @@ mod tests { other => panic!("expected NoCompression on no-tools PAYG body, got {other:?}"), } } + + // ─── PR-E1 tool-array sort: unit tests ──────────────────────── + + #[test] + fn e1_sorts_tools_when_payg_and_no_marker() { + // PAYG, tools out of order, no `cache_control` marker → sort + // should fire. Live-zone dispatcher sees the same `messages` + // structure (no compressible blocks), so we expect + // `Outcome::Compressed` with `tool_array_sort` strategy. + // E3 also fires (no customer marker); after E1 sort, the + // last tool is "zebra" → marker lands on tools[2]. + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + "tools": [ + {"name": "zebra"}, + {"name": "apple"}, + {"name": "mango"}, + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::Payg, + "req-e1-1", + ); + match out { + Outcome::Compressed { + body: new_body, + strategies_applied, + .. + } => { + assert!( + strategies_applied.contains(&"tool_array_sort"), + "expected tool_array_sort strategy, got: {strategies_applied:?}", + ); + let parsed: Value = serde_json::from_slice(&new_body).unwrap(); + let tools = parsed.get("tools").and_then(Value::as_array).unwrap(); + let names: Vec<&str> = tools + .iter() + .map(|t| t.get("name").and_then(Value::as_str).unwrap()) + .collect(); + assert_eq!(names, vec!["apple", "mango", "zebra"]); + } + other => panic!("expected Compressed with sort, got {other:?}"), + } + } + + #[test] + fn e1_passes_through_when_oauth() { + // Same body shape; auth_mode=OAuth → byte-equal passthrough. + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + "tools": [ + {"name": "zebra"}, + {"name": "apple"}, + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::OAuth, + "req-e1-2", + ); + // Non-PAYG → no normalization → live-zone dispatcher sees + // no compressible block → NoCompression. + match out { + Outcome::NoCompression => {} + other => panic!("expected NoCompression for OAuth, got {other:?}"), + } + } + + #[test] + fn e1_passes_through_when_subscription() { + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + "tools": [ + {"name": "zebra"}, + {"name": "apple"}, + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::Subscription, + "req-e1-3", + ); + match out { + Outcome::NoCompression => {} + other => panic!("expected NoCompression for Subscription, got {other:?}"), + } + } + + #[test] + fn e1_skips_when_marker_present() { + // PAYG, but customer placed `cache_control` on a tool → + // skip the sort, byte-equal passthrough. + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + "tools": [ + {"name": "zebra"}, + {"name": "apple", "cache_control": {"type": "ephemeral"}}, + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::Payg, + "req-e1-4", + ); + match out { + Outcome::NoCompression => {} + other => panic!("expected NoCompression when marker present, got {other:?}"), + } + } + + #[test] + fn e1_skips_when_no_tools_field() { + // PAYG, no `tools` field at all → no normalization, no sort + // event, byte-equal passthrough. + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::Payg, + "req-e1-5", + ); + match out { + Outcome::NoCompression => {} + other => panic!("expected NoCompression with no tools, got {other:?}"), + } + } + + #[test] + fn e1_already_sorted_idempotent() { + // Tools in alphabetic order already — E1 sort is a no-op. + // E3 still fires (no customer marker, PAYG, has tools), so + // we still get Outcome::Compressed but only with the + // `e3_anthropic_cache_control` strategy — NOT with + // `tool_array_sort`. + let body = body_of(serde_json::json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "hi"} + ]} + ], + "tools": [ + {"name": "apple"}, + {"name": "mango"}, + {"name": "zebra"}, + ], + })); + let out = compress_anthropic_request( + &body, + CompressionMode::LiveZone, + CacheControlAutoFrozen::Disabled, + RequestAuthMode::Payg, + "req-e1-6", + ); + match out { + Outcome::Compressed { + strategies_applied, .. + } => { + assert!( + !strategies_applied.contains(&"tool_array_sort"), + "expected NO tool_array_sort strategy on already-sorted tools, got: \ + {strategies_applied:?}", + ); + assert!( + strategies_applied.contains(&"e3_anthropic_cache_control"), + "expected e3_anthropic_cache_control on already-sorted PAYG tools, got: \ + {strategies_applied:?}", + ); + } + other => panic!("expected Compressed (E3 fires) for already-sorted tools, got {other:?}"), + } + } } diff --git a/crates/headroom-proxy/src/compression/live_zone_openai.rs b/crates/headroom-proxy/src/compression/live_zone_openai.rs index 7f776932f..65f7f49d5 100644 --- a/crates/headroom-proxy/src/compression/live_zone_openai.rs +++ b/crates/headroom-proxy/src/compression/live_zone_openai.rs @@ -29,11 +29,16 @@ //! block reverts, not the whole request. use bytes::Bytes; +use headroom_core::auth_mode::AuthMode as RequestAuthMode; use headroom_core::transforms::live_zone::DEFAULT_MODEL; use headroom_core::transforms::{ compress_openai_chat_live_zone, AuthMode, BlockAction, LiveZoneError, LiveZoneOutcome, }; +use serde_json::Value; +use crate::cache_stabilization::tool_def_normalize::{ + any_tool_has_cache_control, sort_tools_deterministically, +}; use crate::compression::{Outcome, PassthroughReason}; use crate::config::CompressionMode; @@ -54,6 +59,7 @@ use crate::config::CompressionMode; pub fn compress_openai_chat_request( body: &Bytes, mode: CompressionMode, + auth_mode: RequestAuthMode, request_id: &str, ) -> Outcome { if matches!(mode, CompressionMode::Off) { @@ -119,7 +125,14 @@ pub fn compress_openai_chat_request( .and_then(serde_json::Value::as_str) .unwrap_or(DEFAULT_MODEL); - match compress_openai_chat_live_zone(body, AuthMode::Payg, model) { + // ── Phase E PR-E1: tool-array deterministic sort ──────────── + // Same gate logic as the Anthropic walker (see that module's + // `normalize_tool_definitions` for rationale). PAYG-only, + // skipped when any tool already carries `cache_control`. + let (dispatch_body, normalization_applied) = + normalize_tool_definitions_openai_chat(body, &parsed, auth_mode, request_id); + + match compress_openai_chat_live_zone(&dispatch_body, AuthMode::Payg, model) { Ok(LiveZoneOutcome::NoChange { manifest }) => { tracing::info!( event = "compression_decision", @@ -136,6 +149,15 @@ pub fn compress_openai_chat_request( model = model, "openai chat live-zone dispatch" ); + if normalization_applied.any() { + return Outcome::Compressed { + body: dispatch_body, + tokens_before: 0, + tokens_after: 0, + strategies_applied: normalization_applied.strategies(), + markers_inserted: Vec::new(), + }; + } Outcome::NoCompression } Ok(LiveZoneOutcome::Modified { new_body, manifest }) => { @@ -182,6 +204,13 @@ pub fn compress_openai_chat_request( _ => {} } } + // Stitch in PR-E1 strategy tags so dashboards see the + // tool-array sort separately from live-zone compressors. + for strategy in normalization_applied.strategies() { + if !strategies.contains(&strategy) { + strategies.push(strategy); + } + } let body_bytes_in = body.len(); let new_body_bytes = Bytes::copy_from_slice(new_body.get().as_bytes()); let body_bytes_out = new_body_bytes.len(); @@ -246,6 +275,109 @@ pub fn compress_openai_chat_request( } } +/// Tracks which Phase E normalization steps mutated the dispatch +/// body for the OpenAI Chat path. Mirrors the Anthropic walker's +/// `NormalizationApplied`. Currently carries one flag for PR-E1; +/// PR-E2 lands a second flag in the follow-up commit. +#[derive(Debug, Clone, Copy, Default)] +struct NormalizationApplied { + e1_tool_sort: bool, +} + +impl NormalizationApplied { + fn any(self) -> bool { + self.e1_tool_sort + } + + fn strategies(self) -> Vec<&'static str> { + let mut out = Vec::new(); + if self.e1_tool_sort { + out.push("tool_array_sort"); + } + out + } +} + +/// Apply PR-E1 (tool-array sort) to an OpenAI Chat Completions body +/// when the auth-mode + marker gates clear. Mirrors the Anthropic +/// walker's `normalize_tool_definitions` — same gates, same outcome +/// shape. Module-private; only the dispatcher above calls this. +fn normalize_tool_definitions_openai_chat( + body: &Bytes, + parsed: &Value, + auth_mode: RequestAuthMode, + request_id: &str, +) -> (Bytes, NormalizationApplied) { + if !matches!(auth_mode, RequestAuthMode::Payg) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/chat/completions", + reason = "auth_mode", + auth_mode = auth_mode.as_str(), + "tool-array sort skipped: non-PAYG auth mode passes through byte-equal" + ); + return (body.clone(), NormalizationApplied::default()); + } + + let Some(tools_in) = parsed.get("tools").and_then(Value::as_array) else { + return (body.clone(), NormalizationApplied::default()); + }; + if tools_in.is_empty() { + return (body.clone(), NormalizationApplied::default()); + } + + if any_tool_has_cache_control(tools_in) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/chat/completions", + reason = "marker_present", + tool_count = tools_in.len(), + "tool-array sort skipped: customer cache_control marker present \ + on at least one tool; preserving customer-intentional order" + ); + return (body.clone(), NormalizationApplied::default()); + } + + let mut working = parsed.clone(); + let tools = working + .get_mut("tools") + .and_then(Value::as_array_mut) + .expect("tools array verified above"); + + let mut applied = NormalizationApplied::default(); + applied.e1_tool_sort = sort_tools_deterministically(tools); + if applied.e1_tool_sort { + tracing::info!( + event = "e1_applied", + request_id = %request_id, + path = "/v1/chat/completions", + tool_count = tools.len(), + "tool-array sort applied: tools reordered alphabetically by name" + ); + } + + if !applied.any() { + return (body.clone(), applied); + } + + match serde_json::to_vec(&working) { + Ok(bytes) => (Bytes::from(bytes), applied), + Err(e) => { + tracing::warn!( + event = "tool_def_normalize_serialize_failed", + request_id = %request_id, + path = "/v1/chat/completions", + error = %e, + "tool-def normalization failed at re-serialize; falling back \ + to original body bytes" + ); + (body.clone(), NormalizationApplied::default()) + } + } +} + /// Inspect a Chat Completions request body and return `true` if the /// proxy should skip live-zone compression entirely. /// @@ -313,7 +445,12 @@ mod tests { #[test] fn mode_off_short_circuits() { let body = Bytes::from_static(b"not valid json"); - let out = compress_openai_chat_request(&body, CompressionMode::Off, "req-1"); + let out = compress_openai_chat_request( + &body, + CompressionMode::Off, + RequestAuthMode::Payg, + "req-1", + ); assert!(matches!( out, Outcome::Passthrough { @@ -325,7 +462,12 @@ mod tests { #[test] fn invalid_json_passthrough() { let body = Bytes::from_static(b"\x01\x02 not json"); - let out = compress_openai_chat_request(&body, CompressionMode::LiveZone, "req-2"); + let out = compress_openai_chat_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-2", + ); assert!(matches!( out, Outcome::Passthrough { @@ -337,7 +479,12 @@ mod tests { #[test] fn no_messages_passthrough() { let body = body_of(json!({"model": "gpt-4o"})); - let out = compress_openai_chat_request(&body, CompressionMode::LiveZone, "req-3"); + let out = compress_openai_chat_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-3", + ); assert!(matches!( out, Outcome::Passthrough { @@ -352,7 +499,58 @@ mod tests { "model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}] })); - let out = compress_openai_chat_request(&body, CompressionMode::LiveZone, "req-4"); + let out = compress_openai_chat_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-4", + ); + assert!(matches!(out, Outcome::NoCompression)); + } + + #[test] + fn e1_sorts_tools_when_payg() { + let body = body_of(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"type": "function", "function": {"name": "zebra"}}, + {"type": "function", "function": {"name": "apple"}}, + ], + })); + let out = compress_openai_chat_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-e1", + ); + match out { + Outcome::Compressed { + strategies_applied, .. + } => assert!( + strategies_applied.contains(&"tool_array_sort"), + "expected tool_array_sort, got {strategies_applied:?}", + ), + other => panic!("expected Compressed, got {other:?}"), + } + } + + #[test] + fn e1_passes_through_when_oauth() { + let body = body_of(json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"type": "function", "function": {"name": "zebra"}}, + {"type": "function", "function": {"name": "apple"}}, + ], + })); + let out = compress_openai_chat_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::OAuth, + "req-e1-oauth", + ); assert!(matches!(out, Outcome::NoCompression)); } diff --git a/crates/headroom-proxy/src/compression/live_zone_responses.rs b/crates/headroom-proxy/src/compression/live_zone_responses.rs index 0d4797d76..92aad25d1 100644 --- a/crates/headroom-proxy/src/compression/live_zone_responses.rs +++ b/crates/headroom-proxy/src/compression/live_zone_responses.rs @@ -32,11 +32,16 @@ //! failing block reverts. use bytes::Bytes; +use headroom_core::auth_mode::AuthMode as RequestAuthMode; use headroom_core::transforms::live_zone::DEFAULT_MODEL; use headroom_core::transforms::{ compress_openai_responses_live_zone, AuthMode, BlockAction, LiveZoneError, LiveZoneOutcome, }; +use serde_json::Value; +use crate::cache_stabilization::tool_def_normalize::{ + any_tool_has_cache_control, sort_tools_deterministically, +}; use crate::compression::{Outcome, PassthroughReason}; use crate::config::CompressionMode; @@ -53,6 +58,7 @@ use crate::config::CompressionMode; pub fn compress_openai_responses_request( body: &Bytes, mode: CompressionMode, + auth_mode: RequestAuthMode, request_id: &str, ) -> Outcome { if matches!(mode, CompressionMode::Off) { @@ -131,7 +137,11 @@ pub fn compress_openai_responses_request( .and_then(serde_json::Value::as_str) .unwrap_or(DEFAULT_MODEL); - match compress_openai_responses_live_zone(body, AuthMode::Payg, model) { + // ── Phase E PR-E1: tool-array deterministic sort ──────────── + let (dispatch_body, normalization_applied) = + normalize_tool_definitions_responses(body, &parsed, auth_mode, request_id); + + match compress_openai_responses_live_zone(&dispatch_body, AuthMode::Payg, model) { Ok(LiveZoneOutcome::NoChange { manifest }) => { tracing::info!( event = "compression_decision", @@ -148,6 +158,15 @@ pub fn compress_openai_responses_request( model = model, "openai responses live-zone dispatch" ); + if normalization_applied.any() { + return Outcome::Compressed { + body: dispatch_body, + tokens_before: 0, + tokens_after: 0, + strategies_applied: normalization_applied.strategies(), + markers_inserted: Vec::new(), + }; + } Outcome::NoCompression } Ok(LiveZoneOutcome::Modified { new_body, manifest }) => { @@ -194,6 +213,13 @@ pub fn compress_openai_responses_request( _ => {} } } + // Stitch in PR-E1 strategy tags so dashboards see the + // tool-array sort separately from live-zone compressors. + for strategy in normalization_applied.strategies() { + if !strategies.contains(&strategy) { + strategies.push(strategy); + } + } let body_bytes_in = body.len(); let new_body_bytes = Bytes::copy_from_slice(new_body.get().as_bytes()); let body_bytes_out = new_body_bytes.len(); @@ -258,6 +284,109 @@ pub fn compress_openai_responses_request( } } +/// Tracks which Phase E normalization steps mutated the dispatch +/// body for the Responses path. Sibling of the same struct in +/// `live_zone_anthropic` and `live_zone_openai`. Currently a single +/// flag for PR-E1; PR-E2 lands a second flag in the follow-up commit. +#[derive(Debug, Clone, Copy, Default)] +struct NormalizationApplied { + e1_tool_sort: bool, +} + +impl NormalizationApplied { + fn any(self) -> bool { + self.e1_tool_sort + } + + fn strategies(self) -> Vec<&'static str> { + let mut out = Vec::new(); + if self.e1_tool_sort { + out.push("tool_array_sort"); + } + out + } +} + +/// Apply PR-E1 (tool-array sort) to a Responses request body when +/// the auth-mode + marker gates clear. Mirrors the same gate logic +/// as the Anthropic / Chat Completions walkers — same skip events, +/// same outcome shape, same byte-equal passthrough on non-PAYG. +fn normalize_tool_definitions_responses( + body: &Bytes, + parsed: &Value, + auth_mode: RequestAuthMode, + request_id: &str, +) -> (Bytes, NormalizationApplied) { + if !matches!(auth_mode, RequestAuthMode::Payg) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/responses", + reason = "auth_mode", + auth_mode = auth_mode.as_str(), + "tool-array sort skipped: non-PAYG auth mode passes through byte-equal" + ); + return (body.clone(), NormalizationApplied::default()); + } + + let Some(tools_in) = parsed.get("tools").and_then(Value::as_array) else { + return (body.clone(), NormalizationApplied::default()); + }; + if tools_in.is_empty() { + return (body.clone(), NormalizationApplied::default()); + } + + if any_tool_has_cache_control(tools_in) { + tracing::info!( + event = "e1_skipped", + request_id = %request_id, + path = "/v1/responses", + reason = "marker_present", + tool_count = tools_in.len(), + "tool-array sort skipped: customer cache_control marker present \ + on at least one tool; preserving customer-intentional order" + ); + return (body.clone(), NormalizationApplied::default()); + } + + let mut working = parsed.clone(); + let tools = working + .get_mut("tools") + .and_then(Value::as_array_mut) + .expect("tools array verified above"); + + let mut applied = NormalizationApplied::default(); + applied.e1_tool_sort = sort_tools_deterministically(tools); + if applied.e1_tool_sort { + tracing::info!( + event = "e1_applied", + request_id = %request_id, + path = "/v1/responses", + tool_count = tools.len(), + "tool-array sort applied: tools reordered alphabetically by name" + ); + } + + if !applied.any() { + return (body.clone(), applied); + } + + match serde_json::to_vec(&working) { + Ok(bytes) => (Bytes::from(bytes), applied), + Err(e) => { + tracing::warn!( + event = "tool_def_normalize_serialize_failed", + request_id = %request_id, + path = "/v1/responses", + error = %e, + "tool-def normalization failed at re-serialize; falling back \ + to original body bytes" + ); + (body.clone(), NormalizationApplied::default()) + } + } +} + /// Walk the items array once and emit per-item telemetry. Recognised /// item types are tallied; unknown `type` values trigger a /// `tracing::warn!` `event = responses_unknown_item_type` but never @@ -356,7 +485,12 @@ mod tests { #[test] fn mode_off_short_circuits() { let body = Bytes::from_static(b"not valid json"); - let out = compress_openai_responses_request(&body, CompressionMode::Off, "req-1"); + let out = compress_openai_responses_request( + &body, + CompressionMode::Off, + RequestAuthMode::Payg, + "req-1", + ); assert!(matches!( out, Outcome::Passthrough { @@ -368,7 +502,12 @@ mod tests { #[test] fn invalid_json_passthrough() { let body = Bytes::from_static(b"\x01\x02 not json"); - let out = compress_openai_responses_request(&body, CompressionMode::LiveZone, "req-2"); + let out = compress_openai_responses_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-2", + ); assert!(matches!( out, Outcome::Passthrough { @@ -380,7 +519,12 @@ mod tests { #[test] fn no_input_passthrough() { let body = body_of(json!({"model": "gpt-4o"})); - let out = compress_openai_responses_request(&body, CompressionMode::LiveZone, "req-3"); + let out = compress_openai_responses_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-3", + ); assert!(matches!( out, Outcome::Passthrough { @@ -398,7 +542,67 @@ mod tests { "content": [{"type": "input_text", "text": "hi"}]} ] })); - let out = compress_openai_responses_request(&body, CompressionMode::LiveZone, "req-4"); + let out = compress_openai_responses_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-4", + ); + assert!(matches!(out, Outcome::NoCompression)); + } + + #[test] + fn e1_sorts_tools_when_payg() { + // PAYG, Responses-shape body, tools out of order (and using + // OpenAI's `function`-nested name — same shape as Chat + // Completions). + let body = body_of(json!({ + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", + "content": [{"type": "input_text", "text": "hi"}]} + ], + "tools": [ + {"type": "function", "function": {"name": "zebra"}}, + {"type": "function", "function": {"name": "apple"}}, + ], + })); + let out = compress_openai_responses_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::Payg, + "req-e1-resp", + ); + match out { + Outcome::Compressed { + strategies_applied, .. + } => assert!( + strategies_applied.contains(&"tool_array_sort"), + "expected tool_array_sort, got {strategies_applied:?}", + ), + other => panic!("expected Compressed, got {other:?}"), + } + } + + #[test] + fn e1_passes_through_when_oauth() { + let body = body_of(json!({ + "model": "gpt-4o", + "input": [ + {"type": "message", "role": "user", + "content": [{"type": "input_text", "text": "hi"}]} + ], + "tools": [ + {"type": "function", "function": {"name": "zebra"}}, + {"type": "function", "function": {"name": "apple"}}, + ], + })); + let out = compress_openai_responses_request( + &body, + CompressionMode::LiveZone, + RequestAuthMode::OAuth, + "req-e1-oauth", + ); assert!(matches!(out, Outcome::NoCompression)); } } diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index 1ae12e1ca..e3f7ee732 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -661,6 +661,7 @@ pub(crate) async fn forward_http( compression::compress_openai_chat_request( &buffered, state.config.compression_mode, + auth_mode, &request_id, ) } @@ -674,6 +675,7 @@ pub(crate) async fn forward_http( compression::compress_openai_responses_request( &buffered, state.config.compression_mode, + auth_mode, &request_id, ) } diff --git a/crates/headroom-proxy/tests/integration_tool_sort.rs b/crates/headroom-proxy/tests/integration_tool_sort.rs new file mode 100644 index 000000000..e6dfd71bb --- /dev/null +++ b/crates/headroom-proxy/tests/integration_tool_sort.rs @@ -0,0 +1,257 @@ +//! Integration tests for PR-E1: tool array deterministic sort. +//! +//! Boots a real Rust proxy in front of a wiremock upstream and +//! exercises the three live-zone walkers via the inbound paths the +//! proxy actually serves. Asserts: +//! +//! 1. **PAYG path** (e.g. `x-api-key` on Anthropic): tools arrive +//! at the upstream sorted alphabetically, regardless of the +//! client's input order. +//! 2. **Subscription path** (UA prefix `claude-cli/...`): tools +//! pass through verbatim — bytes the upstream sees match the +//! bytes the client sent (asserted via SHA-256 byte-equality). +//! 3. **Customer-marker path** (PAYG, but at least one tool already +//! carries `cache_control`): tools pass through verbatim — the +//! sort is gated off so the customer's intentional layout wins. +//! +//! The Phase A cache-safety invariant — the proxy NEVER mutates +//! request bytes when a gate skips — is the contract under test. + +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}; + +fn sha256_hex(bytes: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(bytes); + let digest = hasher.finalize(); + let mut s = String::with_capacity(64); + for b in digest { + s.push_str(&format!("{b:02x}")); + } + s +} + +/// Mount a `/v1/messages` handler that captures the upstream-received +/// request body for later inspection. +async fn mount_anthropic_capture(upstream: &MockServer) -> Arc>>> { + let captured: Arc>>> = Arc::new(Mutex::new(None)); + let captured_clone = captured.clone(); + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(move |req: &wiremock::Request| { + *captured_clone.lock().unwrap() = Some(req.body.clone()); + ResponseTemplate::new(200).set_body_string(r#"{"ok":true}"#) + }) + .mount(upstream) + .await; + captured +} + +/// PAYG: send tools in reverse alphabetical order, expect upstream to +/// receive them sorted by `name`. +#[tokio::test] +async fn payg_request_with_unsorted_tools_is_sorted_at_upstream() { + 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 = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"name": "zebra", "description": "z"}, + {"name": "apple", "description": "a"}, + {"name": "mango", "description": "m"}, + ], + }); + let body = serde_json::to_vec(&payload).unwrap(); + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + // PAYG signal: x-api-key header (Anthropic API-key style). + .header("x-api-key", "sk-ant-api03-abc") + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let upstream_body = captured + .lock() + .unwrap() + .clone() + .expect("upstream should have captured"); + let parsed: Value = serde_json::from_slice(&upstream_body).expect("upstream body is JSON"); + let tools = parsed.get("tools").and_then(Value::as_array).unwrap(); + let names: Vec<&str> = tools + .iter() + .map(|t| t.get("name").and_then(Value::as_str).unwrap()) + .collect(); + assert_eq!( + names, + vec!["apple", "mango", "zebra"], + "PAYG path must deliver tools to upstream sorted alphabetically by name", + ); + + proxy.shutdown().await; +} + +/// Subscription: same body shape but with a `claude-cli/...` UA → +/// proxy must NOT mutate. Upstream-received bytes must be byte-equal +/// to client-sent bytes (SHA-256 match). +#[tokio::test] +async fn subscription_request_passes_tools_through_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; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + }) + .await; + + let payload = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"name": "zebra", "description": "z"}, + {"name": "apple", "description": "a"}, + ], + }); + let body = serde_json::to_vec(&payload).unwrap(); + let body_hash = sha256_hex(&body); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + // Subscription signal: claude-cli UA prefix. + .header("user-agent", "claude-cli/1.0.0") + .header("authorization", "Bearer sk-ant-oat-pretend") + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let upstream_body = captured + .lock() + .unwrap() + .clone() + .expect("upstream should have captured"); + assert_eq!( + sha256_hex(&upstream_body), + body_hash, + "Subscription path must pass body bytes through unchanged" + ); + + proxy.shutdown().await; +} + +/// PAYG, but customer placed `cache_control` on a tool → sort is +/// skipped; bytes pass through verbatim. +#[tokio::test] +async fn payg_with_marker_passes_tools_through_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; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + }) + .await; + + let payload = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"name": "zebra", "description": "z"}, + {"name": "apple", "description": "a", "cache_control": {"type": "ephemeral"}}, + ], + }); + let body = serde_json::to_vec(&payload).unwrap(); + let body_hash = sha256_hex(&body); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + .header("x-api-key", "sk-ant-api03-abc") + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let upstream_body = captured + .lock() + .unwrap() + .clone() + .expect("upstream should have captured"); + assert_eq!( + sha256_hex(&upstream_body), + body_hash, + "PAYG with customer cache_control marker must pass body bytes through unchanged" + ); + + proxy.shutdown().await; +} + +/// OAuth: same body shape but with a `Bearer sk-ant-oat-...` token +/// (no claude-cli UA prefix → classified OAuth, not Subscription). +/// Tools pass through verbatim. +#[tokio::test] +async fn oauth_request_passes_tools_through_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; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + }) + .await; + + let payload = json!({ + "model": "claude-3-5-sonnet-20241022", + "max_tokens": 32, + "messages": [{"role": "user", "content": "hi"}], + "tools": [ + {"name": "zebra", "description": "z"}, + {"name": "apple", "description": "a"}, + ], + }); + let body = serde_json::to_vec(&payload).unwrap(); + let body_hash = sha256_hex(&body); + + let resp = reqwest::Client::new() + .post(format!("{}/v1/messages", proxy.url())) + // OAuth signal: Anthropic OAuth token shape. + .header("authorization", "Bearer sk-ant-oat-foo") + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let upstream_body = captured + .lock() + .unwrap() + .clone() + .expect("upstream should have captured"); + assert_eq!( + sha256_hex(&upstream_body), + body_hash, + "OAuth path must pass body bytes through unchanged" + ); + + proxy.shutdown().await; +}