mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
f6398a6476
## 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. |
||
|---|---|---|
| .. | ||
| agent-orchestration.mdx | ||
| agno.mdx | ||
| anthropic-sdk.mdx | ||
| api-reference.mdx | ||
| architecture.mdx | ||
| autogen.mdx | ||
| benchmarks.mdx | ||
| cache-optimization.mdx | ||
| ccr.mdx | ||
| ci-cd-flows.mdx | ||
| claude-code-azure-foundry.mdx | ||
| claude-code-vertex.mdx | ||
| code-compression.mdx | ||
| codex-recovery.mdx | ||
| configuration.mdx | ||
| context-management.mdx | ||
| crewai.mdx | ||
| docker-install.mdx | ||
| errors.mdx | ||
| failure-learning.mdx | ||
| filesystem-contract.mdx | ||
| grok-build.mdx | ||
| how-compression-works.mdx | ||
| image-compression.mdx | ||
| index.mdx | ||
| installation.mdx | ||
| langchain.mdx | ||
| limitations.mdx | ||
| litellm.mdx | ||
| local-llm-prefill.mdx | ||
| mcp.mdx | ||
| memory.mdx | ||
| meta.json | ||
| metrics.mdx | ||
| openai-sdk.mdx | ||
| opencode.mdx | ||
| persistent-installs.mdx | ||
| pipeline-extensions.mdx | ||
| proxy.mdx | ||
| quickstart.mdx | ||
| releases.mdx | ||
| savings.mdx | ||
| shared-context.mdx | ||
| simulation.mdx | ||
| smart-crusher.mdx | ||
| strands.mdx | ||
| text-and-logs.mdx | ||
| troubleshooting.mdx | ||
| vercel-ai-sdk.mdx | ||
| vscode-claude-code.mdx | ||
| vscode-copilot.mdx | ||