fix: PR-E6 cache-bust drift detector telemetry (Phase E)

Per-session SHA-256 fingerprint of the cache hot zone (system / tools /
first 3 messages) with structured-log emission on drift. Detector is
read-only: never mutates request bytes, preserves the Phase A
passthrough invariant. Surfaces invisible cache busts (system prompt
edited mid-session, tools reshuffled, early message changed) without
rewriting them.

* crates/headroom-proxy/src/cache_stabilization/drift_detector.rs:
  StructuralHash (system, tools, early_messages digests),
  compute_structural_hash, observe_drift, derive_session_key,
  DriftState (LRU bounded to 1000 sessions in production).
* Session keys derive from Authorization / x-api-key / client IP /
  (IP, user-agent). Bearer tokens and API keys are SHA-256 hashed
  before they ever reach the log line; the raw secret is never logged.
* Wired into forward_http after the body is buffered, before the
  compression dispatcher runs. Skips paths whose wire shape is not
  Anthropic / OpenAI Chat / OpenAI Responses.
* AppState gains drift_state: DriftState. Bedrock unit-test
  literal-construction sites updated.
* 14 unit tests + 1 integration test covering first-request,
  no-drift, per-dimension drift, multi-dim drift, LRU eviction,
  non-mutation invariant, and bearer-token-never-logged.

Adds lru = "0.12" and promotes sha2 = "0.10" to a normal dependency
on headroom-proxy.
This commit is contained in:
chopratejas 2026-05-04 12:42:32 -07:00
parent c81755c965
commit ce37940d17
8 changed files with 1027 additions and 29 deletions

12
Cargo.lock generated
View file

@ -1729,6 +1729,8 @@ version = "0.15.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1"
dependencies = [
"allocator-api2",
"equivalent",
"foldhash 0.1.5",
]
@ -1825,6 +1827,7 @@ dependencies = [
"humantime",
"hyper",
"hyper-util",
"lru",
"pin-project-lite",
"prometheus",
"proptest",
@ -2469,6 +2472,15 @@ dependencies = [
"imgref",
]
[[package]]
name = "lru"
version = "0.12.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38"
dependencies = [
"hashbrown 0.15.5",
]
[[package]]
name = "lru-slab"
version = "0.1.2"

View file

@ -59,6 +59,18 @@ crc32fast = "1"
# defaults and re-enable nothing — pure registry + counter +
# histogram + text encoder is sufficient.
prometheus = { version = "0.13", default-features = false }
# PR-E6: SHA-256 over canonical bytes of the cache hot zone (system,
# tools, early messages) for cache-bust drift detection. Already in
# the dev-dependencies (and pulled transitively by `aws-sigv4` via
# `aws-smithy-runtime-api`); promoted here to a direct, normal-build
# dependency so the drift detector compiles outside `cfg(test)`.
sha2 = "0.10"
# PR-E6: bounded session-scoped cache of structural hashes. The
# detector evicts the oldest session at 1000 entries — we never want
# unbounded memory growth from a flood of unique session keys. `lru`
# is the de-facto Rust LRU crate; minimal surface, no dependencies of
# our own beyond `hashbrown` (which we already pull transitively).
lru = "0.12"
[dev-dependencies]
tower = { workspace = true, features = ["util"] }

View file

@ -543,6 +543,10 @@ mod tests {
config: std::sync::Arc::new(config),
client: reqwest::Client::new(),
bedrock_credentials: None,
// PR-E6: small capacity is fine — the Bedrock URL builder
// unit test never observes drift, but `AppState` requires
// the field to be populated.
drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8),
};
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
.parse()
@ -569,6 +573,9 @@ mod tests {
config: std::sync::Arc::new(config),
client: reqwest::Client::new(),
bedrock_credentials: None,
// 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),
};
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke"
.parse()

View file

@ -971,6 +971,9 @@ mod tests {
config: std::sync::Arc::new(config),
client: reqwest::Client::new(),
bedrock_credentials: None,
// 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),
};
let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke-with-response-stream"
.parse()

View file

