diff --git a/crates/headroom-proxy/src/bedrock/invoke.rs b/crates/headroom-proxy/src/bedrock/invoke.rs index 2571e8bcb..a30e0027a 100644 --- a/crates/headroom-proxy/src/bedrock/invoke.rs +++ b/crates/headroom-proxy/src/bedrock/invoke.rs @@ -615,6 +615,7 @@ mod tests { client: reqwest::Client::new(), bedrock_credentials: None, drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + beta_sticky: crate::cache_stabilization::beta_sticky::BetaStickyState::new(8), vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( "test".to_string(), )), @@ -649,6 +650,7 @@ mod tests { // unit test never observes drift, but `AppState` requires // the field to be populated. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + beta_sticky: crate::cache_stabilization::beta_sticky::BetaStickyState::new(8), // PR-D4: unit tests for the Bedrock URL builder don't // touch the Vertex route, but `AppState` is one struct // — supply a dummy token source so the test compiles. @@ -684,6 +686,7 @@ mod tests { // PR-E6: see above — drift detector is unused by this // test; we just satisfy the struct shape. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + beta_sticky: crate::cache_stabilization::beta_sticky::BetaStickyState::new(8), // PR-D4: unit tests for the Bedrock URL builder don't // touch the Vertex route, but `AppState` is one struct // — supply a dummy token source so the test compiles. diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index ce66429f9..62501003c 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -1014,6 +1014,7 @@ mod tests { // PR-E6: drift detector is unused by this URL-builder // unit test; small capacity to satisfy the struct shape. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + beta_sticky: crate::cache_stabilization::beta_sticky::BetaStickyState::new(8), // PR-D4: unit tests for the Bedrock URL builder don't // touch the Vertex route, but `AppState` is one struct // — supply a dummy token source so the test compiles. @@ -1056,6 +1057,7 @@ mod tests { client: reqwest::Client::new(), bedrock_credentials: None, drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + beta_sticky: crate::cache_stabilization::beta_sticky::BetaStickyState::new(8), vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( "test".to_string(), )), diff --git a/crates/headroom-proxy/src/cache_stabilization/beta_sticky.rs b/crates/headroom-proxy/src/cache_stabilization/beta_sticky.rs new file mode 100644 index 000000000..a8db6a5bc --- /dev/null +++ b/crates/headroom-proxy/src/cache_stabilization/beta_sticky.rs @@ -0,0 +1,620 @@ +//! Session-sticky provider beta headers — Rust port of the Python +//! proxy's `SessionBetaTracker` (PR-A6, `headroom/proxy/helpers.py`). +//! +//! ## Why +//! +//! Provider beta headers (`anthropic-beta`, `openai-beta`) are part of +//! the request bytes that determine the upstream prefix-cache key. +//! Interactive clients (Claude Code, Codex CLI) MAY drop a beta token +//! between turn N and turn N+1 of the same conversation; the cache hot +//! zone is positional, so the next turn's prefix hashes differently and +//! the prefix-cache read misses — the customer silently pays for a full +//! prompt re-write. The Python proxy defeats this with a bounded LRU +//! tracker that unions the client's tokens with every token previously +//! seen for the same `(provider, session)` and forwards the union. +//! +//! The Rust proxy replaces the Python request path in Phase H, which +//! deletes `SessionBetaTracker` with the rest of +//! `headroom/proxy/helpers.py`. Without this port the protection — +//! and its documented operator contract +//! (`docs/content/docs/configuration.mdx`, "Session Beta Header +//! Tracking") — would silently not survive the migration. +//! +//! ## Behaviour contract (parity with Python) +//! +//! - Union client tokens with previously-seen tokens for the session, +//! preserving first-seen order; case-insensitive dedup where the +//! first-seen casing wins. +//! - Keyed by `(provider, session)` so the same session id against +//! Anthropic and OpenAI upstreams keeps independent token sets. +//! - Bounded LRU (`BETA_TRACKER_CAPACITY` sessions): lookups touch +//! recency, overflow evicts the oldest session. +//! - The tracker only ever records tokens the client itself sent. +//! Headroom-added tokens (e.g. memory-tool betas on the Python +//! path) are NOT recorded — the forwarded union is always a subset +//! of values this client already put on the wire, which is what +//! keeps the mechanism consistent with the subscription-stealth +//! invariant (REALIGNMENT invariant #10: "no beta drift"). +//! +//! The operator opt-out lives at the call site: when +//! `Config::beta_header_sticky` is `disabled` the proxy skips the +//! tracker entirely and forwards the client header verbatim (the +//! Python proxy's `HEADROOM_BETA_HEADER_STICKY=disabled` diagnostic +//! mode). That gate is per REALIGNMENT build constraint #4 an explicit +//! loud opt-in, not a silent fallback. +//! +//! Session identity comes from +//! [`super::drift_detector::derive_session_key`] — the same +//! conversation-aware key the drift detector uses (explicit +//! `x-headroom-session-id` when the client declares it, otherwise +//! credential/IP arms folded with a first-message conversation +//! discriminator). +//! +//! ## Divergence from Python: per-conversation, not per-(model, system) +//! +//! The Python tracker keys on the store session id — explicit header, +//! else a hash of `(model, leading system prompt)` — so all parallel +//! conversations sharing a model + system prompt (a Claude Code +//! session and every one of its subagents) share ONE token union and +//! cross-inherit each other's tokens. This port keys on the drift +//! detector's conversation-aware key instead, so each conversation +//! keeps its own union; the integration test +//! `separate_conversations_do_not_leak_tokens` pins that. Deliberate: +//! the `(model, system)` bucket conflating parallel agentic +//! conversations is the exact defect #2085 / #2193 / #2301 chased out +//! of the other session-sticky subsystems. The cost is losing +//! Python's accidental cross-conversation repair (conversation B +//! turn 1 inheriting a token only conversation A ever sent); each +//! conversation's stickiness now starts from its own first sighting, +//! which is also the only variant that can't leak one tenant-visible +//! experiment token into an unrelated conversation's request bytes. + +use std::collections::HashSet; +use std::num::NonZeroUsize; +use std::sync::{Arc, Mutex}; + +use http::header::{HeaderMap, HeaderValue}; +use lru::LruCache; + +use super::drift_detector::session_key_log_prefix; + +/// Maximum number of `(provider, session)` entries tracked. Sessions +/// are keyed per conversation (see module docs), so the working set is +/// the number of concurrently active conversations — same sizing +/// rationale as the drift detector's capacity. Eviction cost is +/// re-learning a live session's dropped tokens from scratch (the next +/// turn forwards the client value verbatim), not a lost request. +pub const BETA_TRACKER_CAPACITY: usize = 1000; + +/// Upstream namespace for a tracked beta-token set. Mirrors the +/// Python tracker's `provider` string key ("anthropic" / "openai"). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BetaProvider { + /// `/v1/messages` — `anthropic-beta` header. + Anthropic, + /// `/v1/chat/completions` and `/v1/responses` — `openai-beta` + /// header. One namespace for both endpoints, matching the Python + /// proxy's single `provider="openai"` key. + OpenAi, +} + +impl BetaProvider { + /// Stable lower-case label for log fields; matches the Python + /// tracker's provider strings. + pub fn as_str(self) -> &'static str { + match self { + BetaProvider::Anthropic => "anthropic", + BetaProvider::OpenAi => "openai", + } + } + + /// The request header this provider's beta tokens travel in. + pub fn header_name(self) -> &'static str { + match self { + BetaProvider::Anthropic => "anthropic-beta", + BetaProvider::OpenAi => "openai-beta", + } + } +} + +/// Split a comma-separated beta-header value into trimmed, non-empty +/// tokens. Port of the Python `split_beta_tokens` helper. +pub fn split_beta_tokens(value: Option<&str>) -> Vec { + value + .unwrap_or("") + .split(',') + .map(str::trim) + .filter(|t| !t.is_empty()) + .map(str::to_string) + .collect() +} + +/// Per-session ordered token lists, keyed by `(provider, session)`. +type SessionTokenCache = LruCache<(BetaProvider, String), Vec>; + +/// Bounded LRU of beta tokens observed per `(provider, session)`. +/// +/// Cloning shares the underlying map (`Arc`), mirroring +/// [`super::drift_detector::DriftState`] so one instance lives in +/// `AppState` and clones freely into every handler path. +#[derive(Clone)] +pub struct BetaStickyState { + sessions: Arc>, +} + +impl BetaStickyState { + /// Create a tracker bounded to `capacity` sessions. + /// + /// # Panics + /// + /// Panics when `capacity == 0`, mirroring `DriftState::new` (the + /// Python tracker raises `ValueError` on a non-positive bound). + pub fn new(capacity: usize) -> Self { + let cap = NonZeroUsize::new(capacity).expect("BetaStickyState capacity must be > 0"); + Self { + sessions: Arc::new(Mutex::new(LruCache::new(cap))), + } + } + + /// Union `client_value`'s tokens with the session's previously + /// seen tokens, update the session, and return the merged + /// comma-separated value (possibly empty). Port of the Python + /// `SessionBetaTracker.record_and_get_sticky_betas`. + /// + /// On a poisoned lock the tracker fails open: the client value is + /// returned verbatim (trimmed) and state is left untouched — + /// never drop or delay the request for a telemetry-adjacent + /// protection. + pub fn record_and_get_sticky_betas( + &self, + provider: BetaProvider, + session_key: &str, + client_value: Option<&str>, + ) -> String { + let client_tokens = split_beta_tokens(client_value); + + let mut sessions = match self.sessions.lock() { + Ok(guard) => guard, + Err(poisoned) => { + tracing::warn!( + event = "beta_sticky_lock_poisoned", + provider = provider.as_str(), + "beta tracker lock poisoned; forwarding client value verbatim" + ); + drop(poisoned); + return client_tokens.join(","); + } + }; + + let key = (provider, session_key.to_string()); + // `get_mut` touches LRU recency on hit, mirroring the Python + // tracker's move-to-end. + if let Some(merged) = sessions.get_mut(&key) { + // Dedup is case-insensitive with the first-seen casing + // winning. Header values reaching this point are visible + // ASCII (`HeaderValue::to_str` rejects anything else), so + // ASCII lowercasing matches Python's `str.lower()` over + // the reachable domain. + let mut seen: HashSet = merged.iter().map(|t| t.to_ascii_lowercase()).collect(); + for token in client_tokens { + if seen.insert(token.to_ascii_lowercase()) { + merged.push(token); + } + } + return merged.join(","); + } + + let mut merged: Vec = Vec::with_capacity(client_tokens.len()); + let mut seen: HashSet = HashSet::with_capacity(client_tokens.len()); + for token in client_tokens { + if seen.insert(token.to_ascii_lowercase()) { + merged.push(token); + } + } + let joined = merged.join(","); + // `put` on a fresh key evicts the oldest entry once the cache + // is at capacity — the Python tracker's bounded-LRU overflow + // pop. Sessions that never sent a beta token still occupy a + // slot (Python stores their empty list too); the cost is one + // LRU entry, the benefit is identical recency behaviour. + sessions.put(key, merged); + joined + } + + /// Number of tracked sessions (test observability). + #[cfg(test)] + fn active_sessions(&self) -> usize { + self.sessions.lock().map(|c| c.len()).unwrap_or(0) + } +} + +/// Count tokens in a raw header value without allocating a `Vec` +/// (log-field helper; same tokenization as [`split_beta_tokens`]). +fn count_beta_tokens(value: Option<&str>) -> usize { + value + .unwrap_or("") + .split(',') + .filter(|t| !t.trim().is_empty()) + .count() +} + +/// Record the client's beta header for this `(provider, session)` and +/// rewrite the upstream-bound header to the session union when they +/// differ. The full merge site: reads `provider.header_name()` from +/// `outgoing_headers`, unions via the tracker, mutates the map in +/// place. Mirrors the Python handler block (anthropic.py PR-A6): +/// rewrite only when the union is non-empty and differs from the +/// client value; an absent client header gains the union; a session +/// with no tokens anywhere stays header-less. +/// +/// Fail-open contract: a client value that isn't visible ASCII is +/// forwarded verbatim and nothing is recorded (never rewrite what we +/// can't faithfully parse); an unencodable union (unreachable — every +/// token came from a parsed header value) logs and forwards verbatim. +/// +/// Logging: counts only — beta tokens can carry experiment IDs the +/// user hasn't opted to share with Headroom logs (Python +/// `log_beta_header_merge` contract). Python logs every merge at +/// info; here the no-op case drops to debug, matching the drift +/// detector's silent-on-stable precedent, so an info-level +/// `beta_header_merge` always marks an actual cache-affecting +/// rewrite. +pub fn apply_sticky_betas( + tracker: &BetaStickyState, + provider: BetaProvider, + session_key: &str, + outgoing_headers: &mut HeaderMap, + request_id: &str, +) { + let header_name = provider.header_name(); + // Join repeated field lines with "," per RFC 9110 §5.3 list + // semantics BEFORE recording, so a client sending two beta lines + // has both recorded and a later rewrite (which `insert`s a single + // line, dropping the others) can never shrink the upstream token + // set mid-conversation. + let mut parts: Vec<&str> = Vec::new(); + for raw in outgoing_headers.get_all(header_name) { + match raw.to_str() { + Ok(s) => parts.push(s), + Err(_) => { + tracing::debug!( + event = "beta_header_merge_skipped", + request_id = %request_id, + provider = provider.as_str(), + reason = "non_ascii_header_value", + "client beta header is not visible ASCII; forwarding verbatim" + ); + return; + } + } + } + let client_value: Option = if parts.is_empty() { + None + } else { + Some(parts.join(",")) + }; + + let sticky = + tracker.record_and_get_sticky_betas(provider, session_key, client_value.as_deref()); + let rewritten = !sticky.is_empty() && sticky != client_value.as_deref().unwrap_or(""); + if rewritten { + match HeaderValue::from_str(&sticky) { + Ok(value) => { + outgoing_headers.insert(header_name, value); + } + Err(error) => { + tracing::warn!( + event = "beta_header_merge_skipped", + request_id = %request_id, + provider = provider.as_str(), + reason = "unencodable_union", + error = %error, + "sticky beta union not encodable as a header value" + ); + return; + } + } + } + + let client_betas = count_beta_tokens(client_value.as_deref()); + let sticky_betas = count_beta_tokens(Some(&sticky)); + if rewritten { + tracing::info!( + event = "beta_header_merge", + request_id = %request_id, + provider = provider.as_str(), + session_key_hash = %session_key_log_prefix(session_key), + client_betas, + sticky_betas, + "session-sticky beta merge rewrote the upstream header" + ); + } else { + tracing::debug!( + event = "beta_header_merge", + request_id = %request_id, + provider = provider.as_str(), + session_key_hash = %session_key_log_prefix(session_key), + client_betas, + sticky_betas, + "session-sticky beta merge (no-op)" + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ----------------------------------------------------------------- + // split_beta_tokens — port of the Python tokenizer contract. + // ----------------------------------------------------------------- + + #[test] + fn split_none_and_empty_yield_no_tokens() { + assert!(split_beta_tokens(None).is_empty()); + assert!(split_beta_tokens(Some("")).is_empty()); + assert!(split_beta_tokens(Some(" ")).is_empty()); + assert!(split_beta_tokens(Some(",, ,")).is_empty()); + } + + #[test] + fn split_trims_and_drops_empty_segments() { + assert_eq!( + split_beta_tokens(Some(" a , ,b, c-1 ")), + vec!["a".to_string(), "b".to_string(), "c-1".to_string()] + ); + } + + // ----------------------------------------------------------------- + // record_and_get_sticky_betas — tracker semantics ported from + // tests/test_anthropic_beta_session_sticky.py. + // ----------------------------------------------------------------- + + fn tracker() -> BetaStickyState { + BetaStickyState::new(BETA_TRACKER_CAPACITY) + } + + #[test] + fn first_request_returns_client_tokens() { + let t = tracker(); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a,b")); + assert_eq!(got, "a,b"); + } + + #[test] + fn dropped_token_is_reinjected_on_next_turn() { + // The cache-killer this module exists for: turn N sends + // "a,b", turn N+1 drops "b" — the union must restore it. + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a,b")); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a")); + assert_eq!(got, "a,b"); + } + + #[test] + fn union_preserves_first_seen_order() { + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("b,a")); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a,c")); + assert_eq!(got, "b,a,c"); + } + + #[test] + fn dedup_is_case_insensitive_first_casing_wins() { + let t = tracker(); + t.record_and_get_sticky_betas( + BetaProvider::Anthropic, + "s1", + Some("Context-Management-2025-06-27"), + ); + let got = t.record_and_get_sticky_betas( + BetaProvider::Anthropic, + "s1", + Some("context-management-2025-06-27"), + ); + assert_eq!(got, "Context-Management-2025-06-27"); + } + + #[test] + fn duplicate_client_tokens_are_deduped() { + let t = tracker(); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a,a,b,A")); + assert_eq!(got, "a,b"); + } + + #[test] + fn client_whitespace_is_trimmed_in_union() { + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some(" a , b ")); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("c ")); + assert_eq!(got, "a,b,c"); + } + + #[test] + fn absent_client_value_returns_session_union() { + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a")); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", None); + assert_eq!(got, "a"); + } + + #[test] + fn empty_session_and_client_yield_empty_string() { + let t = tracker(); + assert_eq!( + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", None), + "" + ); + } + + #[test] + fn providers_keep_independent_namespaces() { + // Same session id, different providers — token sets must not + // leak across (Python: the (provider, session_id) tuple key). + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("anth-only")); + let got = t.record_and_get_sticky_betas(BetaProvider::OpenAi, "s1", Some("oai-only")); + assert_eq!(got, "oai-only"); + } + + #[test] + fn sessions_keep_independent_token_sets() { + let t = tracker(); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a")); + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s2", Some("b")); + assert_eq!(got, "b"); + } + + #[test] + fn lru_evicts_oldest_session_at_capacity() { + let t = BetaStickyState::new(2); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a")); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s2", Some("b")); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s3", Some("c")); + assert_eq!(t.active_sessions(), 2); + // s1 was evicted: its history is gone, so a bare re-request + // returns only the fresh client value. + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("z")); + assert_eq!(got, "z"); + } + + #[test] + fn lru_hit_touches_recency() { + let t = BetaStickyState::new(2); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("a")); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s2", Some("b")); + // Touch s1 so s2 becomes the eviction candidate. + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", None); + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s3", Some("c")); + // s1 survived the s3 insert… + assert_eq!( + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", None), + "a" + ); + // …and s2 did not. + assert_eq!( + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s2", None), + "" + ); + } + + // ----------------------------------------------------------------- + // apply_sticky_betas — header-map plumbing. + // ----------------------------------------------------------------- + + fn header_map(values: &[&str]) -> HeaderMap { + let mut map = HeaderMap::new(); + for v in values { + map.append("anthropic-beta", HeaderValue::from_str(v).unwrap()); + } + map + } + + fn beta_values(map: &HeaderMap) -> Vec { + map.get_all("anthropic-beta") + .iter() + .map(|v| v.to_str().unwrap().to_string()) + .collect() + } + + #[test] + fn apply_rewrites_dropped_token_to_union() { + let t = tracker(); + let mut turn1 = header_map(&["a,b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn1, "req-1"); + assert_eq!(beta_values(&turn1), vec!["a,b"]); + + let mut turn2 = header_map(&["a"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn2, "req-2"); + assert_eq!(beta_values(&turn2), vec!["a,b"]); + } + + #[test] + fn apply_reinserts_union_when_header_fully_omitted() { + let t = tracker(); + let mut turn1 = header_map(&["a,b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn1, "req-1"); + + let mut turn2 = HeaderMap::new(); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn2, "req-2"); + assert_eq!(beta_values(&turn2), vec!["a,b"]); + } + + #[test] + fn apply_never_invents_a_header() { + let t = tracker(); + let mut map = HeaderMap::new(); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut map, "req-1"); + assert!(map.get("anthropic-beta").is_none()); + } + + #[test] + fn apply_noop_leaves_header_lines_untouched() { + let t = tracker(); + let mut map = header_map(&["a,b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut map, "req-1"); + let mut again = header_map(&["a,b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut again, "req-2"); + assert_eq!(beta_values(&again), vec!["a,b"]); + } + + #[test] + fn apply_records_all_repeated_header_lines() { + // RFC 9110 list semantics: two field lines are one list. The + // union must record BOTH lines, so a later rewrite (which + // collapses to a single line) can never shrink the upstream + // token set mid-conversation. + let t = tracker(); + let mut turn1 = header_map(&["a,x", "b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn1, "req-1"); + // No rewrite on turn 1 (union == joined client list): both + // lines pass through untouched. + assert_eq!(beta_values(&turn1), vec!["a,x", "b"]); + + // Turn 2 drops "x" from the first line: the rewrite must + // carry the full set from both turn-1 lines. + let mut turn2 = header_map(&["a", "b"]); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut turn2, "req-2"); + assert_eq!(beta_values(&turn2), vec!["a,x,b"]); + } + + #[test] + fn apply_skips_non_ascii_value_and_records_nothing() { + let t = tracker(); + let mut map = HeaderMap::new(); + map.insert( + "anthropic-beta", + HeaderValue::from_bytes(&[0xfa, 0xfb]).unwrap(), + ); + apply_sticky_betas(&t, BetaProvider::Anthropic, "s1", &mut map, "req-1"); + // Wire bytes untouched… + assert_eq!(map.get("anthropic-beta").unwrap().as_bytes(), &[0xfa, 0xfb]); + // …and nothing recorded: the next ASCII turn sees only its + // own tokens. + let got = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some("y")); + assert_eq!(got, "y"); + } + + #[test] + fn concurrent_unions_lose_no_tokens() { + // Port of the Python thread-hammering test: concurrent turns + // on one session must never drop a recorded token. + let t = tracker(); + std::thread::scope(|s| { + for i in 0..8 { + let t = t.clone(); + s.spawn(move || { + let token = format!("tok-{i}"); + for _ in 0..50 { + t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", Some(&token)); + } + }); + } + }); + let merged = t.record_and_get_sticky_betas(BetaProvider::Anthropic, "s1", None); + let tokens: HashSet<&str> = merged.split(',').collect(); + for i in 0..8 { + assert!(tokens.contains(format!("tok-{i}").as_str())); + } + } +} diff --git a/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs b/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs index c288d3131..12ca1f709 100644 --- a/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs +++ b/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs @@ -399,8 +399,9 @@ pub fn observe_drift(state: &DriftState, session_key: &str, current: StructuralH /// 16-char hex prefix of SHA-256(session_key). Bounds the log line /// width and never reveals the raw key (which may be a bearer token -/// or API key — see `derive_session_key`). -fn session_key_log_prefix(session_key: &str) -> String { +/// or API key — see `derive_session_key`). `pub(crate)` so the +/// beta-sticky merge site logs the same session identity the same way. +pub(crate) fn session_key_log_prefix(session_key: &str) -> String { let mut hasher = Sha256::new(); hasher.update(session_key.as_bytes()); let digest = hasher.finalize(); diff --git a/crates/headroom-proxy/src/cache_stabilization/mod.rs b/crates/headroom-proxy/src/cache_stabilization/mod.rs index 815c3af21..b5f91a21b 100644 --- a/crates/headroom-proxy/src/cache_stabilization/mod.rs +++ b/crates/headroom-proxy/src/cache_stabilization/mod.rs @@ -12,9 +12,14 @@ //! - **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. +//! These mutate *body* bytes only when the auth-mode gate and +//! per-policy preconditions (e.g. no customer `cache_control` +//! marker) all clear; for body mutations, OAuth and Subscription +//! always passthrough. +//! - **Re-echo** client-sent state ([`beta_sticky`]): mutate request +//! *headers* only, on every auth mode, and only ever with values +//! the same client already put on the wire — anti-drift repair of +//! the client's own signal, never injection of Headroom state. //! //! Currently shipped: //! @@ -49,6 +54,16 @@ //! `(model, system, tools)` and inject it so the upstream pins //! cache lookup to a tenant-stable identity. **Mutates the body** //! (only on PAYG) — see its docs for the gating contract. +//! - [`beta_sticky`] — parity port of the Python proxy's PR-A6 +//! `SessionBetaTracker`: per-`(provider, session)` LRU that unions +//! `anthropic-beta` / `openai-beta` tokens across turns so a client +//! dropping a token mid-conversation doesn't rotate the upstream +//! prefix-cache key. **Mutates request headers, never the body**; +//! applies to all auth modes exactly like the Python path (the +//! union only ever contains tokens this client itself sent, so +//! subscription stealth — invariant #10 "no beta drift" — is +//! preserved by construction). Operator opt-out: +//! `--beta-header-sticky disabled`. //! //! Sibling PRs hang additional submodules off this `mod.rs`. Conflict //! resolution between parallel Phase E PRs is intentionally trivial: @@ -56,6 +71,7 @@ //! `mod.rs`'s `pub mod` list. pub mod anthropic_cache_control; +pub mod beta_sticky; pub mod drift_detector; pub mod openai_cache_key; pub mod tool_def_normalize; diff --git a/crates/headroom-proxy/src/config.rs b/crates/headroom-proxy/src/config.rs index 8cc0ae593..5c2dfffc2 100644 --- a/crates/headroom-proxy/src/config.rs +++ b/crates/headroom-proxy/src/config.rs @@ -181,6 +181,48 @@ impl CompressionMode { } } +/// Session-sticky provider beta headers (parity port of the Python +/// proxy's `HEADROOM_BETA_HEADER_STICKY`; see +/// `cache_stabilization::beta_sticky`). +/// +/// When `enabled` (default), the proxy unions each request's +/// `anthropic-beta` / `openai-beta` tokens with the tokens previously +/// seen for the same conversation and forwards the union, so a client +/// dropping a beta token mid-conversation doesn't rotate the upstream +/// prefix-cache key. +/// +/// When `disabled`, the client header is forwarded verbatim and no +/// per-session token state is kept. Diagnostic operator opt-in — NOT +/// a fallback per realignment build constraint #4. +/// +/// Source priority: CLI flag → `HEADROOM_PROXY_BETA_HEADER_STICKY` +/// env var → default (`enabled`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +#[clap(rename_all = "snake_case")] +pub enum BetaHeaderSticky { + /// Union beta tokens per conversation and forward the union. + /// Default. Matches the Python proxy's default behaviour. + Enabled, + /// Forward the client's beta header verbatim; keep no state. + /// Diagnostic-only. + Disabled, +} + +impl BetaHeaderSticky { + /// Stable snake_case name suitable for log fields. + pub fn as_str(self) -> &'static str { + match self { + BetaHeaderSticky::Enabled => "enabled", + BetaHeaderSticky::Disabled => "disabled", + } + } + + /// Convenience: is the sticky union switched on? + pub fn is_enabled(self) -> bool { + matches!(self, BetaHeaderSticky::Enabled) + } +} + #[derive(Debug, Clone, Parser)] #[command( name = "headroom-proxy", @@ -320,6 +362,28 @@ pub struct CliArgs { )] pub strip_internal_headers: StripInternalHeaders, + /// Session-sticky provider beta headers: union `anthropic-beta` / + /// `openai-beta` tokens per conversation so a client dropping a + /// token mid-conversation doesn't bust the upstream prefix cache. + /// Parity port of the Python proxy's `SessionBetaTracker` (PR-A6). + /// Default `enabled`; `disabled` is a diagnostic operator opt-in. + /// + /// Active only when the compression interceptor is on + /// (`--compression` / `HEADROOM_PROXY_COMPRESSION=1`): with the + /// interceptor off the proxy is a strict byte-pipe and never + /// mutates headers. Startup logs a warning when this is `enabled` + /// while `--compression` is off. + /// + /// Source priority: CLI flag → `HEADROOM_PROXY_BETA_HEADER_STICKY` + /// env var → default (`enabled`). + #[arg( + long = "beta-header-sticky", + env = "HEADROOM_PROXY_BETA_HEADER_STICKY", + value_enum, + default_value_t = BetaHeaderSticky::Enabled, + )] + pub beta_header_sticky: BetaHeaderSticky, + /// Phase C PR-C4: enable the `/v1/responses` SSE streaming /// pipeline. When `true` (default), `Accept: text/event-stream` /// requests on `/v1/responses` flow through the byte-level SSE @@ -517,6 +581,9 @@ pub struct Config { /// upstream-bound requests. PR-A5 default-on guard against /// fingerprinting / leakage of internal flags. pub strip_internal_headers: StripInternalHeaders, + /// Session-sticky provider beta headers (parity port of the + /// Python `SessionBetaTracker`, PR-A6). Default `enabled`. + pub beta_header_sticky: BetaHeaderSticky, /// PR-C4: enable the `/v1/responses` streaming pipeline (SSE /// state-machine + telemetry tee). Default `true`. pub enable_responses_streaming: bool, @@ -578,6 +645,7 @@ impl Config { cache_control_auto_frozen: args.cache_control_auto_frozen, auth_mode_policy_enforcement: args.auth_mode_policy_enforcement, strip_internal_headers: args.strip_internal_headers, + beta_header_sticky: args.beta_header_sticky, enable_responses_streaming: args.enable_responses_streaming, enable_conversations_passthrough: args.enable_conversations_passthrough, enable_bedrock_native: args.enable_bedrock_native, @@ -621,6 +689,9 @@ impl Config { // from upstream-bound requests. Tests opt out per-case via // `start_proxy_with`. strip_internal_headers: StripInternalHeaders::Enabled, + // Production default: sticky beta-header union per + // conversation (Python-parity). Tests opt out per-case. + beta_header_sticky: BetaHeaderSticky::Enabled, // PR-C4: streaming pipeline + conversations passthrough // both default-on so tests exercise the same paths // production traffic will hit. diff --git a/crates/headroom-proxy/src/main.rs b/crates/headroom-proxy/src/main.rs index 19e748522..3a19d98ab 100644 --- a/crates/headroom-proxy/src/main.rs +++ b/crates/headroom-proxy/src/main.rs @@ -31,6 +31,22 @@ async fn main() -> Result<(), Box> { "headroom-proxy starting" ); + // Session-sticky beta headers only run inside the compression + // interceptor: with `--compression` off the proxy is a strict + // byte-pipe and never mutates headers. Say so loudly at startup — + // an operator reading `beta_header_sticky=enabled` (the default) + // must not believe the protection is active when it isn't. + if config.beta_header_sticky.is_enabled() && !config.compression { + tracing::warn!( + event = "beta_header_sticky_inactive", + beta_header_sticky = config.beta_header_sticky.as_str(), + compression = config.compression, + "beta-header stickiness is enabled but the compression \ + interceptor is off; enable --compression (or \ + HEADROOM_PROXY_COMPRESSION=1) to activate it" + ); + } + let mut state = AppState::new(config.clone())?; // PR-D1: resolve AWS credentials at startup via the `aws-config` diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index 7f3d81cd3..4c14f9900 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -17,6 +17,7 @@ use futures_util::{StreamExt as _, TryStreamExt}; use http_body_util::BodyExt; use crate::cache_stabilization; +use crate::cache_stabilization::beta_sticky::BetaProvider; use crate::cache_stabilization::drift_detector::{ compute_structural_hash, derive_session_key, observe_drift, ApiKind, DriftState, }; @@ -66,6 +67,13 @@ pub struct AppState { /// request body — so this can be cloned freely into every handler /// path that buffers the body. pub drift_state: DriftState, + /// Session-sticky beta-header tracker (parity port of the Python + /// `SessionBetaTracker`, PR-A6): per-`(provider, session)` LRU of + /// `anthropic-beta` / `openai-beta` tokens, unioned across turns + /// so a client dropping a token mid-conversation doesn't rotate + /// the upstream prefix-cache key. Shares the drift detector's + /// session identity (same `derive_session_key` output). + pub beta_sticky: cache_stabilization::beta_sticky::BetaStickyState, /// PR-D4: GCP ADC bearer-token source for Vertex routes. Default: /// [`crate::vertex::adc::GcpAdcTokenSource`] constructed lazily; /// the actual ADC chain is only resolved when the first Vertex @@ -111,6 +119,9 @@ impl AppState { client, bedrock_credentials: None, drift_state: DriftState::new(DRIFT_DETECTOR_CAPACITY), + beta_sticky: cache_stabilization::beta_sticky::BetaStickyState::new( + cache_stabilization::beta_sticky::BETA_TRACKER_CAPACITY, + ), vertex_token_source, }) } @@ -707,6 +718,41 @@ pub(crate) async fn forward_http( let session_key = derive_session_key(headers, &client_addr, &parsed, kind); let hash = compute_structural_hash(&parsed, kind); observe_drift(&state.drift_state, &session_key, hash); + + // Session-sticky provider beta headers — port of the + // Python PR-A6 `SessionBetaTracker`. Beta headers are + // part of the bytes that determine the upstream + // prefix-cache key; a client dropping a token between + // turns rotates the key and re-writes the whole + // prefix at the customer's cost. Forward the + // per-conversation union instead. See + // `cache_stabilization::beta_sticky` for the behavior + // contract, the auth-mode rationale (applies to every + // mode, like the Python handler), and the one + // documented divergence from Python (per-conversation + // keying). Reuses the drift detector's `session_key` + // so both cache-stability subsystems agree on + // conversation identity. Mutates upstream-bound + // HEADERS only; body bytes stay untouched (Phase-A + // cache-safety invariant). + if state.config.beta_header_sticky.is_enabled() { + let provider = match endpoint { + compression::CompressibleEndpoint::AnthropicMessages => { + BetaProvider::Anthropic + } + compression::CompressibleEndpoint::OpenAiChatCompletions + | compression::CompressibleEndpoint::OpenAiResponses => { + BetaProvider::OpenAi + } + }; + cache_stabilization::beta_sticky::apply_sticky_betas( + &state.beta_sticky, + provider, + &session_key, + &mut outgoing_headers, + &request_id, + ); + } } } let outcome = match endpoint { diff --git a/crates/headroom-proxy/tests/integration_beta_header_sticky.rs b/crates/headroom-proxy/tests/integration_beta_header_sticky.rs new file mode 100644 index 000000000..d1712be96 --- /dev/null +++ b/crates/headroom-proxy/tests/integration_beta_header_sticky.rs @@ -0,0 +1,529 @@ +//! End-to-end coverage for session-sticky provider beta headers +//! (`cache_stabilization::beta_sticky` — Rust port of the Python +//! proxy's PR-A6 `SessionBetaTracker`). +//! +//! The scenario every test guards: a client (Claude Code, Codex CLI) +//! sends `anthropic-beta: a,b` on turn 1 and drops `b` on turn 2 of +//! the SAME conversation. Beta headers are part of the bytes that +//! determine the upstream prefix-cache key, so the drop rotates the +//! key and the provider re-writes the whole prefix at the customer's +//! cost. The proxy must forward the per-conversation union instead. +//! +//! These tests boot a real Rust proxy in front of a wiremock upstream +//! and assert on the headers/bytes the upstream actually receives: +//! +//! - dropped tokens are re-injected on later turns (Anthropic, +//! OpenAI Chat, OpenAI Responses — all three intercepted routes); +//! - conversation identity works both via the explicit +//! `x-headroom-session-id` opt-in AND via the body-derived +//! conversation discriminator (no explicit header — the realistic +//! Claude Code shape); +//! - the union NEVER invents tokens the client didn't send: no beta +//! header in → no beta header out, and separate conversations don't +//! leak tokens into each other; +//! - `--beta-header-sticky disabled` forwards the client value +//! verbatim (diagnostic opt-out, Python +//! `HEADROOM_BETA_HEADER_STICKY=disabled` parity); +//! - the body is forwarded byte-equal (SHA-256) while the header is +//! rewritten — the mechanism mutates request headers, never body +//! bytes (Phase-A cache-safety contract). + +mod common; + +use common::start_proxy_with; +use serde_json::json; +use sha2::{Digest, Sha256}; +use std::sync::{Arc, Mutex}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +/// Everything the upstream saw for one request: selected header +/// values (lower-case names) + raw body bytes. +#[derive(Clone)] +struct Seen { + beta: Option, + session_id_header: Option, + body: Vec, +} + +type Captures = Arc>>; + +/// Mount a capture-everything mock for `route` on the upstream. The +/// `beta_header` name is which provider beta header to record +/// (`anthropic-beta` / `openai-beta`). +async fn mount_capture(upstream: &MockServer, route: &str, beta_header: &'static str) -> Captures { + let captured: Captures = Arc::new(Mutex::new(Vec::new())); + let captured_clone = captured.clone(); + Mock::given(method("POST")) + .and(path(route)) + .respond_with(move |req: &wiremock::Request| { + let get = |name: &str| { + req.headers + .get(name) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + }; + captured_clone.lock().unwrap().push(Seen { + beta: get(beta_header), + session_id_header: get("x-headroom-session-id"), + body: req.body.clone(), + }); + ResponseTemplate::new(200).set_body_string(r#"{"ok":true}"#) + }) + .mount(upstream) + .await; + captured +} + +fn anthropic_body(turns: &[(&str, &str)]) -> Vec { + let messages: Vec = turns + .iter() + .map(|(role, content)| json!({"role": role, "content": content})) + .collect(); + serde_json::to_vec(&json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 32, + "messages": messages, + })) + .unwrap() +} + +fn openai_chat_body(turns: &[(&str, &str)]) -> Vec { + let messages: Vec = turns + .iter() + .map(|(role, content)| json!({"role": role, "content": content})) + .collect(); + serde_json::to_vec(&json!({ + "model": "gpt-4o", + "messages": messages, + })) + .unwrap() +} + +fn openai_responses_body(text: &str) -> Vec { + serde_json::to_vec(&json!({ + "model": "gpt-4o", + "input": [{"role": "user", "content": text}], + })) + .unwrap() +} + +async fn post( + client: &reqwest::Client, + url: String, + body: Vec, + headers: &[(&str, &str)], +) -> reqwest::Response { + let mut req = client + .post(url) + .header("content-type", "application/json") + .body(body); + for (name, value) in headers { + req = req.header(*name, *value); + } + req.send().await.expect("proxy reachable") +} + +#[tokio::test] +async fn anthropic_dropped_beta_token_reinjected_with_explicit_session_header() { + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + // Turn 1: two beta tokens. + let resp = post( + &client, + url.clone(), + anthropic_body(&[("user", "hello")]), + &[ + ( + "anthropic-beta", + "context-management-2025-06-27,interleaved-thinking-2025-05-14", + ), + ("x-headroom-session-id", "conv-explicit-1"), + ], + ) + .await; + assert_eq!(resp.status(), 200); + + // Turn 2, same conversation: the client dropped the second token. + let resp = post( + &client, + url, + anthropic_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]), + &[ + ("anthropic-beta", "context-management-2025-06-27"), + ("x-headroom-session-id", "conv-explicit-1"), + ], + ) + .await; + assert_eq!(resp.status(), 200); + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[0].beta.as_deref(), + Some("context-management-2025-06-27,interleaved-thinking-2025-05-14"), + "turn 1 forwards the client value unchanged" + ); + assert_eq!( + seen[1].beta.as_deref(), + Some("context-management-2025-06-27,interleaved-thinking-2025-05-14"), + "turn 2 must re-inject the dropped token so the upstream \ + prefix-cache key stays byte-stable" + ); + // PR-A5 invariant intact: the internal session header never + // crosses the upstream boundary. + assert!(seen.iter().all(|s| s.session_id_header.is_none())); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn anthropic_conversation_keyed_without_explicit_session_header() { + // The realistic Claude Code shape: no `x-headroom-session-id`; + // conversation identity comes from the credential arm + the + // first-message discriminator inside `derive_session_key`. + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + let auth = ("authorization", "Bearer oauth-workspace-token"); + post( + &client, + url.clone(), + anthropic_body(&[("user", "conversation opener")]), + &[("anthropic-beta", "a,b"), auth], + ) + .await; + // Same conversation (same opener, grown transcript), token "b" + // dropped. + post( + &client, + url, + anthropic_body(&[ + ("user", "conversation opener"), + ("assistant", "reply"), + ("user", "follow-up"), + ]), + &[("anthropic-beta", "a"), auth], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[1].beta.as_deref(), Some("a,b")); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn openai_chat_dropped_beta_token_reinjected() { + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/chat/completions", "openai-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/chat/completions", proxy.url()); + + post( + &client, + url.clone(), + openai_chat_body(&[("user", "hello")]), + &[ + ("openai-beta", "assistants=v2,realtime=v1"), + ("x-headroom-session-id", "conv-oai-1"), + ], + ) + .await; + post( + &client, + url, + openai_chat_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]), + &[ + ("openai-beta", "assistants=v2"), + ("x-headroom-session-id", "conv-oai-1"), + ], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[1].beta.as_deref(), Some("assistants=v2,realtime=v1")); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn openai_responses_dropped_beta_token_reinjected() { + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/responses", "openai-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/responses", proxy.url()); + + post( + &client, + url.clone(), + openai_responses_body("hello"), + &[ + ("openai-beta", "responses=v1,tools=v2"), + ("x-headroom-session-id", "conv-resp-1"), + ], + ) + .await; + post( + &client, + url, + openai_responses_body("hello again"), + &[ + ("openai-beta", "responses=v1"), + ("x-headroom-session-id", "conv-resp-1"), + ], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[1].beta.as_deref(), Some("responses=v1,tools=v2")); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn anthropic_fully_omitted_beta_header_regains_union() { + // The headline docs claim: "sends a token in turn N and omits it + // in turn N+1" — here the whole header disappears, not just one + // token, and the union must be re-added through real axum/reqwest + // plumbing. + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + post( + &client, + url.clone(), + anthropic_body(&[("user", "hello")]), + &[ + ("anthropic-beta", "context-management-2025-06-27"), + ("x-headroom-session-id", "conv-omit-1"), + ], + ) + .await; + // Turn 2: no anthropic-beta header at all. + post( + &client, + url, + anthropic_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]), + &[("x-headroom-session-id", "conv-omit-1")], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[1].beta.as_deref(), + Some("context-management-2025-06-27"), + "a fully omitted beta header must be restored from session state" + ); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn disabled_flag_forwards_client_value_verbatim() { + use headroom_proxy::config::BetaHeaderSticky; + + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + c.beta_header_sticky = BetaHeaderSticky::Disabled; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + post( + &client, + url.clone(), + anthropic_body(&[("user", "hello")]), + &[ + ("anthropic-beta", "a,b"), + ("x-headroom-session-id", "conv-d1"), + ], + ) + .await; + post( + &client, + url, + anthropic_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]), + &[ + ("anthropic-beta", "a"), + ("x-headroom-session-id", "conv-d1"), + ], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + seen[1].beta.as_deref(), + Some("a"), + "disabled mode must forward the dropped-token value verbatim \ + and keep no session state" + ); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn no_client_beta_header_is_never_invented() { + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + for body in [ + anthropic_body(&[("user", "hello")]), + anthropic_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]), + ] { + post( + &client, + url.clone(), + body, + &[("x-headroom-session-id", "conv-n1")], + ) + .await; + } + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert!( + seen.iter().all(|s| s.beta.is_none()), + "a session that never sent a beta header must never gain one" + ); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn separate_conversations_do_not_leak_tokens() { + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + post( + &client, + url.clone(), + anthropic_body(&[("user", "conversation A")]), + &[ + ("anthropic-beta", "token-a"), + ("x-headroom-session-id", "conv-A"), + ], + ) + .await; + post( + &client, + url, + anthropic_body(&[("user", "conversation B")]), + &[ + ("anthropic-beta", "token-b"), + ("x-headroom-session-id", "conv-B"), + ], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[0].beta.as_deref(), Some("token-a")); + assert_eq!( + seen[1].beta.as_deref(), + Some("token-b"), + "conversation B must not inherit conversation A's tokens" + ); + + proxy.shutdown().await; +} + +#[tokio::test] +async fn body_bytes_stay_byte_equal_while_header_is_rewritten() { + // Cache-safety contract: the sticky union mutates request + // HEADERS only. The forwarded body must remain byte-identical + // (SHA-256) to what the client sent — same assertion idiom as the + // model-sanitizer integration tests. + let upstream = MockServer::start().await; + let captured = mount_capture(&upstream, "/v1/messages", "anthropic-beta").await; + let proxy = start_proxy_with(&upstream.uri(), |c| { + c.compression = true; + }) + .await; + let client = reqwest::Client::new(); + let url = format!("{}/v1/messages", proxy.url()); + + let turn1 = anthropic_body(&[("user", "hello")]); + let turn2 = anthropic_body(&[("user", "hello"), ("assistant", "hi"), ("user", "next")]); + + post( + &client, + url.clone(), + turn1, + &[ + ("anthropic-beta", "a,b"), + ("x-headroom-session-id", "conv-bb"), + ], + ) + .await; + post( + &client, + url, + turn2.clone(), + &[ + ("anthropic-beta", "a"), + ("x-headroom-session-id", "conv-bb"), + ], + ) + .await; + + let seen = captured.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + // Header was rewritten to the union… + assert_eq!(seen[1].beta.as_deref(), Some("a,b")); + // …but the body bytes are untouched. + assert_eq!( + Sha256::digest(&seen[1].body), + Sha256::digest(&turn2), + "sticky beta union must never mutate body bytes" + ); + + proxy.shutdown().await; +} diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index 40ee56530..f4d2956c1 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -294,6 +294,7 @@ headroom proxy --learn --min-evidence 3 | `HEADROOM_REQUEST_TIMEOUT` | Request timeout in seconds | `300` | | `HEADROOM_BETA_HEADER_STICKY` | Controls per-session `anthropic-beta` / `OpenAI-Beta` re-echo. `enabled` (default): the proxy unions beta tokens across turns within a session — if the client sends a token in turn N and omits it in turn N+1, the proxy re-injects it to preserve prefix-cache stability. `disabled`: the client's value is forwarded verbatim with no accumulation. Any other value raises at request time. See [Session Beta Header Tracking](/docs/configuration#session-beta-header-tracking). | `enabled` | | `HEADROOM_BETA_TRACKER_MAX_SESSIONS` | LRU capacity of the in-memory session beta tracker. Once full, the oldest session entry is evicted. | `1000` | +| `HEADROOM_PROXY_BETA_HEADER_STICKY` | Rust proxy: same per-conversation beta-token union as `HEADROOM_BETA_HEADER_STICKY`, applied to `anthropic-beta` / `openai-beta` on the intercepted `/v1/messages`, `/v1/chat/completions`, and `/v1/responses` routes. Requires the compression interceptor (`HEADROOM_PROXY_COMPRESSION=1`) — with it off the Rust proxy is a strict byte-pipe and this flag has no effect (startup warns). Unlike the Python tracker (keyed on model + system prompt), sessions are keyed per conversation, shared with the cache-drift detector — parallel conversations never inherit each other's tokens. `enabled` default; `disabled` forwards the client value verbatim and keeps no state. Tracker capacity is fixed at 1000 sessions. | `enabled` | | `HEADROOM_MODEL_ROUTER_ENABLED` | Enable cost-aware model routing. `1`/`true`/`yes`/`on`/`enabled` turns it on and requires `HEADROOM_MODEL_ROUTES`. See [Cost-aware model routing](/docs/configuration#cost-aware-model-routing). | `off` | | `HEADROOM_MODEL_ROUTES` | JSON array of ordered routing rules for cost-aware model routing (schema below). | -- | | `HEADROOM_THINKING_COMPACT` | Compact plain-text reasoning that models re-send every turn (Kimi/GLM/DeepSeek `reasoning_content` / inline ``): Kompress it on warm turns, drop it on cold turns. No-op for Claude/Codex/OpenAI (encrypted reasoning). See [Cold-prefix hook](#cold-prefix-hook--reasoning-compaction). | `off` |