github-featured/chopratejas--headroom
2
0
Fork 0
Code Issues Pull Requests Actions 16 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix(proxy): port session-sticky beta headers to the Rust proxy (#2381)

Browse Source 
## Description

The Python proxy protects prompt caches with `SessionBetaTracker`
(PR-A6, `headroom/proxy/helpers.py`): interactive clients (Claude Code,
Codex CLI) may drop an `anthropic-beta` / `openai-beta` token between
turn N and turn N+1 of the same conversation, and since beta headers are
part of the bytes that determine the upstream prefix-cache key, the drop
rotates the key and the provider re-writes the whole prefix at the
customer's cost. The tracker unions the client's tokens with everything
previously seen for that `(provider, session)` and forwards the union —
a documented operator contract (`docs/configuration.mdx`, "Session Beta
Header Tracking").

The Rust proxy has no equivalent, and Phase H (#2258) deletes the
tracker together with `helpers.py` and its test file
(`tests/test_anthropic_beta_session_sticky.py`). None of the Phase A–G
plans port it (Phase F consumes beta headers for auth-mode
classification only), so the protection would silently not survive the
migration — and the Phase-H gate "Cache-hit-rate parity with direct
upstream confirmed" can't catch the loss, because re-injection makes
proxied traffic *beat* direct upstream on cache hits; when the mechanism
disappears, proxied traffic degrades *to* direct-upstream levels, which
that comparison reads as parity.

This PR ports the tracker semantics into the Rust proxy so the
protection lives in the codebase Phase H keeps.

Closes #2380

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] 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)

(New Rust functionality, but a parity port of already-shipped,
already-documented Python behavior — the PR title uses `fix:` per
`REALIGNMENT/INDEX.md`: "Commit prefix: `fix:` for Rust-migration phase
commits".)

## Changes Made

- **`cache_stabilization/beta_sticky.rs`** — the tracker: bounded LRU
(1000 sessions, same sizing rationale and `# Panics` contract as the
drift detector's capacity) keyed by `(provider, session)`, storing the
per-session ordered token list. Union preserves first-seen order; dedup
is case-insensitive with first-seen casing winning; lookups touch
recency; overflow evicts the oldest — mirroring the Python tracker. The
header-plumbing lives in the module too (`apply_sticky_betas`), so the
merge is unit-testable without booting a proxy.
- **`proxy.rs` wiring** — on the intercepted POST routes
(`/v1/messages`, `/v1/chat/completions`, `/v1/responses`), right after
the drift-detector observation, reusing the drift detector's
`derive_session_key` output so both cache-stability subsystems agree on
conversation identity.
- **`config.rs`** — `--beta-header-sticky` /
`HEADROOM_PROXY_BETA_HEADER_STICKY` (`enabled` default; `disabled`
forwards the client value verbatim and keeps no state), mirroring the
`StripInternalHeaders` flag pattern and the existing `HEADROOM_*` →
`HEADROOM_PROXY_*` Python→Rust env pairing. Since the merge runs inside
the compression interceptor, startup logs a warning when the flag is
`enabled` while `--compression` is off, and both the CLI doc and the
docs row state the dependency.
- **`tests/integration_beta_header_sticky.rs`** — 9 end-to-end tests
against a wiremock upstream asserting the headers/bytes the upstream
actually receives; 21 unit tests port the behavioral contract from
`tests/test_anthropic_beta_session_sticky.py` and cover the header-map
plumbing.
- **`docs/content/docs/configuration.mdx`** — one row for
`HEADROOM_PROXY_BETA_HEADER_STICKY` next to the existing Python/Rust
flag pairs.

## Testing

- [x] Unit tests pass (`cargo test -p headroom-proxy`; Python side via
`make ci-precheck-python` — `pytest` subset, 174 passed)
- [x] Linting passes (`cargo clippy --all-targets` — 0 warnings; `cargo
fmt --check` clean; Rust-only change, so `ruff`/`mypy` are covered by
the untouched-Python `ci-precheck-python` build)
- [ ] Type checking passes (`mypy headroom`) — N/A, no Python files
touched
- [x] New tests added for new functionality
- [x] Manual testing performed (RED/GREEN before-and-after runs below)

### Test Output

```text
$ cargo test -p headroom-proxy --lib beta_sticky
test result: ok. 21 passed; 0 failed; 0 ignored; 0 measured; 248 filtered out; finished in 0.03s

$ cargo test -p headroom-proxy --test integration_beta_header_sticky
test result: ok. 9 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.07s

$ cargo test -p headroom-proxy            # full crate: 37 suites, all ok
$ cargo clippy -p headroom-proxy --all-targets   # 0 warnings
$ make ci-precheck-rust ci-precheck-python ci-precheck-commitlint   # green
```

## Real Behavior Proof

- Environment: macOS arm64 (Darwin 24.6), `rustc 1.95.0`, real Rust
proxy booted on an ephemeral port in front of a wiremock upstream
(`tests/common::start_proxy_with`, `compression = true`).
- Exact command / steps: two-turn conversation through the proxy — turn
1 `POST /v1/messages` with `anthropic-beta:
context-management-2025-06-27,interleaved-thinking-2025-05-14`; turn 2,
same conversation, client drops the second token. The wiremock responder
captures the headers the upstream actually receives (`cargo test -p
headroom-proxy --test integration_beta_header_sticky`).
- Observed result: **before** the port (test written first, run against
the unmodified proxy) the upstream sees the shrunken token set and the
prefix-cache key rotates —

  ```text
assertion `left == right` failed: turn 2 must re-inject the dropped
token so the upstream
  prefix-cache key stays byte-stable
    left: Some("context-management-2025-06-27")
right:
Some("context-management-2025-06-27,interleaved-thinking-2025-05-14")
  ```

**After** the port the same scenario passes: the upstream receives the
full union on turn 2, the internal `x-headroom-session-id` never crosses
the upstream boundary, and the forwarded body is SHA-256-identical to
what the client sent (asserted by
`body_bytes_stay_byte_equal_while_header_is_rewritten`).
- Not tested: live traffic against a real provider upstream (wiremock
only); the WebSocket path and Bedrock/Vertex routes (out of scope — see
Additional Notes).

## 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 did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Screenshots (if applicable)

N/A (proxy behavior; see Real Behavior Proof).

## Additional Notes

Design decisions, and where I'd like reviewer judgment:

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