@ -0,0 +1,683 @@
//! PR-E6: cache-bust drift detector.
//!
//! # What it does
//!
//! For every inbound request on a known LLM endpoint, compute a
//! [`StructuralHash`] over the **cache hot zone**:
//!
//! - `system` — SHA-256 of the canonical system-prompt bytes (Anthropic
//! `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
//! live-zone tail where mutation is expected and benign.
//!
//! Track the previous hash 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.
//!
//! # 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.
//!
//! # Cost
//!
//! - One SHA-256 update over each of (system, tools, early messages).
//! Total ~200us on a 8 KB system prompt.
//! - 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
//! accept session-scoped counters without a cardinality explosion).
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use std::net::SocketAddr;
use std::num::NonZeroUsize;
use std::sync::{Arc, Mutex};
use axum::http::HeaderMap;
use lru::LruCache;
use sha2::{Digest, Sha256};
/// Which provider's body shape we're hashing. The walker is shaped
/// per provider because the cache hot zone lives in different fields:
/// Anthropic uses `body.system`/`body.tools`/`body.messages`, OpenAI
/// Chat threads `system` into the first message, and OpenAI Responses
/// uses `body.instructions`/`body.tools`/`body.input`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ApiKind {
/// `POST /v1/messages` (Anthropic).
Anthropic,
/// `POST /v1/chat/completions` (OpenAI).
OpenAiChat,
/// `POST /v1/responses` (OpenAI Responses API).
OpenAiResponses,
}
/// 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.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StructuralHash {
pub system: [u8; 32],
pub tools: [u8; 32],
pub early_messages: [u8; 32],
}
/// How many message-shaped items count as the "early" prefix that
/// feeds `early_messages_hash`. Anything past this is the live zone
/// (where mutation is expected; we deliberately ignore it).
const EARLY_MESSAGES_WINDOW: usize = 3;
/// Compute a [`StructuralHash`] for the body shape implied by `kind`.
///
/// `body` is borrowed; **this function never mutates it**. The
/// `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));
StructuralHash {
system,
tools,
early_messages,
}
}
/// 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
/// well-defined.
fn extract_system(body: &serde_json::Value, kind: ApiKind) -> serde_json::Value {
match kind {
ApiKind::Anthropic => body
.get("system")
.cloned()
.unwrap_or(serde_json::Value::Null),
ApiKind::OpenAiChat => {
// First message with `role == "system"` is the OpenAI
// Chat hot-zone equivalent. There can be at most one in
// practice (newer requests use a `developer` role; that's
// not the system axis and we deliberately don't conflate).
body.get("messages")
.and_then(|v| v.as_array())
.and_then(|arr| {
arr.iter().find(|m| {
m.get("role")
.and_then(|r| r.as_str())
.map(|s| s == "system")
.unwrap_or(false)
})
})
.cloned()
.unwrap_or(serde_json::Value::Null)
}
ApiKind::OpenAiResponses => body
.get("instructions")
.cloned()
.unwrap_or(serde_json::Value::Null),
}
}
/// Extract the "tools" axis as a `serde_json::Value`. The same
/// `tools` array key is used by all three providers in practice.
fn extract_tools(body: &serde_json::Value) -> serde_json::Value {
body.get("tools")
.cloned()
.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 {
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 early: Vec<serde_json::Value> = match kind {
ApiKind::OpenAiChat => messages
.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)
}
/// SHA-256 over `serde_json::to_vec(value)`. Re-serializing the
/// borrowed `Value` defends against trivial whitespace differences
/// from the wire — operators care about *semantic* drift, not
/// formatter drift.
fn hash_value(value: &serde_json::Value) -> [u8; 32] {
// `serde_json::to_vec` on a `Value` cannot fail except on a
// pathological recursion, which the upstream API would itself
// reject; on the impossible failure path we hash the empty byte
// string so the digest is still stable rather than panicking and
// taking the request down.
let bytes = serde_json::to_vec(value).unwrap_or_default();
let mut hasher = Sha256::new();
hasher.update(&bytes);
let digest = hasher.finalize();
let mut out = [0u8; 32];
out.copy_from_slice(&digest);
out
}
/// Bounded session → last-seen `StructuralHash` map. Wrapped in
/// `Arc<Mutex<…>>` so it can be cloned freely into `AppState` without
/// duplicating the underlying LRU.
#[derive(Clone)]
pub struct DriftState {
cache: Arc<Mutex<LruCache<String, StructuralHash>>>,
}
impl DriftState {
/// Build a new `DriftState` bounded to `capacity` sessions. The
/// production capacity is 1000; tests pass small values so the
/// LRU eviction path is exercised cheaply.
///
/// # Panics
///
/// Panics if `capacity == 0`. The detector is meaningless without
/// at least one slot — use `LruCache::new(NonZeroUsize::MIN)` if
/// you need a "remember nothing" mode.
pub fn new(capacity: usize) -> Self {
let cap = NonZeroUsize::new(capacity).expect("DriftState capacity must be > 0");
Self {
cache: Arc::new(Mutex::new(LruCache::new(cap))),
}
}
}
impl std::fmt::Debug for DriftState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let len = self.cache.lock().map(|c| c.len()).unwrap_or(0);
f.debug_struct("DriftState").field("len", &len).finish()
}
}
/// Compare `current` against the last-seen hash for `session_key` and
/// emit a structured `tracing` event accordingly. Always updates the
/// LRU to `current` before returning so the next call sees the most
/// recent fingerprint.
///
/// Logging contract:
///
/// - 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 →
/// `tracing::warn!(event = "cache_drift_observed", drift_dims =
/// "<comma-joined>", previous_hash_prefix, current_hash_prefix, …)`.
pub fn observe_drift(state: &DriftState, session_key: &str, current: StructuralHash) {
let session_prefix = session_key_log_prefix(session_key);
let mut cache = match state.cache.lock() {
Ok(c) => c,
Err(poisoned) => {
// Mutex was poisoned by a panicking writer in another
// task. Recover the inner data — the only thing we lose
// is one stale entry, and continuing the request is
// strictly preferable to failing closed.
tracing::warn!(
event = "cache_drift_state_mutex_poisoned",
"drift detector mutex was poisoned by a panicking task; recovering"
);
poisoned.into_inner()
}
};
match cache.get(session_key).copied() {
None => {
tracing::info!(
event = "cache_drift_first_request",
session_key_hash = %session_prefix,
current_hash_prefix = %structural_hash_log_prefix(&current),
"cache_drift detector observed a new session"
);
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, &current);
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(&current),
"cache_drift detector observed structural change between turns of the same session"
);
cache.put(session_key.to_string(), current);
}
}
}
/// 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 {
let mut hasher = Sha256::new();
hasher.update(session_key.as_bytes());
let digest = hasher.finalize();
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.
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);
let digest = hasher.finalize();
hex_prefix(&digest, 12)
}
/// Lowercase hex of the first `take` bytes of `bytes`. Allocates a
/// `String` once per call.
fn hex_prefix(bytes: &[u8], take: usize) -> String {
let take = take.min(bytes.len());
let mut out = String::with_capacity(take * 2);
for b in &bytes[..take] {
// Manual hex; avoids pulling `hex` for one call site.
const HEX: &[u8; 16] = b"0123456789abcdef";
out.push(HEX[(b >> 4) as usize] as char);
out.push(HEX[(b & 0xf) as usize] as char);
}
out
}
/// Comma-joined list of which dimensions drifted between `prev` and
/// `curr`. The order is fixed (`system`, `tools`, `early_messages`)
/// so log queries can match deterministically.
fn drift_dims(prev: &StructuralHash, curr: &StructuralHash) -> String {
let mut dims: Vec<&'static str> = Vec::with_capacity(3);
if prev.system != curr.system {
dims.push("system");
}
if prev.tools != curr.tools {
dims.push("tools");
}
if 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:
///
/// 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
/// bucketization gives us *some* discrimination when many
/// anonymous clients sit behind the same NAT.
///
/// 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 {
if let Some(token) = headers
.get(axum::http::header::AUTHORIZATION)
.and_then(|v| v.to_str().ok())
{
return format!("auth:{}", 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));
}
let ip = client_addr.ip().to_string();
if let Some(ua) = headers
.get(axum::http::header::USER_AGENT)
.and_then(|v| v.to_str().ok())
{
// Hash the (ip, ua) tuple so the resulting key remains opaque
// and does not leak full UA strings into downstream logs that
// forget our "log only the prefix" contract.
let mut h = DefaultHasher::new();
ip.hash(&mut h);
ua.hash(&mut h);
return format!("ipua:{:016x}", h.finish());
}
format!("ip:{ip}")
}
/// SHA-256 of `secret`, truncated to 16 hex characters. Sufficient
/// to discriminate sessions while pinning that the raw secret never
/// reaches the log line. We do **not** use the full digest because
/// even a hashed bearer that ends up in many log entries leaks
/// fingerprintable information; the 16-char prefix bounds that.
fn hash_secret(secret: &str) -> String {
let mut hasher = Sha256::new();
hasher.update(secret.as_bytes());
let digest = hasher.finalize();
hex_prefix(&digest, 16)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
use std::net::{IpAddr, Ipv4Addr};
fn anthropic_body(
system: &str,
tools: serde_json::Value,
msgs: Vec<&str>,
) -> serde_json::Value {
let messages: Vec<serde_json::Value> = msgs
.into_iter()
.map(|t| json!({"role": "user", "content": t}))
.collect();
json!({
"model": "claude-3-5-sonnet-20241022",
"system": system,
"tools": tools,
"messages": messages,
})
}
fn make_state() -> DriftState {
DriftState::new(8)
}
#[test]
fn first_request_emits_first_request_event() {
let state = make_state();
let body = anthropic_body("you are an assistant", json!([]), vec!["hi"]);
let h = compute_structural_hash(&body, ApiKind::Anthropic);
// Before observation: empty cache.
assert_eq!(state.cache.lock().unwrap().len(), 0);
observe_drift(&state, "session-A", h);
// After observation: 1 entry, equal to the input hash.
let cache = state.cache.lock().unwrap();
assert_eq!(cache.len(), 1);
assert_eq!(cache.peek("session-A"), Some(&h));
}
#[test]
fn same_hash_emits_no_event() {
let state = make_state();
let body = anthropic_body("sys-A", json!([]), vec!["m1"]);
let h = compute_structural_hash(&body, ApiKind::Anthropic);
observe_drift(&state, "sess", h);
// Second observation with identical hash: still 1 entry, same hash.
observe_drift(&state, "sess", h);
let cache = state.cache.lock().unwrap();
assert_eq!(cache.len(), 1);
assert_eq!(cache.peek("sess"), Some(&h));
}
#[test]
fn system_drift_detected_with_correct_dim() {
let state = make_state();
let h1 = compute_structural_hash(
&anthropic_body("sys-A", json!([]), vec!["m1"]),
ApiKind::Anthropic,
);
let h2 = compute_structural_hash(
&anthropic_body("sys-B", json!([]), vec!["m1"]),
ApiKind::Anthropic,
);
assert_ne!(h1.system, h2.system);
assert_eq!(h1.tools, h2.tools);
assert_eq!(h1.early_messages, h2.early_messages);
assert_eq!(drift_dims(&h1, &h2), "system");
observe_drift(&state, "sess", h1);
observe_drift(&state, "sess", h2);
}
#[test]
fn tools_drift_detected_with_correct_dim() {
let h1 = compute_structural_hash(
&anthropic_body("sys", json!([{"name": "a"}]), vec!["m1"]),
ApiKind::Anthropic,
);
let h2 = compute_structural_hash(
&anthropic_body("sys", json!([{"name": "b"}]), vec!["m1"]),
ApiKind::Anthropic,
);
assert_eq!(h1.system, h2.system);
assert_ne!(h1.tools, h2.tools);
assert_eq!(h1.early_messages, h2.early_messages);
assert_eq!(drift_dims(&h1, &h2), "tools");
}
#[test]
fn early_messages_drift_detected_with_correct_dim() {
let h1 = compute_structural_hash(
&anthropic_body("sys", json!([]), vec!["m1"]),
ApiKind::Anthropic,
);
let h2 = compute_structural_hash(
&anthropic_body("sys", json!([]), vec!["DIFFERENT"]),
ApiKind::Anthropic,
);
assert_eq!(h1.system, h2.system);
assert_eq!(h1.tools, h2.tools);
assert_ne!(h1.early_messages, h2.early_messages);
assert_eq!(drift_dims(&h1, &h2), "early_messages");
}
#[test]
fn multi_dim_drift_lists_all_changed_dims() {
let h1 = compute_structural_hash(
&anthropic_body("sys-A", json!([{"name": "a"}]), vec!["m1"]),
ApiKind::Anthropic,
);
let h2 = compute_structural_hash(
&anthropic_body("sys-B", json!([{"name": "b"}]), vec!["X"]),
ApiKind::Anthropic,
);
assert_eq!(drift_dims(&h1, &h2), "system,tools,early_messages");
}
#[test]
fn lru_evicts_at_capacity() {
// Capacity 2: inserting a 3rd session evicts the LRU.
let state = DriftState::new(2);
let h = compute_structural_hash(
&anthropic_body("s", json!([]), vec!["m"]),
ApiKind::Anthropic,
);
observe_drift(&state, "s1", h);
observe_drift(&state, "s2", h);
observe_drift(&state, "s3", h);
let cache = state.cache.lock().unwrap();
assert_eq!(cache.len(), 2);
// s1 was the least-recently-used; should have been evicted.
assert!(!cache.contains("s1"));
assert!(cache.contains("s2"));
assert!(cache.contains("s3"));
}
#[test]
fn does_not_mutate_input() {
let body = anthropic_body(
"sys",
json!([{"name": "t1", "input_schema": {"type": "object"}}]),
vec!["m1", "m2", "m3", "m4"],
);
let original_bytes = serde_json::to_vec(&body).expect("serialize");
// Compute the hash twice — across the three ApiKind shapes —
// to exercise every branch that *could* mutate the input.
let _ = compute_structural_hash(&body, ApiKind::Anthropic);
let _ = compute_structural_hash(&body, ApiKind::OpenAiChat);
let _ = compute_structural_hash(&body, ApiKind::OpenAiResponses);
let after_bytes = serde_json::to_vec(&body).expect("re-serialize");
assert_eq!(original_bytes, after_bytes);
}
#[test]
fn session_key_hashes_authorization_does_not_log_raw() {
let mut headers = HeaderMap::new();
headers.insert(
axum::http::header::AUTHORIZATION,
"Bearer sk-ant-very-secret-token-do-not-log-me"
.parse()
.unwrap(),
);
let addr: SocketAddr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 1234);
let key = derive_session_key(&headers, &addr);
// 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.
assert!(
!key.contains("sk-ant"),
"session key leaked raw secret prefix: {key}"
);
assert!(
!key.contains("very-secret"),
"session key leaked raw secret middle: {key}"
);
assert!(
!key.contains("Bearer"),
"session key leaked the auth scheme: {key}"
);
// The key SHOULD be the auth-scoped envelope, so we know the
// `Authorization` arm was taken (not the IP fallback).
assert!(key.starts_with("auth:"), "expected auth-scoped key: {key}");
// And the log prefix must also not leak the raw secret.
let log_prefix = session_key_log_prefix(&key);
assert!(!log_prefix.contains("sk-ant"));
assert!(!log_prefix.contains("very-secret"));
assert!(!log_prefix.contains("Bearer"));
assert_eq!(log_prefix.len(), 32); // 16 bytes × 2 hex chars
}
#[test]
fn session_key_hashes_x_api_key_does_not_log_raw() {
let mut headers = HeaderMap::new();
headers.insert(
"x-api-key",
"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);
assert!(!key.contains("sk-very-private"));
assert!(key.starts_with("apikey:"));
}
#[test]
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);
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);
assert!(with_ua.starts_with("ipua:"));
assert_ne!(bare, with_ua);
}
#[test]
fn openai_chat_extracts_first_system_message() {
let body = json!({
"model": "gpt-4",
"messages": [
{"role": "system", "content": "you are a helpful assistant"},
{"role": "user", "content": "hi"},
],
"tools": [],
});
let h1 = compute_structural_hash(&body, ApiKind::OpenAiChat);
// Same body but a different system message → system axis drifts.
let body2 = json!({
"model": "gpt-4",
"messages": [
{"role": "system", "content": "you are a different assistant"},
{"role": "user", "content": "hi"},
],
"tools": [],
});
let h2 = compute_structural_hash(&body2, ApiKind::OpenAiChat);
assert_ne!(h1.system, h2.system);
// user message identical → early-messages stays identical.
assert_eq!(h1.early_messages, h2.early_messages);
}
#[test]
fn openai_responses_uses_instructions_and_input() {
let body = json!({
"model": "gpt-4",
"instructions": "be brief",
"tools": [],
"input": [
{"type": "message", "role": "user", "content": "hello"},
],
});
let h1 = compute_structural_hash(&body, ApiKind::OpenAiResponses);
let body2 = json!({
"model": "gpt-4",
"instructions": "be verbose",
"tools": [],
"input": [
{"type": "message", "role": "user", "content": "hello"},
],
});
let h2 = compute_structural_hash(&body2, ApiKind::OpenAiResponses);
assert_ne!(h1.system, h2.system);
assert_eq!(h1.early_messages, h2.early_messages);
}
#[test]
fn early_messages_window_caps_at_three() {
// 5 messages: hash should depend only on the first 3.
let h1 = compute_structural_hash(
&anthropic_body("s", json!([]), vec!["a", "b", "c", "d", "e"]),
ApiKind::Anthropic,
);
// Mutating message 4 only must NOT drift the early_messages hash.
let h2 = compute_structural_hash(
&anthropic_body("s", json!([]), vec!["a", "b", "c", "DIFFERENT", "e"]),
ApiKind::Anthropic,
);
assert_eq!(h1.early_messages, h2.early_messages);
// But mutating message 1 must drift it.
let h3 = compute_structural_hash(
&anthropic_body("s", json!([]), vec!["DIFFERENT", "b", "c", "d", "e"]),
ApiKind::Anthropic,
);
assert_ne!(h1.early_messages, h3.early_messages);
}
}

View file

@ -1,22 +1,32 @@
//! Cache-stabilization observability surface (Phase E).
//! Phase E cache-stabilization surface.
//!
//! Phase E PRs in this module sit *next to* the request path —
//! never on it. They observe inbound bodies and emit structured
//! warnings so customers can see why their prompt-cache hit rate
//! is degrading. Nothing in here mutates request bytes; the
//! cache-safety invariant from Phase A still holds.
//! The realignment plan (`REALIGNMENT/07-phase-E-cache-stabilization.md`)
//! groups every cache-stabilization mechanism behind one module so
//! operators searching for "what does Headroom do to keep prompt
//! caches warm" land in one place. Phase E PRs in this module sit
//! *next to* the request path — never on it. They observe inbound
//! bodies and emit structured logs so customers can see why their
//! prompt-cache hit rate is degrading. Nothing in here mutates
//! request bytes; the cache-safety invariant from Phase A still holds.
//!
//! Currently shipped:
//!
//! - [`volatile_detector`] (PR-E5): scans inbound bodies for
//! patterns that bust prompt-cache hits (timestamps, UUIDs,
//! ID-named fields) and emits one structured WARN log per
//! finding so customers know what to move out of the cached
//! prefix.
//! - [`volatile_detector`] — PR-E5: scans inbound bodies for patterns
//! that bust prompt-cache hits (ISO 8601 timestamps, UUID v4s,
//! ID-named fields) and emits one structured WARN log per finding
//! so customers know what to move out of the cached prefix.
//! - [`drift_detector`] — PR-E6: per-session SHA-256 fingerprint of
//! 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.
//!
//! Sibling PRs (PR-E6, ...) will land additional detectors here.
//! Conflict resolution between parallel PRs is intentionally
//! trivial: each detector lives in its own file, the only shared
//! surface is this `mod.rs`'s `pub mod` list.
//! Future PRs (E1 — tool-array sort, E2 — JSON Schema key sort, E3 —
//! `cache_control` auto-placement, E4 — `prompt_cache_key` injection)
//! hang sibling submodules off this same `mod.rs`. Conflict
//! resolution between parallel Phase E PRs is intentionally trivial:
//! each detector lives in its own file, the only shared surface is
//! this `mod.rs`'s `pub mod` list.
pub mod drift_detector;
pub mod volatile_detector;

View file

@ -17,6 +17,9 @@ use futures_util::{StreamExt as _, TryStreamExt};
use http_body_util::BodyExt;
use crate::cache_stabilization;
use crate::cache_stabilization::drift_detector::{
compute_structural_hash, derive_session_key, observe_drift, ApiKind, DriftState,
};
use crate::compression;
use crate::config::Config;
use crate::error::ProxyError;
@ -48,8 +51,22 @@ pub struct AppState {
/// log so failures are LOUD — no silent fallback to unsigned
/// requests.
pub bedrock_credentials: Option<Arc<aws_credential_types::Credentials>>,
/// PR-E6: per-session structural-hash LRU for the cache-bust
/// drift detector. Bounded to 1000 sessions in production. The
/// detector is read-only — observing it never mutates the
/// request body — so this can be cloned freely into every handler
/// path that buffers the body.
pub drift_state: DriftState,
}
/// 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).
const DRIFT_DETECTOR_CAPACITY: usize = 1000;
impl AppState {
pub fn new(config: Config) -> Result<Self, ProxyError> {
let client = reqwest::Client::builder()
@ -67,6 +84,7 @@ impl AppState {
config: Arc::new(config),
client,
bedrock_credentials: None,
drift_state: DriftState::new(DRIFT_DETECTOR_CAPACITY),
})
}
@ -430,6 +448,19 @@ pub(crate) async fn forward_http(
&& compression::is_compressible_path(uri.path())
&& is_application_json(req.headers());
// PR-E6: capture a header snapshot BEFORE the body is consumed so
// the drift detector can derive a per-session key from
// `Authorization`/`x-api-key`/`User-Agent`. `req` will be moved
// into either `to_bytes(req.into_body())` (buffered branch) or
// `req.into_body().into_data_stream()` (streaming branch); both
// discard the headers along with the body. Snapshot here keeps
// both branches clean.
let headers_snapshot = if should_intercept {
Some(req.headers().clone())
} else {
None
};
let reqwest_method = reqwest::Method::from_bytes(method.as_str().as_bytes())
.map_err(|e| ProxyError::InvalidHeader(e.to_string()))?;
@ -498,27 +529,57 @@ pub(crate) async fn forward_http(
let endpoint = compression::classify_compressible_path(uri.path())
.expect("is_compressible_path guarded above");
// PR-E5: volatile-content detector. Parses the buffered
// body once and walks it read-only, emitting one structured
// WARN log per finding (capped at 10) when the customer's
// cached prefix contains content that busts prompt-cache
// hits (timestamps, UUIDs, ID-named fields). Strictly
// observation-only — never mutates the body. Cheap parse
// failure (malformed JSON) is silently skipped here; the
// dispatcher below logs its own parse-error decision. This
// call is intentionally placed BEFORE dispatch so detection
// runs regardless of whether the dispatcher returns
// `NoCompression`, `Compressed`, or `Passthrough`.
// PR-E5 + PR-E6: cache-stabilization observability hooks.
// Both run READ-ONLY against the buffered body and emit
// structured logs only — passthrough invariant from Phase A
// is preserved. Parsing happens once and is shared. Cheap
// parse failure (malformed JSON) silently skips both
// detectors; the dispatcher below logs its own parse-error
// decision. The hooks run regardless of whether the
// dispatcher returns `NoCompression`, `Compressed`, or
// `Passthrough`.
//
// Bedrock and other shape-mismatched paths skip the drift
// detector specifically; their wire shape is different
// enough that a canonical-bytes hash would compare apples
// to oranges. The volatile detector handles its own
// shape-dispatch via `ApiKind::from_endpoint`.
if let Ok(parsed) = serde_json::from_slice::<serde_json::Value>(&buffered) {
let api_kind = cache_stabilization::volatile_detector::ApiKind::from_endpoint(endpoint);
let findings =
cache_stabilization::volatile_detector::detect_volatile_content(&parsed, api_kind);
// PR-E5: volatile-content detector. Emits one WARN per
// finding (capped at 10) for content that busts cache
// (timestamps, UUIDs, ID-named fields).
let volatile_kind =
cache_stabilization::volatile_detector::ApiKind::from_endpoint(endpoint);
let findings = cache_stabilization::volatile_detector::detect_volatile_content(
&parsed,
volatile_kind,
);
if !findings.is_empty() {
cache_stabilization::volatile_detector::emit_volatile_warnings(
&findings,
&request_id,
);
}
// PR-E6: cache-bust drift detector. SHA-256 fingerprints
// the cache hot zone (system / tools / first 3 messages);
// a mismatch between consecutive turns of the same session
// emits a `cache_drift_observed` event so operators see
// invisible cache busts.
let drift_kind = match endpoint {
compression::CompressibleEndpoint::AnthropicMessages => Some(ApiKind::Anthropic),
compression::CompressibleEndpoint::OpenAiChatCompletions => {
Some(ApiKind::OpenAiChat)
}
compression::CompressibleEndpoint::OpenAiResponses => {
Some(ApiKind::OpenAiResponses)
}
};
if let (Some(kind), Some(headers)) = (drift_kind, headers_snapshot.as_ref()) {
let session_key = derive_session_key(headers, &client_addr);
let hash = compute_structural_hash(&parsed, kind);
observe_drift(&state.drift_state, &session_key, hash);
}
}
let outcome = match endpoint {
compression::CompressibleEndpoint::AnthropicMessages => {

View file

@ -0,0 +1,210 @@
//! Integration tests for the PR-E6 cache-bust drift detector.
//!
//! Boots a real Rust proxy in front of a wiremock upstream, sends two
//! requests on the same `Authorization` (= same session), and asserts
//! that:
//!
//! 1. A second request with a *different* system prompt produces a
//! `cache_drift_observed` warn-level event whose `drift_dims`
//! field includes `system`.
//! 2. The proxy still forwards bytes byte-equal to upstream — the
//! detector is read-only.
//! 3. The session key is hashed in the log line; the raw bearer token
//! (`sk-test-this-is-a-secret`) never appears anywhere in the
//! captured log buffer.
mod common;
use common::start_proxy_with;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::sync::{Arc, Mutex};
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// SHA-256 hex of `bytes`. Used to assert byte-faithful passthrough.
fn sha256_hex(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
hasher
.finalize()
.iter()
.fold(String::with_capacity(64), |mut acc, b| {
use std::fmt::Write as _;
let _ = write!(acc, "{b:02x}");
acc
})
}
/// Mount a /v1/messages handler that captures every body that arrives
/// at upstream into the returned `Vec<Vec<u8>>` for later assertions.
async fn mount_anthropic_capture_all(upstream: &MockServer) -> Arc<Mutex<Vec<Vec<u8>>>> {
let captured: Arc<Mutex<Vec<Vec<u8>>>> = Arc::new(Mutex::new(Vec::new()));
let captured_clone = captured.clone();
Mock::given(method("POST"))
.and(path("/v1/messages"))
.respond_with(move |req: &wiremock::Request| {
captured_clone.lock().unwrap().push(req.body.clone());
ResponseTemplate::new(200).set_body_string(r#"{"ok":true}"#)
})
.mount(upstream)
.await;
captured
}
fn anthropic_payload(system: &str) -> Value {
json!({
"model": "claude-3-5-sonnet-20241022",
"max_tokens": 1024,
"system": system,
"messages": [
{"role": "user", "content": "hello"},
],
})
}
/// The cache-drift integration test installs a global JSON tracing
/// subscriber. Running it in its own `#[test]` (not `#[tokio::test]`)
/// would deadlock the wiremock client; instead we keep it in a
/// dedicated module that owns the OnceLock'd subscriber and is the
/// only async test in this binary.
mod tracing_capture {
use super::*;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::OnceLock;
use tracing_subscriber::fmt::MakeWriter;
#[derive(Clone)]
struct CaptureWriter {
inner: Arc<StdMutex<Vec<u8>>>,
}
impl std::io::Write for CaptureWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.inner.lock().unwrap().extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'a> MakeWriter<'a> for CaptureWriter {
type Writer = Self;
fn make_writer(&'a self) -> Self::Writer {
self.clone()
}
}
fn buffer() -> &'static Arc<StdMutex<Vec<u8>>> {
static BUFFER: OnceLock<Arc<StdMutex<Vec<u8>>>> = OnceLock::new();
BUFFER.get_or_init(|| {
let buf = Arc::new(StdMutex::new(Vec::new()));
let writer = CaptureWriter { inner: buf.clone() };
// INFO level so we also catch `cache_drift_first_request`,
// not just the warn-level `cache_drift_observed`.
let subscriber = tracing_subscriber::fmt()
.json()
.with_writer(writer)
.with_max_level(tracing::Level::INFO)
.finish();
let _ = tracing::subscriber::set_global_default(subscriber);
buf
})
}
#[tokio::test]
async fn cache_drift_observed_when_system_prompt_changes_mid_session() {
let buf = buffer();
buf.lock().unwrap().clear();
let upstream = MockServer::start().await;
let captured = mount_anthropic_capture_all(&upstream).await;
let proxy = start_proxy_with(&upstream.uri(), |c| {
// Drift detection runs inside the buffered branch — the
// master `compression` switch must be ON, which it is in
// every realistic deployment.
c.compression = true;
})
.await;
// Same Authorization header → same session_key. Different
// system prompts on each turn → drift_dims=system on turn 2.
let secret = "Bearer sk-test-this-is-a-secret";
let client = reqwest::Client::new();
let body1 = serde_json::to_vec(&anthropic_payload("you are an expert assistant")).unwrap();
let r1 = client
.post(format!("{}/v1/messages", proxy.url()))
.header("authorization", secret)
.header("content-type", "application/json")
.body(body1.clone())
.send()
.await
.unwrap();
assert_eq!(r1.status(), 200);
let body2 = serde_json::to_vec(&anthropic_payload("you are now a poet")).unwrap();
let r2 = client
.post(format!("{}/v1/messages", proxy.url()))
.header("authorization", secret)
.header("content-type", "application/json")
.body(body2.clone())
.send()
.await
.unwrap();
assert_eq!(r2.status(), 200);
// Byte-faithful passthrough: each upstream-received body must
// SHA-256 match the corresponding inbound body.
let received = captured.lock().unwrap().clone();
assert_eq!(received.len(), 2, "upstream should have seen 2 requests");
assert_eq!(
sha256_hex(&body1),
sha256_hex(&received[0]),
"request 1 byte-faithful passthrough violated",
);
assert_eq!(
sha256_hex(&body2),
sha256_hex(&received[1]),
"request 2 byte-faithful passthrough violated",
);
// Logs: a `cache_drift_observed` event must be present and
// include `system` in `drift_dims`.
let logs = String::from_utf8(buf.lock().unwrap().clone()).expect("logs are utf-8");
assert!(
logs.contains(r#""event":"cache_drift_first_request""#),
"expected first_request event in logs: {logs}",
);
assert!(
logs.contains(r#""event":"cache_drift_observed""#),
"expected drift_observed event in logs: {logs}",
);
// `drift_dims` should include `system` when only the system
// prompt mutated. Find any `cache_drift_observed` line and
// assert its `drift_dims` contains `system`.
let drift_line = logs
.lines()
.find(|line| line.contains(r#""event":"cache_drift_observed""#))
.expect("drift_observed line missing");
assert!(
drift_line.contains(r#""drift_dims":"system""#),
"expected drift_dims=system in drift line: {drift_line}",
);
// Privacy invariant: the raw bearer secret must NEVER appear
// anywhere in the captured logs.
assert!(
!logs.contains("sk-test-this-is-a-secret"),
"raw bearer secret leaked into logs",
);
assert!(
!logs.contains("Bearer sk-test"),
"raw 'Bearer ...' leaked into logs",
);
proxy.shutdown().await;
}
}