mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
c81755c965
commit
ce37940d17
8 changed files with 1027 additions and 29 deletions
12
Cargo.lock
generated
12
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
683
crates/headroom-proxy/src/cache_stabilization/drift_detector.rs
Normal file
683
crates/headroom-proxy/src/cache_stabilization/drift_detector.rs
Normal 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(¤t),
|
||||
"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, ¤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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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 => {
|
||||
|
|
|
|||
210
crates/headroom-proxy/tests/integration_cache_drift.rs
Normal file
210
crates/headroom-proxy/tests/integration_cache_drift.rs
Normal 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;
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue