From 6744833afed9aba7bdadd2aa59ed3a6fa99afc15 Mon Sep 17 00:00:00 2001 From: Andrei Boldyrev <53853163+abcgco@users.noreply.github.com> Date: Fri, 17 Jul 2026 02:34:26 +0500 Subject: [PATCH] fix(proxy): key drift detector on conversations, not credentials; canonicalize drift hashes (#2301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description The Rust proxy's cache-bust drift detector (`crates/headroom-proxy/src/cache_stabilization/drift_detector.rs`, PR-E6) cannot currently tell drift from normal operation on interactive agentic traffic, so it warns on nearly every turn and a real bust drowns in the noise. Three compounding defects, all verified against live Claude Code traffic: 1. `derive_session_key` stops at the credential hash — Claude Code sends one OAuth bearer for every conversation, so all concurrent conversations share one LRU slot and every conversation switch logs a false `cache_drift_observed` (with `drift_dims` computed against the wrong conversation's baseline). 2. The `early_messages` axis hashes the raw first-3-messages window, so a lone conversation's normal growth (1 → 3 messages) and the client relocating its `cache_control` breakpoint to the newest block both fire a false `early_messages` drift at turn 2–3 of essentially every session. 3. `x-headroom-session-id` — the explicit session identity the Python proxy honors everywhere session-sticky state exists — is ignored on the Rust path. This PR makes the detector's session identity conversation-scoped and its comparison canonical, the same shape as the merged Python-side fix for #2085 (`SessionTrackerStore.resolve_tracker` lineage resolution + `_canonicalize_for_prefix_compare`): - **`derive_session_key`**: honors `x-headroom-session-id` first (hashed, like every other key input), then folds a conversation discriminator into the credential/network arms: a 16-hex-char SHA-256 fingerprint of `(model, canonicalized first message)`. Provider prompt caches are per-model, so a small-model sidecar call (title generation) that reuses a conversation's opener stays a separate session instead of false-drifting on `system`. - **`canonicalize_for_hash`** on all axes and the discriminator: objects rebuilt with sorted keys (this workspace enables serde_json `preserve_order`, so a plain re-serialize would keep client wire order and leave the hashes key-order sensitive) and `cache_control` stripped outside opaque tool payloads (`input`/`arguments`/`json`/`input_schema` — mirroring the Python canonicalizer's `_OPAQUE_PAYLOAD_KEYS`, so a user field that happens to be *named* `cache_control` still counts as drift). - **`early_messages`** becomes per-message hashes (`[Option<[u8; 32]>; 3]`) with a prefix-aware comparison: growing into the window is benign; a settled message changing or disappearing under a stable session key is still drift. `observe_drift` now gates the warning on drifted dimensions rather than raw hash inequality. True positives are preserved (`system`/`tools` changes, in-place history rewrites under a pinned identity), and the detector remains a pure observer — no forwarded byte changes, `does_not_mutate_input` still pins that. **Documented trade-off** (module doc + `conversation_discriminator` doc): without the explicit header, a client that rewrites its first message (history compaction, rolling-window truncation, Responses chained mode) re-keys to a fresh session — the rewrite surfaces as `cache_drift_first_request` rather than `cache_drift_observed` against the old baseline. That is deliberate: the credential-keyed alternative false-warned on every conversation switch, which buried those same events anyway. `x-headroom-session-id` pins the identity and reports rewrites as drift. Byte-identical openers on the same model under one credential still conflate (rare; documented). Closes #2300 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `derive_session_key`: `x-headroom-session-id` (hashed) wins; credential/IP arms fold in `conversation_discriminator` — `(model, canonicalized first message)`, 16 hex chars - New `canonicalize_for_hash`: sorted-key object rebuild + `cache_control` stripped outside `OPAQUE_PAYLOAD_KEYS`; applied to the `system`/`tools`/`early_messages` axes and the discriminator - `StructuralHash.early_messages`: `[u8; 32]` → `[Option<[u8; 32]>; EARLY_MESSAGES_WINDOW]` per-message hashes; `drift_dims`/`early_window_drifted` implement the prefix-aware rule; `observe_drift` warns on non-empty dims instead of `!=` - `conversation_messages` shape guard: bare-string message containers only count for the Responses `input` sugar - Docs: module header (canonicalization, trade-off, honest cost), `conversation_discriminator` rationale + blind spots, `DRIFT_DETECTOR_CAPACITY` cardinality note (per-conversation keys, 163-byte entry), `structural_hash_log_prefix` hex-length fix - Tests: 13 new unit tests (conversation separation, turn-growth key stability, explicit header priority, marker relocation + growth not drift, rewrite/shrink still drift, per-model separation, key-order neutrality, opaque-payload fields still count, Responses/Chat discriminator shapes, string-container gating) - `CHANGELOG.md`: Unreleased → Fixed entry ## Testing - [x] Unit tests pass (`cargo test -p headroom-proxy` — full crate: lib + integration suites) - [x] Linting passes (`cargo clippy -p headroom-proxy --all-targets` — zero warnings; `cargo fmt --check` clean) - [ ] Type checking passes (`mypy headroom`) — n/a, no Python files touched - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ cargo test -p headroom-proxy --lib drift_detector test result: ok. 27 passed; 0 failed; 0 ignored; 0 measured; 221 filtered out $ cargo test -p headroom-proxy (all suites) test result: ok. 248 passed (lib) + integration suites, 0 failed $ cargo clippy -p headroom-proxy --all-targets (no warnings) ``` ## Real Behavior Proof - Environment: macOS 15 (arm64), rustc 1.95.0, repo @ 718c8dc + this branch - Exact command / steps: captured two real Claude Code conversations × two turns through a local proxy (`ANTHROPIC_BASE_URL=http://localhost:8791 claude -p …` / `--resume …`), rebuilt the wire bodies, and replayed them through the real `derive_session_key` / `compute_structural_hash` / `drift_dims` in a local `cargo test` harness — before and after this change. - Observed result: **before** — all four requests share one `auth:` key, and the raw early-window hash flips between turn 1 and turn 2 of the *same* conversation (false `early_messages` drift; interleaving also flips `system`). **After** — turn 1/turn 2 map to one stable key with `drift_dims == ""`, the two conversations map to distinct keys, and a rewritten/shrunk settled window still reports `early_messages`. - Not tested: live OpenAI Chat/Responses traffic (shape-level unit tests only); log pipeline consumers (event names/fields unchanged). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [x] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) n/a — log-only telemetry change. ## Additional Notes - `StructuralHash` is `pub`, but the workspace has no external consumers (checked `sdk/`, `plugins/`, Python, docs) — the field-type change is contained to `proxy.rs` and the module tests. `[Option<[u8; 32]>; 3]` keeps `Copy` for the LRU and adds no dependency. - LRU cardinality: keys moved per-credential → per-conversation; `DRIFT_DETECTOR_CAPACITY`'s comment now documents the working set, the ~250-byte entry, and the graceful eviction failure mode (repeated `cache_drift_first_request`, telemetry-only). - Not in scope, noted for follow-up: keying Responses chained mode (`previous_response_id`) as a lineage; surfacing mid-history `role:"system"` insertions on the OpenAI Chat shape (pre-existing blind spot on all axes). --------- Co-authored-by: Tejas Chopra --- .../src/cache_stabilization/drift_detector.rs | 688 ++++++++++++++++-- .../src/cache_stabilization/mod.rs | 3 +- crates/headroom-proxy/src/proxy.rs | 17 +- 3 files changed, 621 insertions(+), 87 deletions(-) diff --git a/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs b/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs index 1302ad3a4..c288d3131 100644 --- a/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs +++ b/crates/headroom-proxy/src/cache_stabilization/drift_detector.rs @@ -9,31 +9,63 @@ //! `body.system`; OpenAI Chat first `role=system` message; //! OpenAI Responses `body.instructions`). //! - `tools` — SHA-256 of the canonical bytes of `body.tools`. -//! - `early_messages` — SHA-256 of the canonical bytes of the first 3 -//! message-shaped items (or all, if fewer than 3). Skips the +//! - `early_messages` — per-message SHA-256 of the first 3 +//! conversation messages (or all, if fewer than 3). Skips the //! live-zone tail where mutation is expected and benign. //! -//! Track the previous hash per session in a bounded LRU. When a +//! All axes are canonicalized before hashing ([`canonicalize_for_hash`]): +//! objects are rebuilt with sorted keys (whitespace- and key-order- +//! neutral — the workspace's `preserve_order` feature would otherwise +//! keep wire order in the hash) and `cache_control` members are +//! stripped outside opaque tool payloads — clients relocate cache +//! breakpoints to the newest block every turn, and moving a breakpoint +//! never invalidates a previously cached prefix, so markers are +//! placement metadata rather than structure. +//! +//! Track the previous fingerprint per session in a bounded LRU. When a //! subsequent request on the same session disagrees on any dimension, //! emit a `cache_drift_observed` log line listing the drifted -//! dimensions. **Never mutates the request body** — the detector is a -//! pure observer and the proxy's "passthrough is sacred" invariant -//! (Phase A) is preserved by construction. +//! dimensions. The `early_messages` comparison is prefix-aware: a +//! conversation *growing into* the window (turn 2 appends messages +//! after turn 1's) is append-only and benign, while a settled early +//! message that changes or disappears busted the provider's prefix +//! cache and is drift. **Never mutates the request body** — the +//! detector is a pure observer and the proxy's "passthrough is sacred" +//! invariant (Phase A) is preserved by construction. +//! +//! Known trade-off: without an explicit `x-headroom-session-id`, the +//! session identity is anchored on the conversation's first message +//! (see [`conversation_discriminator`]), so a client that *rewrites* +//! that message (history compaction, rolling-window truncation) +//! re-keys to a fresh session and the rewrite surfaces as +//! `cache_drift_first_request` rather than `cache_drift_observed` — +//! traded deliberately against the credential-keyed alternative, which +//! false-warned on every conversation switch. With the explicit header +//! the identity is pinned and rewrites are reported as drift. //! //! # Privacy //! -//! The session key is derived from the strongest available client -//! identifier (`Authorization`, `x-api-key`, client IP, finally -//! `(client_ip, user_agent)`). Bearer tokens and API keys are -//! **hashed before they ever leave this module**; the raw secret is -//! never logged, never stored, and is overwritten in transit (truncated -//! to a 16-character hex prefix). The log line itself only includes a -//! short prefix of the SHA-256 hex of the session key. +//! The session key prefers the client's explicit +//! `x-headroom-session-id` header — the same opt-in the Python proxy +//! honors for all session-sticky state — and otherwise combines the +//! strongest available client identifier (`Authorization`, `x-api-key`, +//! client IP, finally `(client_ip, user_agent)`) with a fingerprint of +//! the conversation's first message, so concurrent conversations that +//! share one credential do not share one drift session. Bearer tokens, +//! API keys, and the session-id header value are **hashed before they +//! ever leave this module**; the raw secret is never logged and never +//! stored. The conversation fingerprint is likewise a truncated SHA-256 +//! — no message content appears in the key. The log line itself only +//! includes a short prefix of the SHA-256 hex of the session key. //! //! # Cost //! -//! - One SHA-256 update over each of (system, tools, early messages). -//! Total ~200us on a 8 KB system prompt. +//! - Up to six SHA-256 digests per request (system, tools, up to +//! three early messages, and the session key's conversation +//! fingerprint), each over a canonicalized clone (filtered, +//! key-sorted rebuild) of the corresponding subtree. Still well +//! under a millisecond on a typical agentic request; the detector +//! stays log-only and off the forwarded-bytes path. //! - One LRU lookup + insert. `lru = "0.12"` is O(1) amortised. //! - One `tracing::info!` or `tracing::warn!`. No metric emission yet //! (left for Phase F PR-F* when the global Prometheus registry can @@ -66,22 +98,27 @@ pub enum ApiKind { /// Three-axis structural fingerprint of the cache hot zone. /// -/// Each axis is the SHA-256 of the canonical bytes at that position -/// (we re-serialize via `serde_json::to_vec` so whitespace and key -/// order through the original network bytes do not perturb the hash). -/// All three are required for "no drift"; any one differing flags -/// drift on that dimension. +/// Each axis is the SHA-256 of canonical bytes at that position (see +/// [`canonicalize_for_hash`]): all three axes must be *stable* for +/// "no drift", and each drifting axis is named in the emitted event. +/// +/// `early_messages` holds one hash per settled slot of the early +/// window (`None` = the conversation hasn't grown that far yet), so +/// the comparison can tell "grew into the window" (benign) apart from +/// "a settled message changed" (drift). Note that the derived `==` is +/// stricter than the drift predicate — growth compares unequal but is +/// not drift; use [`drift_dims`]'s emptiness for drift decisions. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct StructuralHash { pub system: [u8; 32], pub tools: [u8; 32], - pub early_messages: [u8; 32], + pub early_messages: [Option<[u8; 32]>; EARLY_MESSAGES_WINDOW], } /// How many message-shaped items count as the "early" prefix that -/// feeds `early_messages_hash`. Anything past this is the live zone +/// feeds [`early_message_hashes`]. Anything past this is the live zone /// (where mutation is expected; we deliberately ignore it). -const EARLY_MESSAGES_WINDOW: usize = 3; +pub const EARLY_MESSAGES_WINDOW: usize = 3; /// Compute a [`StructuralHash`] for the body shape implied by `kind`. /// @@ -89,9 +126,9 @@ const EARLY_MESSAGES_WINDOW: usize = 3; /// `does_not_mutate_input` test in the module below pins this with a /// clone-and-compare assertion. pub fn compute_structural_hash(body: &serde_json::Value, kind: ApiKind) -> StructuralHash { - let system = hash_value(&extract_system(body, kind)); - let tools = hash_value(&extract_tools(body)); - let early_messages = hash_value(&extract_early_messages(body, kind)); + let system = hash_value(&canonicalize_for_hash(&extract_system(body, kind), false)); + let tools = hash_value(&canonicalize_for_hash(&extract_tools(body), false)); + let early_messages = early_message_hashes(body, kind); StructuralHash { system, tools, @@ -99,6 +136,60 @@ pub fn compute_structural_hash(body: &serde_json::Value, kind: ApiKind) -> Struc } } +/// Keys whose values are opaque user payloads: tool-call arguments and +/// schemas. Never strip inside them — a user field that happens to be +/// named `cache_control` is data, and dropping it would make two +/// genuinely different payloads hash identically, masking real drift. +/// `input`/`arguments`/`json` mirror the Python canonicalizer's +/// `_OPAQUE_PAYLOAD_KEYS`; `input_schema` extends the same rule to +/// tool definitions, which the Python comparator never hashes but the +/// `tools` axis here does. +const OPAQUE_PAYLOAD_KEYS: [&str; 4] = ["input", "arguments", "json", "input_schema"]; + +/// Canonicalize a JSON tree for hashing: rebuild every object with +/// sorted keys, dropping `cache_control` members outside opaque +/// payloads. +/// +/// Key sorting is what makes the hashes genuinely key-order neutral: +/// this workspace builds `serde_json` with `preserve_order`, so a +/// plain re-serialize would keep the client's wire order and a +/// serializer-side reordering would read as drift (and would rotate +/// the conversation fingerprint). Sorting is a pure reordering — no +/// information is lost, so distinct payloads never conflate. +/// +/// Cache-breakpoint markers are placement metadata, not structure: +/// clients relocate them to the newest block every turn (observed live +/// from Claude Code), and moving a breakpoint never invalidates a +/// previously cached prefix. Hashing them would flag drift on every +/// relocation. +fn canonicalize_for_hash(value: &serde_json::Value, in_opaque: bool) -> serde_json::Value { + match value { + serde_json::Value::Object(map) => { + let mut entries: Vec<(&String, &serde_json::Value)> = map + .iter() + .filter(|(k, _)| in_opaque || k.as_str() != "cache_control") + .collect(); + entries.sort_by_key(|(k, _)| k.as_str()); + serde_json::Value::Object( + entries + .into_iter() + .map(|(k, v)| { + let opaque = in_opaque || OPAQUE_PAYLOAD_KEYS.contains(&k.as_str()); + (k.clone(), canonicalize_for_hash(v, opaque)) + }) + .collect(), + ) + } + serde_json::Value::Array(items) => serde_json::Value::Array( + items + .iter() + .map(|v| canonicalize_for_hash(v, in_opaque)) + .collect(), + ), + other => other.clone(), + } +} + /// Extract the "system" axis as a `serde_json::Value`. Returns /// `Value::Null` when the dimension is absent — Null still hashes to /// a stable 32-byte digest so first-request comparisons are @@ -142,38 +233,54 @@ fn extract_tools(body: &serde_json::Value) -> serde_json::Value { .unwrap_or(serde_json::Value::Null) } -/// Extract the first [`EARLY_MESSAGES_WINDOW`] message-shaped items -/// as an array `Value`. Skips the system message in the OpenAI Chat -/// shape (the system axis already hashes that separately). -fn extract_early_messages(body: &serde_json::Value, kind: ApiKind) -> serde_json::Value { +/// Collect the conversation-scoped message items for `kind`: Anthropic +/// `messages`, OpenAI Chat `messages` minus `role:"system"` entries +/// (the system axis already hashes those separately), OpenAI Responses +/// `input`. A Responses string-form `input` counts as one item. +fn conversation_messages(body: &serde_json::Value, kind: ApiKind) -> Vec<&serde_json::Value> { let array_key = match kind { ApiKind::Anthropic => "messages", ApiKind::OpenAiChat => "messages", ApiKind::OpenAiResponses => "input", }; - let messages = match body.get(array_key).and_then(|v| v.as_array()) { - Some(arr) => arr, - None => return serde_json::Value::Null, + let items = match (kind, body.get(array_key)) { + (_, Some(serde_json::Value::Array(arr))) => arr.iter().collect::>(), + // `input` may be a bare string in the Responses API. The other + // shapes only accept arrays; anything else is malformed and + // will be rejected upstream, so contribute nothing here. + (ApiKind::OpenAiResponses, Some(s @ serde_json::Value::String(_))) => vec![s], + _ => return Vec::new(), }; - let early: Vec = match kind { - ApiKind::OpenAiChat => messages - .iter() + match kind { + ApiKind::OpenAiChat => items + .into_iter() .filter(|m| { m.get("role") .and_then(|r| r.as_str()) .map(|s| s != "system") .unwrap_or(true) }) - .take(EARLY_MESSAGES_WINDOW) - .cloned() .collect(), - _ => messages - .iter() - .take(EARLY_MESSAGES_WINDOW) - .cloned() - .collect(), - }; - serde_json::Value::Array(early) + _ => items, + } +} + +/// Hash each of the first [`EARLY_MESSAGES_WINDOW`] conversation +/// messages individually (canonicalized). Slots the conversation has +/// not grown into yet stay `None`. +fn early_message_hashes( + body: &serde_json::Value, + kind: ApiKind, +) -> [Option<[u8; 32]>; EARLY_MESSAGES_WINDOW] { + let mut out = [None; EARLY_MESSAGES_WINDOW]; + for (slot, msg) in conversation_messages(body, kind) + .into_iter() + .take(EARLY_MESSAGES_WINDOW) + .enumerate() + { + out[slot] = Some(hash_value(&canonicalize_for_hash(msg, false))); + } + out } /// SHA-256 over `serde_json::to_vec(value)`. Re-serializing the @@ -238,8 +345,9 @@ impl std::fmt::Debug for DriftState { /// - First time a session is seen → `tracing::info!(event = /// "cache_drift_first_request", …)` with a 16-char prefix of the /// SHA-256 hex of `session_key`. -/// - Subsequent requests with all three hashes equal → no event. -/// - Subsequent requests with any dimension differing → +/// - Subsequent requests with no drifted dimension (append-only +/// growth into the early window included) → no event. +/// - Subsequent requests with any dimension drifting → /// `tracing::warn!(event = "cache_drift_observed", drift_dims = /// "", previous_hash_prefix, current_hash_prefix, …)`. pub fn observe_drift(state: &DriftState, session_key: &str, current: StructuralHash) { @@ -268,21 +376,23 @@ pub fn observe_drift(state: &DriftState, session_key: &str, current: StructuralH ); cache.put(session_key.to_string(), current); } - Some(previous) if previous == current => { - // Stable. No event. Update LRU recency by reinserting. - cache.put(session_key.to_string(), current); - } Some(previous) => { let dims = drift_dims(&previous, ¤t); - tracing::warn!( - event = "cache_drift_observed", - session_key_hash = %session_prefix, - drift_dims = %dims, - previous_hash_prefix = %structural_hash_log_prefix(&previous), - current_hash_prefix = %structural_hash_log_prefix(¤t), - "cache_drift detector observed structural change between turns of the same session" - ); - cache.put(session_key.to_string(), current); + if dims.is_empty() { + // Stable (append-only growth included). No event. + // Update LRU recency by reinserting. + cache.put(session_key.to_string(), current); + } else { + tracing::warn!( + event = "cache_drift_observed", + session_key_hash = %session_prefix, + drift_dims = %dims, + previous_hash_prefix = %structural_hash_log_prefix(&previous), + current_hash_prefix = %structural_hash_log_prefix(¤t), + "cache_drift detector observed structural change between turns of the same session" + ); + cache.put(session_key.to_string(), current); + } } } } @@ -297,14 +407,24 @@ fn session_key_log_prefix(session_key: &str) -> String { hex_prefix(&digest, 16) } -/// 12-char hex prefix of the concatenated structural hash. Useful as -/// a compact "did the prefix change" indicator in logs without -/// printing the entire 96-char digest tuple. +/// 24-char hex prefix (12 bytes) of a digest over the concatenated +/// axis hashes. Useful as a compact "did the prefix change" indicator +/// in logs without printing every axis digest in full. fn structural_hash_log_prefix(hash: &StructuralHash) -> String { let mut hasher = Sha256::new(); hasher.update(hash.system); hasher.update(hash.tools); - hasher.update(hash.early_messages); + for slot in &hash.early_messages { + // Length-prefix the slots so `[Some(h), None]` and `[None, + // Some(h)]` cannot collide. + match slot { + Some(h) => { + hasher.update([1u8]); + hasher.update(h); + } + None => hasher.update([0u8]), + } + } let digest = hasher.finalize(); hex_prefix(&digest, 12) } @@ -334,35 +454,75 @@ fn drift_dims(prev: &StructuralHash, curr: &StructuralHash) -> String { if prev.tools != curr.tools { dims.push("tools"); } - if prev.early_messages != curr.early_messages { + if early_window_drifted(&prev.early_messages, &curr.early_messages) { dims.push("early_messages"); } dims.join(",") } -/// Derive a stable per-session key from the request headers and -/// client address. Priority order: +/// Prefix-aware early-window comparison. A settled slot changing or +/// disappearing is drift (the previously sent prefix was rewritten — +/// the provider's cache is busted); the conversation growing into a +/// previously empty slot is append-only and benign. +fn early_window_drifted( + prev: &[Option<[u8; 32]>; EARLY_MESSAGES_WINDOW], + curr: &[Option<[u8; 32]>; EARLY_MESSAGES_WINDOW], +) -> bool { + prev.iter().zip(curr.iter()).any(|slots| match slots { + (Some(p), Some(c)) => p != c, + (Some(_), None) => true, + (None, _) => false, + }) +} + +/// Derive a stable per-session key from the request headers, client +/// address, and body. Priority order: /// -/// 1. `Authorization` header (hashed; never logged raw). -/// 2. `x-api-key` header (hashed; never logged raw). -/// 3. Client IP address. -/// 4. `(client_ip, user_agent)` synthetic tuple — the user-agent +/// 1. `x-headroom-session-id` header (hashed) — the explicit opt-in +/// the Python proxy already honors for every session-sticky +/// subsystem (prefix tracker, beta-header tracker). When the +/// client declares its session, believe it. +/// 2. `Authorization` header (hashed; never logged raw). +/// 3. `x-api-key` header (hashed; never logged raw). +/// 4. Client IP address. +/// 5. `(client_ip, user_agent)` synthetic tuple — the user-agent /// bucketization gives us *some* discrimination when many /// anonymous clients sit behind the same NAT. /// +/// Arms 2–5 identify a *tenant*, not a conversation: interactive +/// clients (Claude Code, Codex CLI) send the same bearer for every +/// concurrent conversation. Each of those arms therefore also folds in +/// [`conversation_discriminator`] — a fingerprint of the +/// conversation's first message — so parallel conversations do not +/// alternate over one LRU slot and log false `cache_drift_observed` +/// events on every switch. +/// /// The returned string is opaque; never log it directly. Callers /// should pass it straight to [`observe_drift`], which logs only a /// hashed prefix. -pub fn derive_session_key(headers: &HeaderMap, client_addr: &SocketAddr) -> String { +pub fn derive_session_key( + headers: &HeaderMap, + client_addr: &SocketAddr, + body: &serde_json::Value, + kind: ApiKind, +) -> String { + if let Some(sid) = headers + .get("x-headroom-session-id") + .and_then(|v| v.to_str().ok()) + .filter(|s| !s.is_empty()) + { + return format!("session:{}", hash_secret(sid)); + } + let conv = conversation_discriminator(body, kind); if let Some(token) = headers .get(axum::http::header::AUTHORIZATION) .and_then(|v| v.to_str().ok()) { - return format!("auth:{}", hash_secret(token)); + return format!("auth:{}:{conv}", hash_secret(token)); } // `x-api-key` is the Anthropic/OpenAI-Responses convention. if let Some(key) = headers.get("x-api-key").and_then(|v| v.to_str().ok()) { - return format!("apikey:{}", hash_secret(key)); + return format!("apikey:{}:{conv}", hash_secret(key)); } let ip = client_addr.ip().to_string(); if let Some(ua) = headers @@ -375,9 +535,55 @@ pub fn derive_session_key(headers: &HeaderMap, client_addr: &SocketAddr) -> Stri let mut h = DefaultHasher::new(); ip.hash(&mut h); ua.hash(&mut h); - return format!("ipua:{:016x}", h.finish()); + return format!("ipua:{:016x}:{conv}", h.finish()); + } + format!("ip:{ip}:{conv}") +} + +/// 16-hex-char fingerprint of `(model, first conversation message)`, +/// the message canonicalized via [`canonicalize_for_hash`] so a +/// relocated `cache_control` marker does not rotate the conversation's +/// identity between turns (interactive clients resend the history each +/// turn with the opener otherwise byte-stable). `-` when the body +/// carries no conversation messages. +/// +/// The model is folded in because provider prompt caches are +/// per-model: an auxiliary small-model call that reuses a +/// conversation's opener (title generation, summarization sidecars) +/// must not share the conversation's drift baseline, and a +/// mid-conversation model switch genuinely starts a new provider +/// cache lineage. +/// +/// Deliberately excludes the system prompt and tools: those are the +/// *measured* axes, and agentic clients legitimately mutate them +/// mid-conversation. An identity built from mutating content would +/// rotate exactly when the detector should be reporting drift instead. +/// +/// Known trade-offs: a client that *rewrites* its first message +/// (history compaction, rolling-window truncation, Responses chained +/// mode sending delta-only `input`) re-keys to a fresh session — the +/// rewrite surfaces as `cache_drift_first_request` on the new key +/// rather than `cache_drift_observed` against the old baseline. An +/// explicit `x-headroom-session-id` pins the identity and reports +/// those rewrites as drift. Conversations sharing one credential AND +/// a byte-identical opener on the same model still conflate. +fn conversation_discriminator(body: &serde_json::Value, kind: ApiKind) -> String { + let model = body.get("model").and_then(|m| m.as_str()).unwrap_or(""); + match conversation_messages(body, kind).first() { + Some(first) => { + let canonical = canonicalize_for_hash(first, false); + let mut hasher = Sha256::new(); + hasher.update(model.as_bytes()); + // NUL separator: domain-separate the model from the + // message bytes so no (model, message) pair can alias + // another by shifting bytes across the boundary. + hasher.update([0u8]); + hasher.update(serde_json::to_vec(&canonical).unwrap_or_default()); + let digest = hasher.finalize(); + hex_prefix(&digest, 8) + } + None => "-".to_string(), } - format!("ip:{ip}") } /// SHA-256 of `secret`, truncated to 16 hex characters. Sufficient @@ -556,7 +762,7 @@ mod tests { .unwrap(), ); let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 1234); - let key = derive_session_key(&headers, &addr); + let key = derive_session_key(&headers, &addr, &json!({}), ApiKind::Anthropic); // The key MUST NOT contain the raw bearer string anywhere — // not the secret token, not the literal "Bearer", not even // any 8+ char substring of the secret. @@ -591,7 +797,7 @@ mod tests { "sk-very-private-api-key-12345".parse().unwrap(), ); let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), 1234); - let key = derive_session_key(&headers, &addr); + let key = derive_session_key(&headers, &addr, &json!({}), ApiKind::Anthropic); assert!(!key.contains("sk-very-private")); assert!(key.starts_with("apikey:")); } @@ -600,12 +806,12 @@ mod tests { fn session_key_falls_back_to_ip_then_ip_ua() { let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 3)), 5555); // No headers → ip-only. - let bare = derive_session_key(&HeaderMap::new(), &addr); + let bare = derive_session_key(&HeaderMap::new(), &addr, &json!({}), ApiKind::Anthropic); assert!(bare.starts_with("ip:")); // With UA → ipua-tuple. let mut headers = HeaderMap::new(); headers.insert(axum::http::header::USER_AGENT, "ua-test".parse().unwrap()); - let with_ua = derive_session_key(&headers, &addr); + let with_ua = derive_session_key(&headers, &addr, &json!({}), ApiKind::Anthropic); assert!(with_ua.starts_with("ipua:")); assert_ne!(bare, with_ua); } @@ -660,6 +866,328 @@ mod tests { assert_eq!(h1.early_messages, h2.early_messages); } + /// Shape observed from live Claude Code traffic: turn 1 sends a single + /// user message whose *last* block carries the `cache_control` marker; + /// on turn 2 the same message returns byte-identical except the marker + /// moved to the newest message's first block. + fn cc_turn1_body() -> serde_json::Value { + json!({ + "model": "claude-haiku-4-5", + "system": [{"type": "text", "text": "agent preamble"}], + "tools": [{"name": "bash"}], + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "opening prompt"}, + {"type": "text", "text": "project context", "cache_control": {"type": "ephemeral"}}, + ]}, + ], + }) + } + + fn cc_turn2_body() -> serde_json::Value { + json!({ + "model": "claude-haiku-4-5", + "system": [{"type": "text", "text": "agent preamble"}], + "tools": [{"name": "bash"}], + "messages": [ + {"role": "user", "content": [ + {"type": "text", "text": "opening prompt"}, + {"type": "text", "text": "project context"}, + ]}, + {"role": "assistant", "content": [{"type": "text", "text": "reply"}]}, + {"role": "user", "content": [ + {"type": "text", "text": "second prompt", "cache_control": {"type": "ephemeral"}}, + ]}, + ], + }) + } + + #[test] + fn different_conversations_on_one_credential_get_distinct_session_keys() { + // Interactive clients (Claude Code, Codex CLI) send the same + // `Authorization` bearer for every concurrent conversation. If the + // session key stops at the credential, their alternating requests + // ping-pong one LRU slot and every switch logs a false + // `cache_drift_observed`. + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let conv_a = anthropic_body("sys", json!([]), vec!["conversation A opener"]); + let conv_b = anthropic_body("sys", json!([]), vec!["conversation B opener"]); + let key_a = derive_session_key(&headers, &addr, &conv_a, ApiKind::Anthropic); + let key_b = derive_session_key(&headers, &addr, &conv_b, ApiKind::Anthropic); + assert_ne!( + key_a, key_b, + "two conversations sharing one credential must not share a drift session" + ); + } + + #[test] + fn same_conversation_next_turn_keeps_its_session_key() { + // The discriminator must survive normal turn-to-turn growth: the + // opener is byte-identical on turn 2 except its relocated + // `cache_control` marker. + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let key_t1 = derive_session_key(&headers, &addr, &cc_turn1_body(), ApiKind::Anthropic); + let key_t2 = derive_session_key(&headers, &addr, &cc_turn2_body(), ApiKind::Anthropic); + assert_eq!( + key_t1, key_t2, + "turn growth and cache_control relocation must not rotate the session key" + ); + } + + #[test] + fn explicit_headroom_session_header_wins_over_credentials() { + // The Python proxy honors `x-headroom-session-id` as the highest- + // priority session identity (prefix tracker, beta-header tracker); + // the drift detector must respect the same explicit opt-in. + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + headers.insert("x-headroom-session-id", "conv-42".parse().unwrap()); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let body_a = anthropic_body("sys", json!([]), vec!["conversation A opener"]); + let body_b = anthropic_body("sys", json!([]), vec!["conversation B opener"]); + let key_a = derive_session_key(&headers, &addr, &body_a, ApiKind::Anthropic); + assert!( + key_a.starts_with("session:"), + "explicit session header must define the key, got: {key_a}" + ); + assert!( + !key_a.contains("conv-42"), + "session keys stay opaque even for non-secret ids: {key_a}" + ); + // The explicit id pins the session across body variance… + let key_b = derive_session_key(&headers, &addr, &body_b, ApiKind::Anthropic); + assert_eq!(key_a, key_b); + // …and different ids mean different sessions. + let mut headers2 = headers.clone(); + headers2.insert("x-headroom-session-id", "conv-43".parse().unwrap()); + let key_c = derive_session_key(&headers2, &addr, &body_a, ApiKind::Anthropic); + assert_ne!(key_a, key_c); + } + + #[test] + fn cache_control_relocation_and_growth_are_not_drift() { + // Turn 1 → turn 2 of a single conversation: the early window gains + // messages and the client relocates its `cache_control` marker to + // the newest block. Neither busts the provider's prefix cache, so + // neither is drift. + let h1 = compute_structural_hash(&cc_turn1_body(), ApiKind::Anthropic); + let h2 = compute_structural_hash(&cc_turn2_body(), ApiKind::Anthropic); + assert_eq!( + drift_dims(&h1, &h2), + "", + "append-only growth with marker relocation must not be drift" + ); + } + + #[test] + fn append_only_growth_without_markers_is_not_drift() { + let h1 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["m1"]), + ApiKind::Anthropic, + ); + let h2 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["m1", "m2", "m3"]), + ApiKind::Anthropic, + ); + assert_eq!( + drift_dims(&h1, &h2), + "", + "a conversation growing into the early window must not be drift" + ); + } + + #[test] + fn rewritten_early_message_is_still_drift() { + let h1 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["m1", "m2", "m3"]), + ApiKind::Anthropic, + ); + let h2 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["REWRITTEN", "m2", "m3"]), + ApiKind::Anthropic, + ); + assert_eq!(drift_dims(&h1, &h2), "early_messages"); + } + + #[test] + fn shrunk_history_is_still_drift() { + // Fewer messages than previously observed *under the same + // session key* means the settled prefix was rewritten in place + // — that IS a cache bust. (Reachable when the identity is + // pinned, e.g. an explicit x-headroom-session-id; on the + // credential arms a first-message rewrite re-keys instead — + // see conversation_discriminator's trade-off note.) + let h1 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["m1", "m2", "m3"]), + ApiKind::Anthropic, + ); + let h2 = compute_structural_hash( + &anthropic_body("s", json!([]), vec!["m1"]), + ApiKind::Anthropic, + ); + assert_eq!(drift_dims(&h1, &h2), "early_messages"); + } + + #[test] + fn same_opener_on_different_model_gets_distinct_session_keys() { + // Auxiliary small-model calls (title generation, summaries) + // reuse a conversation's opener under the same credential. + // Provider prompt caches are per-model, so these are separate + // cache lineages and must not share a drift baseline — the + // sidecar's different system prompt would otherwise false-warn. + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let mut main_conv = anthropic_body("sys", json!([]), vec!["shared opener"]); + let mut sidecar = anthropic_body("sys", json!([]), vec!["shared opener"]); + main_conv["model"] = json!("opus-large"); + sidecar["model"] = json!("haiku-small"); + let key_main = derive_session_key(&headers, &addr, &main_conv, ApiKind::Anthropic); + let key_side = derive_session_key(&headers, &addr, &sidecar, ApiKind::Anthropic); + assert_ne!(key_main, key_side); + } + + #[test] + fn key_order_variation_does_not_perturb_hashes_or_identity() { + // The workspace's serde_json enables `preserve_order`, so two + // serializations of the same message with different key order + // stay distinct through Value round-trips. Canonicalization + // must neutralize that for both the axes and the session key. + let body_a: serde_json::Value = serde_json::from_str( + r#"{"model":"m","system":"s","tools":[], + "messages":[{"role":"user","content":"hello"}]}"#, + ) + .unwrap(); + let body_b: serde_json::Value = serde_json::from_str( + r#"{"model":"m","system":"s","tools":[], + "messages":[{"content":"hello","role":"user"}]}"#, + ) + .unwrap(); + assert_eq!( + compute_structural_hash(&body_a, ApiKind::Anthropic), + compute_structural_hash(&body_b, ApiKind::Anthropic), + ); + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + assert_eq!( + derive_session_key(&headers, &addr, &body_a, ApiKind::Anthropic), + derive_session_key(&headers, &addr, &body_b, ApiKind::Anthropic), + ); + } + + #[test] + fn cache_control_named_fields_inside_opaque_payloads_still_count() { + // A tool schema property (or tool-call argument) that happens + // to be NAMED cache_control is user data, not a cache marker. + // Changing it must still read as drift on the affected axis. + let with_schema = |ty: &str| { + json!({ + "model": "m", "system": "s", + "tools": [{"name": "t", "input_schema": + {"properties": {"cache_control": {"type": ty}}}}], + "messages": [{"role": "user", "content": "hi"}], + }) + }; + let h1 = compute_structural_hash(&with_schema("string"), ApiKind::Anthropic); + let h2 = compute_structural_hash(&with_schema("integer"), ApiKind::Anthropic); + assert_eq!(drift_dims(&h1, &h2), "tools"); + + let with_tool_input = |v: &str| { + json!({ + "model": "m", "system": "s", "tools": [], + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "x", "name": "t", + "input": {"cache_control": v}}, + ]}, + ], + }) + }; + let h3 = compute_structural_hash(&with_tool_input("a"), ApiKind::Anthropic); + let h4 = compute_structural_hash(&with_tool_input("b"), ApiKind::Anthropic); + assert_eq!(drift_dims(&h3, &h4), "early_messages"); + } + + #[test] + fn bare_string_messages_only_count_for_responses_input() { + // `input: "text"` is valid Responses sugar; a bare-string + // `messages` on the other shapes is malformed and contributes + // no conversation identity. + let responses = json!({"model": "m", "instructions": "i", "input": "hello"}); + assert_eq!( + conversation_messages(&responses, ApiKind::OpenAiResponses).len(), + 1 + ); + let malformed = json!({"model": "m", "system": "s", "messages": "oops"}); + assert!(conversation_messages(&malformed, ApiKind::Anthropic).is_empty()); + assert!(conversation_discriminator(&malformed, ApiKind::Anthropic) == "-"); + } + + #[test] + fn openai_chat_discriminator_uses_first_non_system_message() { + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let conv = |first_user: &str| { + json!({ + "model": "gpt-4", + "messages": [ + {"role": "system", "content": "shared assistant config"}, + {"role": "user", "content": first_user}, + ], + }) + }; + let key_a = derive_session_key(&headers, &addr, &conv("opener A"), ApiKind::OpenAiChat); + let key_b = derive_session_key(&headers, &addr, &conv("opener B"), ApiKind::OpenAiChat); + assert_ne!(key_a, key_b); + } + + #[test] + fn openai_responses_discriminator_uses_first_input_item() { + let mut headers = HeaderMap::new(); + headers.insert( + axum::http::header::AUTHORIZATION, + "Bearer shared-workspace-token".parse().unwrap(), + ); + let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 9)), 4242); + let conv = |first_input: &str| { + json!({ + "model": "gpt-4", + "instructions": "shared instructions", + "input": [{"type": "message", "role": "user", "content": first_input}], + }) + }; + let key_a = + derive_session_key(&headers, &addr, &conv("opener A"), ApiKind::OpenAiResponses); + let key_b = + derive_session_key(&headers, &addr, &conv("opener B"), ApiKind::OpenAiResponses); + assert_ne!(key_a, key_b); + } + #[test] fn early_messages_window_caps_at_three() { // 5 messages: hash should depend only on the first 3. diff --git a/crates/headroom-proxy/src/cache_stabilization/mod.rs b/crates/headroom-proxy/src/cache_stabilization/mod.rs index 7ae753d4c..815c3af21 100644 --- a/crates/headroom-proxy/src/cache_stabilization/mod.rs +++ b/crates/headroom-proxy/src/cache_stabilization/mod.rs @@ -26,7 +26,8 @@ //! the cache hot zone (system / tools / early messages). Emits //! `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. +//! session drift on any of the three dimensions (append-only +//! conversation growth and `cache_control` relocation are benign). //! - [`tool_def_normalize`] — PR-E1 / PR-E2: sorts `tools[]` //! alphabetically by name (PR-E1) and recursively sorts JSON //! Schema object keys inside each tool's `input_schema` / diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index cb9f1ffb1..7f3d81cd3 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -75,11 +75,16 @@ pub struct AppState { } /// PR-E6: maximum number of sessions tracked by the drift detector -/// LRU. Picked so that a noisy test fleet of 1000 distinct API keys -/// stays in cache for at least one full turn before the oldest -/// evicts. Operators with larger fleets can bump this; the memory -/// cost per entry is ~150 bytes (key string + 96-byte StructuralHash -/// + LRU overhead). +/// LRU. Sessions are keyed per conversation (credential + first- +/// message fingerprint), not per credential, so the working set is +/// the number of *concurrently active conversations* — 1000 keeps a +/// noisy fleet in cache for at least one full turn before the oldest +/// evicts. A burst of short one-shot conversations can cycle the LRU +/// and evict a live session between its turns; the cost is telemetry- +/// only (one repeated `cache_drift_first_request`, no lost requests). +/// Operators with larger fleets can bump this; the memory cost per +/// entry is ~250 bytes (key string + 163-byte StructuralHash + LRU +/// overhead). const DRIFT_DETECTOR_CAPACITY: usize = 1000; impl AppState { @@ -699,7 +704,7 @@ pub(crate) async fn forward_http( } }; if let (Some(kind), Some(headers)) = (drift_kind, headers_snapshot.as_ref()) { - let session_key = derive_session_key(headers, &client_addr); + 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); }