akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy): key drift detector on conversations, not credentials; canonicalize drift hashes (#2301)

Browse source 
## 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 <chopratejas@gmail.com>
This commit is contained in:
Andrei Boldyrev 2026-07-17 02:34:26 +05:00 • committed by GitHub
parent 517bf992cf
commit 6744833afe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 621 additions and 87 deletions
Download patch file Download diff file Expand all files Collapse all files

688
crates/headroom-proxy/src/cache_stabilization/drift_detector.rs
View file

@ -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::<Vec<_>>(),
// `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<serde_json::Value> = 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 =
/// "<comma-joined>", 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, &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);
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(&current),
"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.

3
crates/headroom-proxy/src/cache_stabilization/mod.rs
View file

@ -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` /

17
crates/headroom-proxy/src/proxy.rs
View file

@ -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);
}