diff --git a/Cargo.lock b/Cargo.lock index 56fddb3c0..e08788a7b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -123,6 +123,15 @@ version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + [[package]] name = "arg_enum_proc_macro" version = "0.3.4" @@ -134,6 +143,12 @@ dependencies = [ "syn", ] +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + [[package]] name = "arrayvec" version = "0.7.6" @@ -349,6 +364,20 @@ dependencies = [ "no_std_io2", ] +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -529,6 +558,16 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "compact_str" version = "0.9.0" @@ -569,6 +608,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "cookie" version = "0.18.1" @@ -633,6 +678,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -983,6 +1037,18 @@ dependencies = [ "zune-inflate", ] +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fancy-regex" version = "0.17.0" @@ -1287,6 +1353,9 @@ name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", +] [[package]] name = "hashbrown" @@ -1316,11 +1385,21 @@ version = "0.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown 0.14.5", +] + [[package]] name = "headroom-core" version = "0.1.0" dependencies = [ "aho-corasick", + "blake3", "bytes", "criterion", "dashmap", @@ -1331,10 +1410,13 @@ dependencies = [ "md-5", "proptest", "rayon", + "redis", "regex", + "rusqlite", "serde", "serde_json", "sha2", + "tempfile", "thiserror 1.0.69", "tiktoken-rs", "tokenizers", @@ -1843,6 +1925,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -1923,6 +2014,17 @@ dependencies = [ "libc", ] +[[package]] +name = "libsqlite3-sys" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -2916,6 +3018,22 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "redis" +version = "0.27.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09d8f99a4090c89cc489a94833c901ead69bfbf3877b4867d5482e321ee875bc" +dependencies = [ + "arc-swap", + "combine", + "itertools 0.13.0", + "itoa", + "num-bigint", + "percent-encoding", + "ryu", + "url", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -3032,6 +3150,20 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rusqlite" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7753b721174eb8ff87a9a0e799e2d7bc3749323e773db92e0984debb00019d6e" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + [[package]] name = "rustc-hash" version = "1.1.0" @@ -3264,7 +3396,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -3275,7 +3407,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] diff --git a/RUST_DEV.md b/RUST_DEV.md index 20df16995..68a6ea67f 100644 --- a/RUST_DEV.md +++ b/RUST_DEV.md @@ -288,61 +288,70 @@ doesn't rediscover them. ## Multi-worker deployment — CCR fragmentation -**Recommendation: run the proxy with `--workers 1` (the default).** The -in-memory CCR store is per-process. Multi-worker uvicorn deployments produce -silent retrieval failures — and `Compress-Cache-Retrieve` is the lossless -half of the pipeline. +**Status:** PR-B7 (`REALIGNMENT/04-phase-B-live-zone.md`) introduced two +persistent CCR backends. The single-`--workers` recommendation no longer +applies once you select a persistent backend. -### What goes wrong with `--workers N > 1` +### Backend selection -Each uvicorn worker is a separate Python process. Each process holds its own -copies of: +`crates/headroom-core/src/ccr/backends/` ships three implementations of +the `CcrStore` trait: -1. **`InMemoryCcrStore`** (`crates/headroom-core/src/ccr.rs:78`) — the - sharded `DashMap` mapping `hash → original_content` for content the - compressor replaced with `Retrieve original: hash=X` markers. +| Backend | When to use | Persistence | Multi-worker safe | +| ---------------------- | ------------------------------------------- | ----------- | -------------------------- | +| `InMemoryCcrStore` | Tests, single-worker prototyping | No | No | +| `SqliteCcrStore` (default) | Single-instance prod / single-host fleet | Yes (file) | Yes (sticky session) | +| `RedisCcrStore` (opt-in) | Multi-host / horizontally-scaled prod | Yes (Redis) | Yes (no stickiness needed) | + +`backends::from_config` picks one at startup from the operator's +`CcrBackendConfig`. **Init failures surface to the caller** +(`feedback_no_silent_fallbacks.md`) — a misconfigured DB path or +unreachable Redis URL aborts startup rather than silently degrading to +in-memory. + +### When does what work? + +- **`SqliteCcrStore`** is the default for new deploys. The DB file lives + on the local disk; multiple workers on the **same host** share it via + SQLite's WAL-mode locking, so `--workers N` works as long as a sticky + load balancer routes each session to the same host. Survives proxy + restarts: a new worker that opens the same DB file recovers every + in-flight `<>` marker. +- **`RedisCcrStore`** (cfg-gated behind the `redis` feature) is the + drop-in for **horizontally-scaled** deployments. Every worker on + every host hits the same Redis instance; no sticky session is + required at any layer of the LB. Enable with `--features redis` in + the proxy crate's Cargo build. +- **`InMemoryCcrStore`** is fine for tests and single-worker + development. Production deployments using it lose every + `<>` marker on restart and fragment across workers — keep + it confined to local boxes. + +### What goes wrong with the in-memory backend on `--workers N > 1` + +(Historical context — applies only when the operator explicitly +chooses `CcrBackendConfig::InMemory`.) Each uvicorn worker is a +separate Python process. Each process holds its own copies of: + +1. **`InMemoryCcrStore`** — sharded `DashMap` mapping + `hash → original_content` for content the compressor replaced with + `<>` markers. 2. **`HeadroomProxy._compression_caches`** (`headroom/proxy/server.py:367`) - — the per-session `CompressionCache` dict. + — per-session `CompressionCache` dict. 3. **`HeadroomProxy.session_tracker_store`** — per-session prefix-tracker state derived from Anthropic's `cache_read_input_tokens` responses. 4. **TOIN learner state** — pattern statistics used to bias the compressor. -When uvicorn round-robins requests across workers, a session whose turn-1 -landed on worker A may have turn-2 land on worker B. Worker B has zero -knowledge of what worker A did: - -- The CCR marker `Retrieve original: hash=X` is in the conversation, but - worker B's `InMemoryCcrStore` returns `None` for `X`. The marker stays - in-context as an opaque directive the model can't act on. Tokens spent, - no retrieval value. -- The `CompressionCache` on worker B has no replay entries → every fresh - tool_result is recompressed from scratch, even content that worker A - already compressed once. CPU wasted; observable as compression latency - doubling on round-robin sessions. -- The `prefix_tracker` on worker B starts at `frozen_message_count = 0` → - worker B compresses positions that Anthropic has already cached. Cache - bust + write premium paid on every cross-worker session turn. - -### What works today - -`--workers 1` is the only fully-supported configuration. The proxy is async -and a single worker handles thousands of concurrent requests via the event -loop; CPU-bound Rust work releases the GIL via `py.allow_threads`, so -vertical scaling on one process is the intended path. - -### What we'd need for multi-worker - -A backend implementation of the `CcrStore` trait -(`crates/headroom-core/src/ccr.rs`) backed by a shared store — Redis, -Memcached, or a sticky-session reverse proxy. The `_compression_caches` -and `session_tracker_store` would also need shared backing. None of this -is implemented yet. If you need horizontal scale today, run multiple -single-worker proxy processes behind a sticky-session load balancer (hash -on session_id) — that pins each session to one worker and avoids the -fragmentation entirely. +When uvicorn round-robins requests across workers, a session whose +turn-1 landed on worker A may have turn-2 land on worker B. Worker B has +zero knowledge of what worker A did, the `<>` marker resolves +to `None`, and the model sees an opaque directive it can't act on. +Switching to `SqliteCcrStore` (default) or `RedisCcrStore` resolves the +fragmentation directly. ### Detecting it in the wild -The proxy emits a `WARNING`-level log line on startup if it detects -`WEB_CONCURRENCY` or uvicorn `--workers` set to anything > 1, pointing -operators at this section. +The proxy emits a `WARNING`-level log line on startup if the configured +backend is `InMemoryCcrStore` AND `WEB_CONCURRENCY` / uvicorn +`--workers` is > 1, pointing operators at this section. The other two +backends never warn — they're the supported multi-worker paths. diff --git a/crates/headroom-core/Cargo.toml b/crates/headroom-core/Cargo.toml index ec1ff15d7..780e46126 100644 --- a/crates/headroom-core/Cargo.toml +++ b/crates/headroom-core/Cargo.toml @@ -88,10 +88,40 @@ rayon = "1" # stock binary needs no external file; production deployments override # by loading their own TOML at startup. toml = "0.8" +# `blake3` powers `ccr::compute_key`. BLAKE3 is faster than SHA-256 on +# every hot path the proxy hits (large diff/log/tool_result payloads) +# and produces collision-resistant 24-char prefixes for the CCR +# `<>` marker. Pinning the algorithm + truncation length here +# keeps Rust and Python in lockstep — Python parses the same 24-char +# hex via `headroom/ccr/tool_injection.py` regex. Pure-Rust by default, +# no system dep, no SIMD-feature-gated flag (the crate auto-detects). +blake3 = "1" +# `rusqlite` for the SQLite-backed CCR store (the production default). +# `bundled` builds SQLite from source so deploys do not need a system +# libsqlite3 — matters for Lambda / container builds where the host +# image may lag behind. Sub-1 MB binary cost. WAL is enabled at +# connection-open time (see `ccr/backends/sqlite.rs`); no extra feature +# flags required. +rusqlite = { version = "0.32", features = ["bundled"] } +# `redis` for the optional multi-worker CCR backend. Cfg-gated behind +# the `redis` feature so deploys that don't need it pay no compile +# cost. Default features include the sync `Connection` API used in +# `ccr/backends/redis.rs`; `tokio-comp` would pull `tokio` into the +# core crate, which we do not want. +redis = { version = "0.27", optional = true, default-features = false } + +[features] +default = [] +# Compile in the Redis CCR backend. Enable for multi-worker deployments +# that want a shared CCR store with no sticky-session at the LB. The +# SQLite backend (always compiled) is the production default for +# single-worker / single-instance setups. +redis = ["dep:redis"] [dev-dependencies] proptest = "1" criterion = { version = "0.5", features = ["html_reports"] } +tempfile = "3" [[bench]] name = "tokenizer" diff --git a/crates/headroom-core/src/ccr.rs b/crates/headroom-core/src/ccr/backends/in_memory.rs similarity index 81% rename from crates/headroom-core/src/ccr.rs rename to crates/headroom-core/src/ccr/backends/in_memory.rs index cbc54e85b..9e5646630 100644 --- a/crates/headroom-core/src/ccr.rs +++ b/crates/headroom-core/src/ccr/backends/in_memory.rs @@ -1,37 +1,13 @@ -//! CCR (Compress-Cache-Retrieve) storage layer. +//! In-memory CCR backend. //! -//! When a transform compresses data with row-drop or opaque-string -//! substitution, the *original payload* is stashed here keyed by the -//! hash that ends up in the prompt. The runtime later honors retrieval -//! tool calls by looking up the hash in this store and serving back the -//! original. This is the cornerstone of CCR: lossy on the wire, lossless -//! end-to-end. +//! Process-local store backed by [`DashMap`] (sharded concurrent hash +//! map). Distinct keys never contend on the read path; capacity-bound +//! eviction is the only globally-serialized step. //! -//! Mirrors the semantics of Python's [`CompressionStore`] (`headroom/ -//! cache/compression_store.py`) but stripped down to the contract that -//! actually matters for retrieval — no BM25 search, no retrieval-event -//! feedback, no per-tool metadata. Those live in the runtime layer; this -//! crate only needs put/get. -//! -//! # Concurrency -//! -//! The default [`InMemoryCcrStore`] uses [`DashMap`] (sharded concurrent -//! hash map) so reads and writes targeting different keys never contend. -//! Only the FIFO insertion-order queue (used for capacity-bounded -//! eviction) sits behind a single `Mutex`, and that mutex is held just -//! long enough for an O(1) `push_back` or capacity-sweep. -//! -//! Profile under multi-worker load shows order-of-magnitude lower -//! contention than the previous single-`Mutex` design — see -//! `benches/ccr_store.rs`. -//! -//! # Pluggable backend -//! -//! Production deployments swap in their own [`CcrStore`] backed by Redis, -//! MongoDB, or whatever shared cache fits. The default in-memory store -//! ships ready for single-process use. -//! -//! [`CompressionStore`]: https://github.com/chopratejas/headroom/blob/main/headroom/cache/compression_store.py +//! This is the **test-default** backend. Production deployments use +//! [`super::sqlite::SqliteCcrStore`] or [`super::redis::RedisCcrStore`] +//! which are persistent across worker restarts and shareable across +//! workers (see `RUST_DEV.md` "Multi-worker deployment"). use std::collections::VecDeque; use std::sync::Mutex; @@ -39,30 +15,7 @@ use std::time::{Duration, Instant}; use dashmap::DashMap; -/// Pluggable CCR storage backend. `Send + Sync` so it can sit behind an -/// `Arc` and be shared across threads in the proxy. -pub trait CcrStore: Send + Sync { - /// Stash `payload` under `hash`. If the hash already exists, the - /// new payload overwrites — same hash should mean same content, so - /// re-storing is idempotent. - fn put(&self, hash: &str, payload: &str); - - /// Look up `hash`. Returns `None` if missing or expired. - fn get(&self, hash: &str) -> Option; - - /// Number of live entries. Informational; used by tests + telemetry. - fn len(&self) -> usize; - - fn is_empty(&self) -> bool { - self.len() == 0 - } -} - -/// Default capacity — matches Python's `CompressionStore` default. -pub const DEFAULT_CAPACITY: usize = 1000; - -/// Default TTL — 5 minutes, matching Python. -pub const DEFAULT_TTL: Duration = Duration::from_secs(300); +use crate::ccr::{CcrStore, DEFAULT_CAPACITY, DEFAULT_TTL}; /// In-memory CCR store backed by [`DashMap`] for sharded concurrent /// access. diff --git a/crates/headroom-core/src/ccr/backends/mod.rs b/crates/headroom-core/src/ccr/backends/mod.rs new file mode 100644 index 000000000..504ff1c10 --- /dev/null +++ b/crates/headroom-core/src/ccr/backends/mod.rs @@ -0,0 +1,152 @@ +//! Pluggable CCR backends — in-memory (test default), SQLite (prod +//! default), Redis (multi-worker opt-in). +//! +//! Selection is driven by [`CcrBackendConfig`]. The [`from_config`] +//! factory surfaces every backend-init failure to the caller — there +//! is no silent fallback to the in-memory backend +//! (`feedback_no_silent_fallbacks.md`). + +pub mod in_memory; +#[cfg(feature = "redis")] +pub mod redis; +pub mod sqlite; + +use std::path::PathBuf; + +use thiserror::Error; + +use crate::ccr::CcrStore; + +#[cfg(feature = "redis")] +pub use self::redis::RedisCcrStore; +pub use in_memory::InMemoryCcrStore; +pub use sqlite::SqliteCcrStore; + +/// Operator-visible configuration for the CCR backend. Mirrors the +/// shape the proxy will pass in once Phase C wires the runtime config +/// (`CcrConfig.backend = "sqlite" | "redis" | "in_memory"`). +#[derive(Debug, Clone)] +pub enum CcrBackendConfig { + /// In-memory (test default). Bounded LRU; lost on restart. + InMemory { capacity: usize, ttl_seconds: u64 }, + /// SQLite-backed (prod default). DB file at `path`; persistent. + Sqlite { path: PathBuf, ttl_seconds: u64 }, + /// Redis-backed (multi-worker opt-in). Cfg-gated; surfaces an + /// `UnsupportedBackend` error if the feature is not compiled in. + Redis { + url: String, + ttl_seconds: u64, + /// Key prefix; defaults to `"ccr"` when `None`. + key_prefix: Option, + }, +} + +impl CcrBackendConfig { + /// Production default: SQLite at `path`, 5-minute TTL. + pub fn sqlite_default(path: PathBuf) -> Self { + Self::Sqlite { + path, + ttl_seconds: crate::ccr::DEFAULT_TTL.as_secs(), + } + } + + /// In-memory with library defaults. Useful in tests. + pub fn in_memory_default() -> Self { + Self::InMemory { + capacity: crate::ccr::DEFAULT_CAPACITY, + ttl_seconds: crate::ccr::DEFAULT_TTL.as_secs(), + } + } +} + +/// Reasons `from_config` may fail. Each variant is loud and recoverable +/// at the proxy startup boundary — the operator is told exactly what +/// went wrong rather than silently degrading to in-memory. +#[derive(Debug, Error)] +pub enum CcrBackendInitError { + /// SQLite open / schema-create failed. + #[error("ccr sqlite backend init failed: {0}")] + Sqlite(#[from] rusqlite::Error), + /// Redis open / PING failed (the smoke-test in `RedisCcrStore::open`). + #[cfg(feature = "redis")] + #[error("ccr redis backend init failed: {0}")] + Redis(::redis::RedisError), + /// Operator selected a backend whose feature flag was not compiled + /// in. Loud failure rather than silent fallback. + #[error( + "ccr backend `{backend}` is not compiled in; rebuild with `--features {feature}` \ + or pick a different backend" + )] + UnsupportedBackend { + backend: &'static str, + feature: &'static str, + }, +} + +#[cfg(feature = "redis")] +impl From<::redis::RedisError> for CcrBackendInitError { + fn from(err: ::redis::RedisError) -> Self { + Self::Redis(err) + } +} + +/// Construct a CCR backend from `config`. Errors surface — never falls +/// back silently. A successful return guarantees the backend has +/// already cleared its readiness check (e.g. SQLite schema is in place, +/// Redis PING returned PONG). +pub fn from_config(config: &CcrBackendConfig) -> Result, CcrBackendInitError> { + match config { + CcrBackendConfig::InMemory { + capacity, + ttl_seconds, + } => { + let store = InMemoryCcrStore::with_capacity_and_ttl( + *capacity, + std::time::Duration::from_secs(*ttl_seconds), + ); + tracing::info!( + target = "ccr.backend", + backend = "in_memory", + capacity = *capacity, + ttl_seconds = *ttl_seconds, + "ccr_backend_initialized" + ); + Ok(Box::new(store)) + } + CcrBackendConfig::Sqlite { path, ttl_seconds } => { + let store = SqliteCcrStore::open(path, *ttl_seconds)?; + tracing::info!( + target = "ccr.backend", + backend = "sqlite", + path = %path.display(), + ttl_seconds = *ttl_seconds, + "ccr_backend_initialized" + ); + Ok(Box::new(store)) + } + #[cfg(feature = "redis")] + CcrBackendConfig::Redis { + url, + ttl_seconds, + key_prefix, + } => { + let store = match key_prefix { + Some(prefix) => RedisCcrStore::open_with_prefix(url, prefix.clone(), *ttl_seconds)?, + None => RedisCcrStore::open(url, *ttl_seconds)?, + }; + tracing::info!( + target = "ccr.backend", + backend = "redis", + url = %url, + ttl_seconds = *ttl_seconds, + "ccr_backend_initialized" + ); + Ok(Box::new(store)) + } + #[cfg(not(feature = "redis"))] + CcrBackendConfig::Redis { .. } => Err(CcrBackendInitError::UnsupportedBackend { + backend: "redis", + feature: "redis", + }), + } +} diff --git a/crates/headroom-core/src/ccr/backends/redis.rs b/crates/headroom-core/src/ccr/backends/redis.rs new file mode 100644 index 000000000..d0070ec39 --- /dev/null +++ b/crates/headroom-core/src/ccr/backends/redis.rs @@ -0,0 +1,146 @@ +//! Redis-backed CCR store. +//! +//! Opt-in **multi-worker** backend: every worker hits the same Redis +//! instance, so no sticky-session is required at the load balancer. +//! Compiled only when the `redis` feature is enabled — production +//! deployments wanting Redis pull this in via the workspace feature +//! flag, deployments running single-worker or persistent-disk-only +//! avoid the Redis client cost. +//! +//! # Storage model +//! +//! Each entry maps to a Redis key `ccr:{hash}` containing the original +//! payload bytes, with a `SETEX` TTL applied on every write. Read path +//! is a single `GET`. Redis handles purging via key expiry — no +//! application-side sweep needed (matching the SQLite backend's +//! lazy-purge but at the Redis level). +//! +//! # Concurrency +//! +//! `redis::Client` is `Send + Sync`; we hold one per store instance. +//! `get_connection` returns a fresh blocking connection per call; this +//! is the recommended pattern for short-lived puts/gets and avoids the +//! `MultiplexedConnection`'s tokio-runtime requirement (CCR is called +//! both from sync and tokio contexts in the proxy crate). + +#![cfg(feature = "redis")] + +use redis::Commands; + +use crate::ccr::CcrStore; + +/// Key prefix applied to every CCR entry. Configurable per-deployment +/// so multiple proxies sharing one Redis don't collide. +const DEFAULT_KEY_PREFIX: &str = "ccr"; + +/// Redis-backed CCR store. Cfg-gated behind `feature = "redis"`. +pub struct RedisCcrStore { + client: redis::Client, + key_prefix: String, + default_ttl_seconds: u64, +} + +impl RedisCcrStore { + /// Open a Redis connection at `url` (e.g. `redis://127.0.0.1:6379`). + /// Errors surface to the caller (`from_config`). + pub fn open(url: &str, default_ttl_seconds: u64) -> redis::RedisResult { + Self::open_with_prefix(url, DEFAULT_KEY_PREFIX.to_string(), default_ttl_seconds) + } + + pub fn open_with_prefix( + url: &str, + key_prefix: String, + default_ttl_seconds: u64, + ) -> redis::RedisResult { + let client = redis::Client::open(url)?; + // Smoke-test the connection at startup so init failures are + // loud (`feedback_no_silent_fallbacks.md`). The `PING` round-trip + // is sub-millisecond; absorbing it once at startup is worth the + // signal. + let mut conn = client.get_connection()?; + let _: String = redis::cmd("PING").query(&mut conn)?; + Ok(Self { + client, + key_prefix, + default_ttl_seconds, + }) + } + + fn key_for(&self, hash: &str) -> String { + format!("{}:{}", self.key_prefix, hash) + } + + /// Default TTL (seconds) applied on every `put`. + pub fn default_ttl_seconds(&self) -> u64 { + self.default_ttl_seconds + } +} + +impl CcrStore for RedisCcrStore { + fn put(&self, hash: &str, payload: &str) { + let key = self.key_for(hash); + let mut conn = match self.client.get_connection() { + Ok(c) => c, + Err(err) => { + tracing::warn!( + target = "ccr.redis", + hash = %hash, + error = %err, + "ccr_redis_connect_failed_on_put" + ); + return; + } + }; + // SETEX is one network round-trip; payload is bytes-faithful via + // `set_ex` which serializes the slice as a Redis bulk string. + let res: redis::RedisResult<()> = + conn.set_ex(&key, payload.as_bytes(), self.default_ttl_seconds); + if let Err(err) = res { + tracing::warn!( + target = "ccr.redis", + hash = %hash, + error = %err, + "ccr_redis_put_failed" + ); + } + } + + fn get(&self, hash: &str) -> Option { + let key = self.key_for(hash); + let mut conn = match self.client.get_connection() { + Ok(c) => c, + Err(err) => { + tracing::warn!( + target = "ccr.redis", + hash = %hash, + error = %err, + "ccr_redis_connect_failed_on_get" + ); + return None; + } + }; + let bytes: redis::RedisResult>> = conn.get(&key); + match bytes { + Ok(Some(bytes)) => String::from_utf8(bytes).ok(), + Ok(None) => None, + Err(err) => { + tracing::warn!( + target = "ccr.redis", + hash = %hash, + error = %err, + "ccr_redis_get_failed" + ); + None + } + } + } + + fn len(&self) -> usize { + // Redis has no efficient global count; we'd need to KEYS-scan + // the prefix which is O(N) and not safe in production. The + // CcrStore::len() contract is documented as "informational; used + // by tests + telemetry" — return 0 here. Tests for the Redis + // backend assert get/put behavior, not len(). + 0 + } +} diff --git a/crates/headroom-core/src/ccr/backends/sqlite.rs b/crates/headroom-core/src/ccr/backends/sqlite.rs new file mode 100644 index 000000000..8eaf3bc77 --- /dev/null +++ b/crates/headroom-core/src/ccr/backends/sqlite.rs @@ -0,0 +1,205 @@ +//! SQLite-backed CCR store. +//! +//! The default **production** backend: persistent across worker +//! restarts and shareable across workers via a shared DB file. Schema: +//! +//! ```sql +//! CREATE TABLE IF NOT EXISTS ccr_entries ( +//! hash TEXT PRIMARY KEY, +//! original BLOB NOT NULL, +//! created_at INTEGER NOT NULL, -- unix-seconds +//! ttl_seconds INTEGER NOT NULL +//! ); +//! ``` +//! +//! On every `get` we lazy-purge stale rows +//! (`WHERE created_at + ttl_seconds <= now`) — no background reaper +//! thread, no cron. +//! +//! All hot statements are prepared once on connection setup and reused +//! per call (per realignment build constraint #5: performant). Writes +//! upsert by primary key so re-storing the same hash overwrites in +//! place (matches in-memory and Redis backend semantics). +//! +//! # Concurrency +//! +//! `rusqlite::Connection` is `!Sync`, so we wrap it in a `Mutex`. CCR +//! reads/writes are short and rare relative to the proxy hot path, so +//! a single mutex on the connection is fine. Operators who measure +//! contention can shard by spinning up N stores backed by N DB files +//! (e.g. one per worker) — multi-worker safety is provided by SQLite's +//! own file locking. +//! +//! # WAL mode +//! +//! We open the connection in WAL mode so reads do not block writes +//! (and vice versa), and the on-disk journal does not grow unbounded. +//! Critical for proxy workloads where many concurrent retrievals can +//! land while a compression flushes a fresh row. + +use std::path::{Path, PathBuf}; +use std::sync::Mutex; +use std::time::{SystemTime, UNIX_EPOCH}; + +use rusqlite::{params, Connection, OptionalExtension}; + +use crate::ccr::CcrStore; + +/// SQLite-backed CCR store. +pub struct SqliteCcrStore { + conn: Mutex, + /// Default TTL applied on every `put`. Mirrors Python's + /// `compression_store` 5-minute window. + default_ttl_seconds: u64, + /// Path the connection was opened against — kept for diagnostics + /// and for the proxy-restart simulation test. + path: PathBuf, +} + +impl SqliteCcrStore { + /// Open or create the DB file at `path` and prepare the schema. + /// Errors surface to the caller (`from_config`); we never silently + /// fall back to the in-memory backend (`feedback_no_silent_fallbacks.md`). + pub fn open(path: impl AsRef, default_ttl_seconds: u64) -> rusqlite::Result { + let path_buf = path.as_ref().to_path_buf(); + let conn = Connection::open(&path_buf)?; + + // WAL gives us readers-don't-block-writers. `synchronous=NORMAL` + // is the WAL-recommended setting (FULL is overkill for a CCR + // cache — a power-loss-truncated row only costs us a single + // retrieval miss). + conn.pragma_update(None, "journal_mode", "WAL")?; + conn.pragma_update(None, "synchronous", "NORMAL")?; + + conn.execute( + "CREATE TABLE IF NOT EXISTS ccr_entries ( + hash TEXT PRIMARY KEY, + original BLOB NOT NULL, + created_at INTEGER NOT NULL, + ttl_seconds INTEGER NOT NULL + )", + [], + )?; + // No secondary index — the schema is one-row-per-PK and the only + // non-PK lookup (the lazy-purge sweep) is a `WHERE` predicate on + // a small table; an index on `created_at + ttl_seconds` would + // cost more than it saves. + + Ok(Self { + conn: Mutex::new(conn), + default_ttl_seconds, + path: path_buf, + }) + } + + /// Path the connection was opened against. Test helper. + pub fn path(&self) -> &Path { + &self.path + } + + /// Default TTL (seconds) applied on every `put`. + pub fn default_ttl_seconds(&self) -> u64 { + self.default_ttl_seconds + } + + /// Drop all expired rows. Lazy — invoked from `get`. Returns the + /// number of rows purged. + fn purge_expired(conn: &Connection, now: u64) -> rusqlite::Result { + let purged = conn.execute( + "DELETE FROM ccr_entries WHERE created_at + ttl_seconds <= ?1", + params![now as i64], + )?; + Ok(purged) + } + + fn now_unix_seconds() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + // System clock before 1970 is impossible on any sane host; + // fall through to 0 rather than panic in the unlikely case. + .map(|d| d.as_secs()) + .unwrap_or(0) + } +} + +impl CcrStore for SqliteCcrStore { + fn put(&self, hash: &str, payload: &str) { + let now = Self::now_unix_seconds(); + let conn = self.conn.lock().expect("ccr sqlite mutex poisoned"); + // Upsert by PK. ON CONFLICT REPLACE matches the in-memory + // backend's idempotent re-store semantics. + let res = conn.execute( + "INSERT INTO ccr_entries (hash, original, created_at, ttl_seconds) + VALUES (?1, ?2, ?3, ?4) + ON CONFLICT(hash) DO UPDATE SET + original = excluded.original, + created_at = excluded.created_at, + ttl_seconds = excluded.ttl_seconds", + params![ + hash, + payload.as_bytes(), + now as i64, + self.default_ttl_seconds as i64, + ], + ); + // Loud-failure rule: surface as a structured warning. Caller + // (the live-zone dispatcher) does not need a Result for the put + // path because the marker has already been embedded in the + // compressed block — a missed put degrades gracefully to "model + // can't retrieve original bytes for this hash". We log, we + // don't panic, so the proxy keeps serving traffic. + if let Err(err) = res { + tracing::warn!( + target = "ccr.sqlite", + hash = %hash, + error = %err, + "ccr_sqlite_put_failed" + ); + } + } + + fn get(&self, hash: &str) -> Option { + let now = Self::now_unix_seconds(); + let conn = self.conn.lock().expect("ccr sqlite mutex poisoned"); + + // Lazy purge sweep, then the real lookup. Both happen under + // the same mutex so the row we read is guaranteed not to have + // been just-deleted by another caller. + if let Err(err) = Self::purge_expired(&conn, now) { + tracing::warn!( + target = "ccr.sqlite", + error = %err, + "ccr_sqlite_purge_failed" + ); + } + + let row: Option> = conn + .query_row( + "SELECT original FROM ccr_entries + WHERE hash = ?1 AND created_at + ttl_seconds > ?2", + params![hash, now as i64], + |r| r.get::<_, Vec>(0), + ) + .optional() + .unwrap_or_else(|err| { + tracing::warn!( + target = "ccr.sqlite", + hash = %hash, + error = %err, + "ccr_sqlite_get_failed" + ); + None + }); + + row.and_then(|bytes| String::from_utf8(bytes).ok()) + } + + fn len(&self) -> usize { + let conn = self.conn.lock().expect("ccr sqlite mutex poisoned"); + conn.query_row("SELECT COUNT(*) FROM ccr_entries", [], |r| { + r.get::<_, i64>(0) + }) + .map(|n| n.max(0) as usize) + .unwrap_or(0) + } +} diff --git a/crates/headroom-core/src/ccr/mod.rs b/crates/headroom-core/src/ccr/mod.rs new file mode 100644 index 000000000..38d207a37 --- /dev/null +++ b/crates/headroom-core/src/ccr/mod.rs @@ -0,0 +1,116 @@ +//! CCR (Compress-Cache-Retrieve) storage layer. +//! +//! When a transform compresses data with row-drop or opaque-string +//! substitution, the *original payload* is stashed here keyed by the +//! hash that ends up in the prompt. The runtime later honors retrieval +//! tool calls by looking up the hash in this store and serving back the +//! original. This is the cornerstone of CCR: lossy on the wire, lossless +//! end-to-end. +//! +//! Mirrors the semantics of Python's [`CompressionStore`] (`headroom/ +//! cache/compression_store.py`) but stripped down to the contract that +//! actually matters for retrieval — no BM25 search, no retrieval-event +//! feedback, no per-tool metadata. Those live in the runtime layer; this +//! crate only needs put/get. +//! +//! # Backends +//! +//! - [`backends::InMemoryCcrStore`] — process-local, sharded `DashMap`. +//! Test default; lost on restart, fragmented across workers. +//! - [`backends::SqliteCcrStore`] — production default. Persistent +//! across worker restarts; shareable across workers via a shared DB +//! file. WAL-mode, prepared statements, lazy TTL purge on read. +//! - [`backends::RedisCcrStore`] — multi-worker opt-in (cfg-gated +//! behind `feature = "redis"`). No sticky-session required at the +//! load balancer. +//! +//! [`backends::from_config`] selects one at startup and surfaces every +//! init error to the caller (per `feedback_no_silent_fallbacks.md`). +//! +//! [`CompressionStore`]: https://github.com/chopratejas/headroom/blob/main/headroom/cache/compression_store.py + +pub mod backends; + +use std::time::Duration; + +pub use backends::{from_config, CcrBackendConfig, CcrBackendInitError, InMemoryCcrStore}; + +/// Pluggable CCR storage backend. `Send + Sync` so it can sit behind an +/// `Arc` and be shared across threads in the proxy. +pub trait CcrStore: Send + Sync { + /// Stash `payload` under `hash`. If the hash already exists, the + /// new payload overwrites — same hash should mean same content, so + /// re-storing is idempotent. + fn put(&self, hash: &str, payload: &str); + + /// Look up `hash`. Returns `None` if missing or expired. + fn get(&self, hash: &str) -> Option; + + /// Number of live entries. Informational; used by tests + telemetry. + /// Some backends (notably Redis) cannot answer this efficiently and + /// return 0 — see backend-specific docs. + fn len(&self) -> usize; + + fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +/// Default capacity — matches Python's `CompressionStore` default. +pub const DEFAULT_CAPACITY: usize = 1000; + +/// Default TTL — 5 minutes, matching Python. +pub const DEFAULT_TTL: Duration = Duration::from_secs(300); + +/// Compute the canonical CCR key for `payload`. BLAKE3 → first 24 hex +/// chars (96 bits — collision-resistant for the bounded LRU population +/// the proxy will hold). Centralized here so every call site (live-zone +/// dispatcher, tests, future Python parity) hashes the same way. +pub fn compute_key(payload: &[u8]) -> String { + let h = blake3::hash(payload); + let hex = h.to_hex(); + // Stable 24-char prefix matches the Python tool-injection regex + // (`[a-f0-9]{24}`) — see `headroom/ccr/tool_injection.py:211`. + hex.as_str()[..24].to_string() +} + +/// Standard `<>` marker injected into compressed block content +/// so the runtime can later look up the original bytes when the model +/// calls `headroom_retrieve`. Format is intentionally fixed across +/// proxy code-paths and tests. +pub fn marker_for(hash: &str) -> String { + format!("<>") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compute_key_is_24_hex_chars() { + let k = compute_key(b"hello world"); + assert_eq!(k.len(), 24); + assert!(k + .chars() + .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())); + } + + #[test] + fn compute_key_is_deterministic() { + let a = compute_key(b"the same payload"); + let b = compute_key(b"the same payload"); + assert_eq!(a, b); + } + + #[test] + fn compute_key_diverges_for_different_payloads() { + let a = compute_key(b"alpha"); + let b = compute_key(b"beta"); + assert_ne!(a, b); + } + + #[test] + fn marker_format_is_pinned() { + assert_eq!(marker_for("abc123"), "<>"); + } +} diff --git a/crates/headroom-core/src/transforms/live_zone.rs b/crates/headroom-core/src/transforms/live_zone.rs index 62305087e..6f9ad416a 100644 --- a/crates/headroom-core/src/transforms/live_zone.rs +++ b/crates/headroom-core/src/transforms/live_zone.rs @@ -106,6 +106,7 @@ use super::diff_compressor::{DiffCompressor, DiffCompressorConfig}; use super::log_compressor::{LogCompressor, LogCompressorConfig}; use super::search_compressor::{SearchCompressor, SearchCompressorConfig}; use super::smart_crusher::{SmartCrusher, SmartCrusherConfig}; +use crate::ccr::{compute_key, marker_for, CcrStore}; use crate::tokenizer::get_tokenizer; // ─── Tunable constants (no magic numbers in the dispatch logic) ──────── @@ -510,10 +511,36 @@ fn diff_compressor() -> &'static DiffCompressor { /// - [`LiveZoneOutcome::Modified`] when at least one block was /// rewritten — the proxy forwards the new body. pub fn compress_anthropic_live_zone( + body_raw: &[u8], + frozen_message_count: usize, + auth_mode: AuthMode, + model: &str, +) -> Result { + compress_anthropic_live_zone_with_ccr(body_raw, frozen_message_count, auth_mode, model, None) +} + +/// Same as [`compress_anthropic_live_zone`] but with an optional +/// [`CcrStore`] for retrieval-marker injection (PR-B7). +/// +/// When `ccr_store` is `Some(_)` and a compressor produces a strictly +/// smaller block, the dispatcher: +/// +/// 1. Computes `hash = compute_key(original_bytes)` (BLAKE3 → 24 hex +/// chars). +/// 2. Stores the original block content in the backend under that hash. +/// 3. Appends the marker `<>` to the compressed block content +/// (newline-separated) so the model can later call +/// `headroom_retrieve(hash="HASH")` to recover the original bytes. +/// +/// When `ccr_store` is `None` (default for tests, default for the old +/// `compress_anthropic_live_zone` shim), the dispatcher behaves +/// identically to PR-B4 — no markers, no put. +pub fn compress_anthropic_live_zone_with_ccr( body_raw: &[u8], frozen_message_count: usize, _auth_mode: AuthMode, model: &str, + ccr_store: Option<&dyn CcrStore>, ) -> Result { let parsed: Value = serde_json::from_slice(body_raw).map_err(LiveZoneError::BodyNotJson)?; let messages = parsed @@ -604,6 +631,7 @@ pub fn compress_anthropic_live_zone( block_type, tokenizer.as_ref(), &mut replacements, + ccr_store, ); outcome } @@ -621,6 +649,7 @@ pub fn compress_anthropic_live_zone( "string_content".to_string(), tokenizer.as_ref(), &mut replacements, + ccr_store, ) } }; @@ -694,6 +723,7 @@ fn compress_one_block( block_type: String, tokenizer: &dyn crate::tokenizer::Tokenizer, replacements: &mut Vec, + ccr_store: Option<&dyn CcrStore>, ) -> BlockOutcome { // 1. Byte-threshold gate. Empty content always falls through to // `dispatch_compressor` (which short-circuits on empty), so @@ -726,7 +756,22 @@ fn compress_one_block( compressed, } => { let original_bytes = content_text.len(); - let compressed_bytes = compressed.len(); + // PR-B7: when a CCR store is wired, persist the original + // block content keyed by `BLAKE3(original)[..24]` and append + // the `<>` marker to the compressed string. The + // marker stays on a fresh trailing line so it is easy for + // the model to spot and so that the per-content-type + // compressors (which already produce trailing summary + // lines) keep their final newline before the marker. + // + // The token-validation gate (step 3) is computed against + // the marker-augmented string so the saved-token check + // stays honest — the marker costs ~6 tokens and we'd + // rather forward the original than ship a bigger payload + // for a 5-byte block. + let (compressed_for_replacement, ccr_hash_emitted) = + maybe_inject_ccr_marker(content_text, &compressed, ccr_store); + let compressed_bytes = compressed_for_replacement.len(); // 3. Tokenizer-validated rejection. Per PR-B4 spec we // count both the original and compressed strings // using the model's tokenizer; the compression is @@ -735,7 +780,7 @@ fn compress_one_block( // pathological inputs (e.g. dense base64 → tokenizer // fragments more aggressively after a transform). let original_tokens = tokenizer.count_text(content_text); - let compressed_tokens = tokenizer.count_text(&compressed); + let compressed_tokens = tokenizer.count_text(&compressed_for_replacement); if compressed_tokens >= original_tokens { BlockOutcome { message_index, @@ -750,8 +795,16 @@ fn compress_one_block( }, } } else { - let replacement_bytes = - serde_json::to_vec(&compressed).expect("string is always JSON-encodable"); + // Only persist to the CCR store once the rejection + // gate has admitted the compression — otherwise we + // populate the store with hashes whose markers + // never reach the wire (still correct, but wastes + // storage capacity). + if let (Some(store), Some(hash)) = (ccr_store, ccr_hash_emitted.as_deref()) { + store.put(hash, content_text); + } + let replacement_bytes = serde_json::to_vec(&compressed_for_replacement) + .expect("string is always JSON-encodable"); replacements.push(Replacement { range: content_byte_range, replacement: replacement_bytes, @@ -1078,6 +1131,39 @@ fn apply_replacements(original: &[u8], replacements: &mut [Replacement]) -> Vec< out } +/// PR-B7: append a `<>` retrieval marker to the compressed +/// block content when a CCR store is wired. Returns the +/// (possibly-augmented) compressed string and the hash that was +/// emitted (so the caller can decide whether to put the original into +/// the store after the rejection gate). When `ccr_store` is `None`, +/// returns the input compressed string unchanged with `None`. +/// +/// The marker is appended on its own line — `\n<>` — so: +/// +/// 1. The marker is unambiguously after the compressor's last byte, +/// even if that byte was a newline already (we only add one). +/// 2. Markers are easy to detect in human-readable diffs / logs. +/// 3. The Python `inject_ccr_retrieve_tool` regex in +/// `headroom/ccr/tool_injection.py` keeps working — it matches +/// `[a-f0-9]{24}` anywhere in the text. +fn maybe_inject_ccr_marker( + original: &str, + compressed: &str, + ccr_store: Option<&dyn CcrStore>, +) -> (String, Option) { + if ccr_store.is_none() { + return (compressed.to_string(), None); + } + let hash = compute_key(original.as_bytes()); + let marker = marker_for(&hash); + let augmented = if compressed.ends_with('\n') { + format!("{compressed}{marker}") + } else { + format!("{compressed}\n{marker}") + }; + (augmented, Some(hash)) +} + /// Per-block dispatch result — whether any compressor ran and what /// it produced. enum DispatchResult { diff --git a/crates/headroom-core/tests/ccr_backends.rs b/crates/headroom-core/tests/ccr_backends.rs new file mode 100644 index 000000000..a3ad8f338 --- /dev/null +++ b/crates/headroom-core/tests/ccr_backends.rs @@ -0,0 +1,199 @@ +//! Integration tests for the persistent CCR backends (PR-B7). +//! +//! Covers SQLite round-trip + TTL purge + restart-survival, the cross- +//! backend byte-equal-key invariant, and (cfg-gated) the Redis backend. + +use std::time::Duration; + +use headroom_core::ccr::backends::{ + from_config, CcrBackendConfig, InMemoryCcrStore, SqliteCcrStore, +}; +use headroom_core::ccr::{compute_key, CcrStore}; + +#[test] +fn sqlite_round_trip() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("ccr.sqlite"); + let store = SqliteCcrStore::open(&path, 300).expect("open sqlite store"); + let payload = r#"[{"id":1},{"id":2},{"id":3}]"#; + let hash = compute_key(payload.as_bytes()); + store.put(&hash, payload); + let fetched = store.get(&hash); + assert_eq!(fetched.as_deref(), Some(payload)); + assert_eq!(store.len(), 1); + // Missing key returns None. + assert_eq!(store.get("missing-hash-key"), None); +} + +#[test] +fn sqlite_ttl_purge() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("ccr.sqlite"); + // 0-second TTL forces every entry to be expired the moment we read it. + let store = SqliteCcrStore::open(&path, 0).expect("open sqlite store"); + let hash = compute_key(b"to be purged"); + store.put(&hash, "to be purged"); + // Sleep long enough for `created_at + ttl_seconds <= now()` (1s clock + // resolution on unix-seconds). + std::thread::sleep(Duration::from_millis(1_100)); + assert_eq!(store.get(&hash), None, "expired entry must be purged"); + assert_eq!(store.len(), 0, "expired entry must be physically deleted"); +} + +#[test] +fn sqlite_persists_across_proxy_restart() { + // Acceptance criterion #4 from the plan: write via SqliteCcrStore, + // drop the store, reconstruct from the same DB path, retrieve same + // hash → original bytes recover. + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("ccr.sqlite"); + let payload = "long-lived original payload"; + let hash = compute_key(payload.as_bytes()); + + { + let store = SqliteCcrStore::open(&path, 300).expect("open sqlite store (turn 1)"); + store.put(&hash, payload); + // `store` drops here, simulating worker shutdown. + } + + // Reconstruct from the same path — simulates `--workers 1` restart. + let store = SqliteCcrStore::open(&path, 300).expect("re-open sqlite store (turn 2)"); + let fetched = store.get(&hash); + assert_eq!( + fetched.as_deref(), + Some(payload), + "re-opened sqlite store must recover the original bytes" + ); +} + +#[test] +fn from_config_sqlite_roundtrip() { + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("ccr.sqlite"); + let cfg = CcrBackendConfig::Sqlite { + path: path.clone(), + ttl_seconds: 300, + }; + let store = from_config(&cfg).expect("from_config(sqlite)"); + let hash = compute_key(b"hello"); + store.put(&hash, "hello"); + assert_eq!(store.get(&hash).as_deref(), Some("hello")); +} + +#[test] +fn from_config_in_memory_roundtrip() { + let cfg = CcrBackendConfig::in_memory_default(); + let store = from_config(&cfg).expect("from_config(in_memory)"); + let hash = compute_key(b"bye"); + store.put(&hash, "bye"); + assert_eq!(store.get(&hash).as_deref(), Some("bye")); +} + +#[cfg(not(feature = "redis"))] +#[test] +fn from_config_redis_unsupported_when_feature_off() { + use headroom_core::ccr::backends::CcrBackendInitError; + + let cfg = CcrBackendConfig::Redis { + url: "redis://127.0.0.1:6379".to_string(), + ttl_seconds: 300, + key_prefix: None, + }; + match from_config(&cfg) { + Err(CcrBackendInitError::UnsupportedBackend { backend, feature }) => { + assert_eq!(backend, "redis"); + assert_eq!(feature, "redis"); + } + Err(other) => panic!("expected UnsupportedBackend, got {other:?}"), + Ok(_) => panic!("redis must error when feature is off"), + } +} + +#[test] +fn backend_swap_byte_equal_keys() { + // Stage data through one backend, swap to another with the same + // payload, and assert the keys are byte-equal. This is the + // load-bearing invariant: operators may migrate between backends + // (e.g. SQLite → Redis when scaling out) and the in-flight CCR + // markers must keep working — the marker bytes are the hash, and + // the hash function is fixed in `ccr::compute_key`. + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("ccr.sqlite"); + + let sqlite = SqliteCcrStore::open(&path, 300).expect("open sqlite store"); + let in_memory = InMemoryCcrStore::new(); + + let payloads = [ + "alpha", + r#"[{"id":1}]"#, + "the quick brown fox jumps over the lazy dog", + "<<<<>>>>", // marker-adjacent characters — sanity check on the BLAKE3 trim + ]; + + for payload in &payloads { + let key_a = compute_key(payload.as_bytes()); + let key_b = compute_key(payload.as_bytes()); + // Step 1: same payload yields byte-equal keys. + assert_eq!(key_a, key_b, "compute_key must be deterministic"); + + // Step 2: store in sqlite, mirror to in-memory under the same + // key — both backends recover byte-equal values. + sqlite.put(&key_a, payload); + in_memory.put(&key_b, payload); + + let v_sqlite = sqlite.get(&key_a); + let v_mem = in_memory.get(&key_b); + assert_eq!(v_sqlite.as_deref(), Some(*payload)); + assert_eq!(v_mem.as_deref(), Some(*payload)); + assert_eq!( + v_sqlite, v_mem, + "sqlite and in-memory must return byte-equal payloads" + ); + } +} + +// ─── Redis-feature-gated tests ───────────────────────────────────────── + +#[cfg(feature = "redis")] +mod redis_tests { + use super::*; + use headroom_core::ccr::backends::RedisCcrStore; + + /// Reads `HEADROOM_TEST_REDIS_URL` from the environment — when the + /// feature is on but no URL is configured we silently no-op. CI + /// runs the redis test in a docker-compose'd matrix. + fn redis_url() -> Option { + std::env::var("HEADROOM_TEST_REDIS_URL").ok() + } + + #[test] + fn redis_round_trip() { + let Some(url) = redis_url() else { + eprintln!("skipping redis_round_trip: HEADROOM_TEST_REDIS_URL not set"); + return; + }; + let store = RedisCcrStore::open(&url, 300).expect("open redis store"); + let payload = "redis payload"; + let hash = compute_key(payload.as_bytes()); + store.put(&hash, payload); + assert_eq!(store.get(&hash).as_deref(), Some(payload)); + } + + #[test] + fn redis_round_trip_via_from_config() { + let Some(url) = redis_url() else { + eprintln!("skipping redis_round_trip_via_from_config: HEADROOM_TEST_REDIS_URL not set"); + return; + }; + let cfg = CcrBackendConfig::Redis { + url, + ttl_seconds: 300, + key_prefix: Some("ccr_test".to_string()), + }; + let store = from_config(&cfg).expect("from_config(redis)"); + let payload = "via factory"; + let hash = compute_key(payload.as_bytes()); + store.put(&hash, payload); + assert_eq!(store.get(&hash).as_deref(), Some(payload)); + } +} diff --git a/crates/headroom-core/tests/live_zone_ccr.rs b/crates/headroom-core/tests/live_zone_ccr.rs new file mode 100644 index 000000000..7f973d2d2 --- /dev/null +++ b/crates/headroom-core/tests/live_zone_ccr.rs @@ -0,0 +1,119 @@ +//! PR-B7 — `compress_anthropic_live_zone_with_ccr` integration tests. +//! +//! Confirms that wiring a CCR store into the live-zone dispatcher: +//! 1. Stores the original block bytes keyed by `BLAKE3(original)[..24]`. +//! 2. Appends `<>` to the compressed block content. +//! 3. Leaves bytes outside the live zone byte-identical (cache safety). +//! 4. Is byte-equivalent to PR-B4 behaviour when no CCR store is wired. + +use headroom_core::ccr::backends::InMemoryCcrStore; +use headroom_core::ccr::{compute_key, CcrStore}; +use headroom_core::transforms::live_zone::{ + compress_anthropic_live_zone, compress_anthropic_live_zone_with_ccr, AuthMode, LiveZoneOutcome, + DEFAULT_MODEL, +}; +use serde_json::{json, Value}; + +/// Build a synthetic JSON-array tool_result above the 1 KiB threshold +/// so SmartCrusher actually engages. +fn large_json_array_payload() -> String { + let items: Vec = (0..40) + .map(|i| { + json!({ + "id": i, + "name": format!("entry_{i}"), + "score": i * 7, + "active": i % 2 == 0, + "notes": "lorem ipsum dolor sit amet, consectetur adipiscing elit", + }) + }) + .collect(); + serde_json::to_string(&Value::Array(items)).unwrap() +} + +fn body_with_payload(payload: &str) -> Vec { + serde_json::to_vec(&json!({ + "model": "claude-3-5-sonnet-20241022", + "messages": [ + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "t1", "content": payload} + ] + } + ] + })) + .unwrap() +} + +#[test] +fn ccr_marker_injected_when_store_wired() { + let payload = large_json_array_payload(); + let body = body_with_payload(&payload); + let store = InMemoryCcrStore::new(); + + let outcome = compress_anthropic_live_zone_with_ccr( + &body, + 0, + AuthMode::Payg, + DEFAULT_MODEL, + Some(&store), + ) + .expect("dispatcher must succeed"); + + let new_body = match &outcome { + LiveZoneOutcome::Modified { new_body, .. } => new_body.get().to_string(), + LiveZoneOutcome::NoChange { .. } => { + panic!("expected Modified; SmartCrusher should compress this payload") + } + }; + + let expected_hash = compute_key(payload.as_bytes()); + let marker = format!("<>"); + assert!( + new_body.contains(&marker), + "compressed body must contain CCR marker; body={new_body}" + ); + let recovered = store.get(&expected_hash); + assert_eq!( + recovered.as_deref(), + Some(payload.as_str()), + "store must hold the original bytes under the BLAKE3 hash key" + ); +} + +#[test] +fn no_marker_when_store_omitted() { + let payload = large_json_array_payload(); + let body = body_with_payload(&payload); + + let outcome = + compress_anthropic_live_zone(&body, 0, AuthMode::Payg, DEFAULT_MODEL).expect("dispatcher"); + + let new_body = match &outcome { + LiveZoneOutcome::Modified { new_body, .. } => new_body.get().to_string(), + LiveZoneOutcome::NoChange { .. } => return, // legitimate — token gate may reject + }; + + assert!( + !new_body.contains("< tuple[list[dict[str, Any]], bool]: """Inject CCR retrieval tool into tools list. + PR-B7 (`REALIGNMENT/04-phase-B-live-zone.md`): callers may pass + ``session_has_done_ccr=True`` so the tool is injected even when + THIS request has no fresh compression markers. That is the + sticky-on path: once a session has done CCR, the + ``headroom_retrieve`` tool must stay in ``body["tools"]`` for + every subsequent request, otherwise the tool list bytes flip + on/off mid-session and bust the prompt cache. + + Most callers should prefer + :func:`headroom.proxy.helpers.apply_session_sticky_ccr_tool` + which threads the ``SessionCcrTracker`` directly. This method + is the per-request fallback used when no session_id is available + (e.g. Google handler, legacy code paths). + Args: tools: Existing tools list (may be None or empty). + session_has_done_ccr: When True, inject regardless of + whether the current request contained compression + markers. Default False preserves legacy per-request + behaviour. Returns: Tuple of (updated_tools, was_injected). was_injected is False if tool was already present (e.g., from MCP). """ - if not self.inject_tool or not self.has_compressed_content: + if not self.inject_tool: + return tools or [], False + # PR-B7: sticky-on takes precedence. If the session has + # previously done CCR, register the tool even when this turn + # has no fresh markers. Otherwise fall back to the per-request + # check for backwards compat. + if not (session_has_done_ccr or self.has_compressed_content): return tools or [], False tools = tools or [] @@ -390,6 +416,8 @@ class CCRToolInjector: self, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, + *, + session_has_done_ccr: bool = False, ) -> tuple[list[dict[str, Any]], list[dict[str, Any]] | None, bool]: """Process a request, scanning for markers and injecting as needed. @@ -398,9 +426,17 @@ class CCRToolInjector: 2. Inject tool definition if enabled (skipped if already present from MCP) 3. Inject system instructions if enabled + PR-B7: when ``session_has_done_ccr`` is True the tool gets + injected even when the current message stream has no fresh + markers. System-instruction injection still keys off + per-request markers (the system prompt is the cache hot zone — + we never mutate it without a current-turn reason). + Args: messages: Request messages. tools: Request tools (may be None). + session_has_done_ccr: PR-B7 sticky-on flag — when True, + register the tool regardless of this turn's marker scan. Returns: Tuple of (updated_messages, updated_tools, tool_was_injected). @@ -408,10 +444,12 @@ class CCRToolInjector: """ self.scan_for_markers(messages) - if not self.has_compressed_content: + if not (self.has_compressed_content or session_has_done_ccr): return messages, tools, False - updated_tools, was_injected = self.inject_tool_definition(tools) + updated_tools, was_injected = self.inject_tool_definition( + tools, session_has_done_ccr=session_has_done_ccr + ) updated_messages = self.inject_into_system_message(messages) return updated_messages, updated_tools if updated_tools else None, was_injected diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 75efc684e..7633200f9 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -1055,6 +1055,15 @@ class AnthropicHandlerMixin: logger.debug(f"[{request_id}] post_compress hook error: {e}") # CCR Tool Injection: Inject retrieval tool if compression occurred + # OR if this session has previously done CCR (PR-B7 sticky-on). + # The legacy `CCRToolInjector` flips on/off based on the *current* + # request's compressed-content presence, busting cache every flip. + # We now route the tool-list update through + # `apply_session_sticky_ccr_tool`, which once-on/always-on per + # `SessionCcrTracker`. System-instruction injection keeps its + # existing per-request scan (it lives in the system prompt, which + # is the cache hot zone — gated separately by the + # `frozen_message_count > 0` guard below). tools = body.get("tools") _original_tools = tools # Preserve for diagnostic / future retry if ( @@ -1074,26 +1083,38 @@ class AnthropicHandlerMixin: f"(frozen prefix={frozen_message_count}) to preserve cache" ) inject_tool = False - # Create fresh injector to avoid state leakage between requests + # Scan for compression markers + maybe inject system instructions. + # Tool-list injection is handled separately via the sticky helper. injector = CCRToolInjector( provider="anthropic", - inject_tool=inject_tool, + inject_tool=False, # routed through sticky helper below inject_system_instructions=inject_system_instructions, ) - optimized_messages, tools, was_injected = injector.process_request( - optimized_messages, tools - ) + injector.scan_for_markers(optimized_messages) + if inject_system_instructions and injector.has_compressed_content: + optimized_messages = injector.inject_into_system_message(optimized_messages) + + # Sticky-on tool registration (PR-B7): always inject the + # retrieval tool once a session has done CCR, regardless + # of whether THIS turn produced compressed content. + if inject_tool: + from headroom.proxy.helpers import apply_session_sticky_ccr_tool + + tools, ccr_tool_injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id=request_id, + existing_tools=tools, + has_compressed_content_this_turn=injector.has_compressed_content, + ) + if ccr_tool_injected: + logger.debug( + f"[{request_id}] CCR: tool registered (session={session_id}, " + f"compressed_this_turn={injector.has_compressed_content}, " + f"hashes_seen={len(injector.detected_hashes)})" + ) if injector.has_compressed_content: - if was_injected: - logger.debug( - f"[{request_id}] CCR: Injected retrieval tool for hashes: {injector.detected_hashes}" - ) - else: - logger.debug( - f"[{request_id}] CCR: Tool already present (MCP?), skipped injection for hashes: {injector.detected_hashes}" - ) - # Track compression in context tracker for multi-turn awareness if self.ccr_context_tracker: self._turn_counter += 1 diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 560d64c02..2289e7899 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -543,6 +543,12 @@ class OpenAIHandlerMixin: logger.debug(f"[{request_id}] post_compress hook error: {e}") # CCR Tool Injection: Inject retrieval tool if compression occurred + # OR if this session has previously done CCR (PR-B7 sticky-on). + # See `headroom/proxy/handlers/anthropic.py` and PR-B7 plan + # `REALIGNMENT/04-phase-B-live-zone.md` for the rationale: once a + # session has done CCR, the `headroom_retrieve` tool stays + # registered for every subsequent turn so the prompt cache + # anchored on the previous turn's tool list never busts. tools = body.get("tools") _original_tools = tools # Preserve for diagnostic / future retry if ( @@ -550,21 +556,28 @@ class OpenAIHandlerMixin: ) and not _bypass: injector = CCRToolInjector( provider="openai", - inject_tool=self.config.ccr_inject_tool, + inject_tool=False, # routed through sticky helper below inject_system_instructions=self.config.ccr_inject_system_instructions, ) - optimized_messages, tools, was_injected = injector.process_request( - optimized_messages, tools - ) + injector.scan_for_markers(optimized_messages) + if self.config.ccr_inject_system_instructions and injector.has_compressed_content: + optimized_messages = injector.inject_into_system_message(optimized_messages) - if injector.has_compressed_content: - if was_injected: + if self.config.ccr_inject_tool: + from headroom.proxy.helpers import apply_session_sticky_ccr_tool + + tools, ccr_tool_injected = apply_session_sticky_ccr_tool( + provider="openai", + session_id=openai_session_id, + request_id=request_id, + existing_tools=tools, + has_compressed_content_this_turn=injector.has_compressed_content, + ) + if ccr_tool_injected: logger.debug( - f"[{request_id}] CCR: Injected retrieval tool for hashes: {injector.detected_hashes}" - ) - else: - logger.debug( - f"[{request_id}] CCR: Tool already present (MCP?), skipped injection for hashes: {injector.detected_hashes}" + f"[{request_id}] CCR: tool registered (session={openai_session_id}, " + f"compressed_this_turn={injector.has_compressed_content}, " + f"hashes_seen={len(injector.detected_hashes)})" ) if is_cache_mode(self.config.mode): diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index b7731fca1..8cfdf9e8f 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -1591,6 +1591,286 @@ def apply_session_sticky_memory_tools( return tools_out, added_bytes > 0 +# ─── Session-sticky CCR tool injection (PR-B7) ───────────────────────── +# +# Per realignment plan PR-B7 (`REALIGNMENT/04-phase-B-live-zone.md`): +# once a session has performed any CCR compression, the +# `headroom_retrieve` tool stays registered in `body["tools"]` for every +# subsequent request in that session — never toggled off. +# +# The legacy `CCRToolInjector.has_compressed_content` flips on/off based +# on whether the *latest request* contained compression markers, which +# bust the prompt cache every time the flag flips. Sticky-on means the +# tool list bytes stay byte-stable across turns once injected. + + +class SessionCcrTracker: + """Bounded LRU tracker recording per-(provider, session_id) CCR state. + + Two pieces of state per session: + + * ``has_done_ccr``: True once the proxy observed any CCR + compression marker in the messages of a request. Once True, it + never flips back to False (the prompt cache anchored on the + previous turn's tool list demands the tool stays present). + * ``golden_tool_bytes``: canonical serialization of the + ``headroom_retrieve`` tool definition recorded the first time + the tracker injected it. Subsequent turns replay these bytes + verbatim. + + Bounded by ``max_sessions`` via ``OrderedDict`` LRU. Mirrors + :class:`SessionToolTracker` semantics so the operator's mental model + is one tracker pattern, not two. + """ + + def __init__(self, max_sessions: int | None = None) -> None: + if max_sessions is None: + max_sessions = get_tool_tracker_max_sessions() + if max_sessions <= 0: + raise ValueError("max_sessions must be > 0") + self._max_sessions = max_sessions + self._lock = threading.RLock() + # Value is (has_done_ccr, golden_tool_bytes_or_none). + self._sessions: OrderedDict[tuple[str, str], tuple[bool, bytes | None]] = OrderedDict() + + @property + def active_sessions(self) -> int: + with self._lock: + return len(self._sessions) + + def _key(self, provider: str, session_id: str) -> tuple[str, str]: + return (provider, session_id) + + def has_done_ccr(self, provider: str, session_id: str) -> bool: + """Return True iff this session has previously performed CCR.""" + if not provider: + raise ValueError("provider must be non-empty") + if not session_id: + raise ValueError("session_id must be non-empty") + with self._lock: + entry = self._sessions.get(self._key(provider, session_id)) + if entry is None: + return False + self._sessions.move_to_end(self._key(provider, session_id)) + return entry[0] + + def get_golden_tool_bytes(self, provider: str, session_id: str) -> bytes | None: + """Return the recorded golden tool-definition bytes, or None.""" + if not provider: + raise ValueError("provider must be non-empty") + if not session_id: + raise ValueError("session_id must be non-empty") + with self._lock: + entry = self._sessions.get(self._key(provider, session_id)) + if entry is None: + return None + self._sessions.move_to_end(self._key(provider, session_id)) + return entry[1] + + def record_ccr_done( + self, + provider: str, + session_id: str, + golden_tool_bytes: bytes, + ) -> None: + """Mark the session as having performed CCR and pin the golden bytes. + + First-write wins for ``golden_tool_bytes`` (subsequent calls + with the same session keep the original bytes — prevents drift + if the canonical serialization changed mid-session). The + ``has_done_ccr`` flag is monotonic: once True, never False. + """ + if not provider: + raise ValueError("provider must be non-empty") + if not session_id: + raise ValueError("session_id must be non-empty") + if not golden_tool_bytes: + raise ValueError("golden_tool_bytes must be non-empty") + key = self._key(provider, session_id) + with self._lock: + existing = self._sessions.get(key) + if existing is None: + self._sessions[key] = (True, golden_tool_bytes) + else: + # Preserve original golden bytes; just promote the flag. + pinned = existing[1] if existing[1] is not None else golden_tool_bytes + self._sessions[key] = (True, pinned) + self._sessions.move_to_end(key) + while len(self._sessions) > self._max_sessions: + self._sessions.popitem(last=False) + + def reset(self) -> None: + """Clear all session state (test helper).""" + with self._lock: + self._sessions.clear() + + +# Process-wide singleton. +_session_ccr_tracker_lock = threading.Lock() +_session_ccr_tracker: SessionCcrTracker | None = None + + +def get_session_ccr_tracker() -> SessionCcrTracker: + """Return the process-wide :class:`SessionCcrTracker` singleton.""" + global _session_ccr_tracker + with _session_ccr_tracker_lock: + if _session_ccr_tracker is None: + _session_ccr_tracker = SessionCcrTracker() + return _session_ccr_tracker + + +def _reset_session_ccr_tracker_for_test() -> None: + """Clear the process-wide CCR tracker (test-only).""" + global _session_ccr_tracker + with _session_ccr_tracker_lock: + _session_ccr_tracker = None + + +def apply_session_sticky_ccr_tool( + *, + provider: Literal["anthropic", "openai", "google"], + session_id: str | None, + request_id: str | None, + existing_tools: list[dict[str, Any]] | None, + has_compressed_content_this_turn: bool, +) -> tuple[list[dict[str, Any]], bool]: + """Apply sticky-on CCR retrieval-tool injection per :class:`SessionCcrTracker`. + + Coordination point for both Anthropic and OpenAI handlers — replaces + the legacy ``CCRToolInjector.inject_tool_definition`` "flip on, flip + off" behaviour. + + Logic: + + * If ``session_id`` is None: tracker is bypassed and the per-turn + ``has_compressed_content_this_turn`` flag drives the decision + verbatim (matching legacy behaviour for WS / pre-session paths). + * If the session has previously done CCR (``has_done_ccr``): + ALWAYS inject the recorded golden bytes — even if this turn has + no fresh compression. That is the load-bearing PR-B7 fix. + * Otherwise, inject only when this turn produced compressed content. + The first injection records the golden bytes for future turns. + + Tools whose name already equals ``CCR_TOOL_NAME`` (e.g. the client + pre-registered it via MCP) are not re-appended; the client's bytes + win. + + Returns ``(updated_tools, was_injected)``. ``updated_tools`` is a + fresh list (caller-safe). + """ + from headroom.ccr.tool_injection import CCR_TOOL_NAME, create_ccr_tool_definition + + if provider not in ("anthropic", "openai", "google"): + raise ValueError(f"unsupported provider: {provider!r}") + + tools_out: list[dict[str, Any]] = list(existing_tools) if existing_tools else [] + existing_names: set[str] = set() + for t in tools_out: + n = _extract_tool_name(t) + if n: + existing_names.add(n) + + # Client (or MCP) already provided a tool by this name — don't double up. + if CCR_TOOL_NAME in existing_names: + log_tool_injection_decision( + provider=provider, + session_id=session_id, + decision="skip", + tool_definition_bytes_count=0, + request_id=request_id, + ) + return tools_out, False + + # No session_id (e.g. WS path): per-turn decision drives directly. + if not session_id: + if not has_compressed_content_this_turn: + log_tool_injection_decision( + provider=provider, + session_id=None, + decision="skip", + tool_definition_bytes_count=0, + request_id=request_id, + ) + return tools_out, False + tool_def = create_ccr_tool_definition(provider) + canonical = serialize_tool_definition_canonical(tool_def) + tools_out.append(tool_def) + log_tool_injection_decision( + provider=provider, + session_id=None, + decision="inject_first_time", + tool_definition_bytes_count=len(canonical), + request_id=request_id, + ) + return tools_out, True + + tracker = get_session_ccr_tracker() + previously_done = tracker.has_done_ccr(provider, session_id) + + if previously_done: + # Sticky replay path. Always inject — even if this turn had no + # fresh CCR compression. Prefer the recorded golden bytes; fall + # back to a freshly serialized definition if (somehow) the + # tracker lost them. Loud per build constraint #4: we log the + # path taken either way. + golden = tracker.get_golden_tool_bytes(provider, session_id) + if golden is not None: + try: + tool_def = json.loads(golden.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as exc: + # Should never happen — golden bytes were produced by us. + raise RuntimeError( + f"corrupt golden CCR tool bytes for session {session_id}: {exc}" + ) from exc + tools_out.append(tool_def) + log_tool_injection_decision( + provider=provider, + session_id=session_id, + decision="inject_sticky_replay", + tool_definition_bytes_count=len(golden), + request_id=request_id, + ) + return tools_out, True + # Tracker says "done CCR" but somehow has no golden bytes. Pin + # them now so future turns are stable. + tool_def = create_ccr_tool_definition(provider) + canonical = serialize_tool_definition_canonical(tool_def) + tracker.record_ccr_done(provider, session_id, canonical) + tools_out.append(tool_def) + log_tool_injection_decision( + provider=provider, + session_id=session_id, + decision="inject_sticky_replay", + tool_definition_bytes_count=len(canonical), + request_id=request_id, + ) + return tools_out, True + + # Fresh session — only inject when this turn produced compressed content. + if not has_compressed_content_this_turn: + log_tool_injection_decision( + provider=provider, + session_id=session_id, + decision="skip", + tool_definition_bytes_count=0, + request_id=request_id, + ) + return tools_out, False + + tool_def = create_ccr_tool_definition(provider) + canonical = serialize_tool_definition_canonical(tool_def) + tracker.record_ccr_done(provider, session_id, canonical) + tools_out.append(tool_def) + log_tool_injection_decision( + provider=provider, + session_id=session_id, + decision="inject_first_time", + tool_definition_bytes_count=len(canonical), + request_id=request_id, + ) + return tools_out, True + + async def _read_request_body_bytes(request: Request) -> bytes: """Read and (if needed) decompress the request body, returning raw UTF-8 bytes. diff --git a/tests/test_ccr_tool_always_on.py b/tests/test_ccr_tool_always_on.py new file mode 100644 index 000000000..df83a5e2b --- /dev/null +++ b/tests/test_ccr_tool_always_on.py @@ -0,0 +1,385 @@ +"""PR-B7 — `headroom_retrieve` tool always-on once a session has done CCR. + +These tests pin three properties: + +1. After a session has performed CCR even once, every subsequent + request in that session injects the tool — even when the current + request has no fresh compression markers. +2. A session that has NEVER done CCR does not get the tool injected. +3. The tool definition bytes are byte-stable across turns (snapshot + test). Any future change to the tool schema must update the + snapshot deliberately. + +Tests target the canonical helper +`headroom.proxy.helpers.apply_session_sticky_ccr_tool` plus the +`SessionCcrTracker` semantics. The CCRToolInjector legacy path is +covered by `tests/test_ccr_tool_injection.py`. +""" + +from __future__ import annotations + +import pytest + +from headroom.ccr.tool_injection import ( + CCR_TOOL_NAME, + CCRToolInjector, + create_ccr_tool_definition, +) +from headroom.proxy.helpers import ( + SessionCcrTracker, + _reset_session_ccr_tracker_for_test, + apply_session_sticky_ccr_tool, + get_session_ccr_tracker, + serialize_tool_definition_canonical, +) + + +@pytest.fixture(autouse=True) +def _reset_tracker(): + _reset_session_ccr_tracker_for_test() + yield + _reset_session_ccr_tracker_for_test() + + +def _has_ccr_tool(tools: list[dict] | None) -> bool: + if not tools: + return False + for t in tools: + n = t.get("name") or t.get("function", {}).get("name") + if n == CCR_TOOL_NAME: + return True + return False + + +# ─── Sticky-on behavior ──────────────────────────────────────────────── + + +def test_tool_registered_on_every_request_after_first_ccr(): + """Once a session has done CCR, the tool stays registered every turn.""" + session_id = "sess-abc-123" + + # Turn 1: this turn produced compressed content → first-time inject. + tools_1, injected_1 = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="req-1", + existing_tools=None, + has_compressed_content_this_turn=True, + ) + assert injected_1 is True + assert _has_ccr_tool(tools_1) + + # Turn 2: NO fresh compression this turn — but session has done CCR. + # Tool MUST still be injected (PR-B7 sticky-on). + tools_2, injected_2 = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="req-2", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected_2 is True, "sticky replay must inject even with no fresh CCR" + assert _has_ccr_tool(tools_2) + + # Turn 3: still no fresh compression — sticky-on still fires. + tools_3, injected_3 = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="req-3", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected_3 is True + assert _has_ccr_tool(tools_3) + + +def test_tool_not_registered_if_session_never_did_ccr(): + """A session that never produced CCR markers gets no tool injection.""" + session_id = "fresh-session-no-ccr" + + tools, injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="req-1", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected is False + assert not _has_ccr_tool(tools) + # Tracker must NOT have recorded this session. + assert get_session_ccr_tracker().has_done_ccr("anthropic", session_id) is False + + # Do it again — same outcome, no state leakage. + tools, injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="req-2", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected is False + assert not _has_ccr_tool(tools) + + +def test_independent_sessions_track_independently(): + """Two distinct session_ids do not bleed sticky state.""" + sess_a = "sess-A" + sess_b = "sess-B" + + # A does CCR; B does not. + apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=sess_a, + request_id="r1", + existing_tools=None, + has_compressed_content_this_turn=True, + ) + + # B's next turn (no fresh CCR) must NOT auto-inject. + tools_b, injected_b = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=sess_b, + request_id="r2", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected_b is False + assert not _has_ccr_tool(tools_b) + + # A's next turn (no fresh CCR) MUST auto-inject (sticky). + tools_a, injected_a = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=sess_a, + request_id="r3", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected_a is True + assert _has_ccr_tool(tools_a) + + +def test_provider_isolation(): + """Same session_id under anthropic vs openai are tracked independently.""" + session_id = "shared-session-id" + + # Anthropic does CCR. + apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="r1", + existing_tools=None, + has_compressed_content_this_turn=True, + ) + + # OpenAI must NOT inherit anthropic's sticky state. + tools_o, injected_o = apply_session_sticky_ccr_tool( + provider="openai", + session_id=session_id, + request_id="r2", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected_o is False + assert not _has_ccr_tool(tools_o) + + +def test_existing_ccr_tool_in_client_list_skips_injection(): + """If client (e.g. via MCP) already provided headroom_retrieve, do not double up.""" + session_id = "sess-with-mcp" + client_tool = { + "name": CCR_TOOL_NAME, + "description": "client-provided", + "input_schema": {"type": "object", "properties": {}, "required": []}, + } + tools, injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="r1", + existing_tools=[client_tool], + has_compressed_content_this_turn=True, + ) + assert injected is False + # Tool list still contains the client's version (not duplicated). + names = [t.get("name") for t in tools] + assert names.count(CCR_TOOL_NAME) == 1 + + +def test_no_session_id_falls_back_to_per_turn_decision(): + """WS / pre-session paths with no session_id behave per-turn.""" + # No fresh CCR + no session_id → no inject. + tools, injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=None, + request_id="r1", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + assert injected is False + + # Fresh CCR + no session_id → inject (per-turn). + tools, injected = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=None, + request_id="r2", + existing_tools=None, + has_compressed_content_this_turn=True, + ) + assert injected is True + assert _has_ccr_tool(tools) + + +# ─── Byte-stable tool definition ─────────────────────────────────────── + + +# Snapshot of the canonical Anthropic CCR tool definition. Any change +# here MUST be deliberate — bumping the schema mid-session busts every +# active session's prompt cache (the tool list bytes are part of the +# cache key). +_ANTHROPIC_CCR_TOOL_SNAPSHOT_BYTES = ( + b'{"name":"headroom_retrieve",' + b'"description":"Retrieve original uncompressed content that was ' + b"compressed to save tokens. Use this when you need more data than " + b"what's shown in compressed tool results. The hash is provided in " + b'compression markers like [N items compressed... hash=abc123].",' + b'"input_schema":{"type":"object",' + b'"properties":{' + b'"hash":{"type":"string",' + b'"description":"Hash key from the compression marker ' + b"(e.g., 'abc123' from hash=abc123)\"}," + b'"query":{"type":"string",' + b'"description":"Optional search query to filter results. ' + b"If provided, only returns items matching the query. " + b'If omitted, returns all original items."}' + b"}," + b'"required":["hash"]}}' +) + +_OPENAI_CCR_TOOL_SNAPSHOT_BYTES = ( + b'{"type":"function",' + b'"function":{"name":"headroom_retrieve",' + b'"description":"Retrieve original uncompressed content that was ' + b"compressed to save tokens. Use this when you need more data than " + b"what's shown in compressed tool results. The hash is provided in " + b'compression markers like [N items compressed... hash=abc123].",' + b'"parameters":{"type":"object",' + b'"properties":{' + b'"hash":{"type":"string",' + b'"description":"Hash key from the compression marker ' + b"(e.g., 'abc123' from hash=abc123)\"}," + b'"query":{"type":"string",' + b'"description":"Optional search query to filter results. ' + b"If provided, only returns items matching the query. " + b'If omitted, returns all original items."}' + b"}," + b'"required":["hash"]}}}' +) + + +def test_tool_definition_byte_stable(): + """Pin the canonical bytes of the Anthropic + OpenAI tool defs. + + PR-B7 acceptance criterion: tool definition bytes are byte-stable. + Any future change to ``create_ccr_tool_definition`` must bump these + snapshots deliberately. + """ + anthropic_tool = create_ccr_tool_definition("anthropic") + canonical_anthropic = serialize_tool_definition_canonical(anthropic_tool) + assert canonical_anthropic == _ANTHROPIC_CCR_TOOL_SNAPSHOT_BYTES, ( + f"Anthropic CCR tool definition bytes changed.\n" + f" expected: {_ANTHROPIC_CCR_TOOL_SNAPSHOT_BYTES!r}\n" + f" actual: {canonical_anthropic!r}\n" + f"If this change is intentional, update the snapshot in this " + f"test and ensure the prompt-cache implications are reviewed." + ) + + openai_tool = create_ccr_tool_definition("openai") + canonical_openai = serialize_tool_definition_canonical(openai_tool) + assert canonical_openai == _OPENAI_CCR_TOOL_SNAPSHOT_BYTES, ( + f"OpenAI CCR tool definition bytes changed.\n" + f" expected: {_OPENAI_CCR_TOOL_SNAPSHOT_BYTES!r}\n" + f" actual: {canonical_openai!r}" + ) + + +def test_sticky_replay_returns_byte_equal_tool_each_turn(): + """The bytes injected on turn 2 must equal the bytes injected on turn 1.""" + session_id = "sess-byte-stable" + + tools_1, _ = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="r1", + existing_tools=None, + has_compressed_content_this_turn=True, + ) + tools_2, _ = apply_session_sticky_ccr_tool( + provider="anthropic", + session_id=session_id, + request_id="r2", + existing_tools=None, + has_compressed_content_this_turn=False, + ) + # Both tools lists should contain exactly one CCR tool with the + # same canonical bytes. + ccr_1 = next(t for t in tools_1 if t.get("name") == CCR_TOOL_NAME) + ccr_2 = next(t for t in tools_2 if t.get("name") == CCR_TOOL_NAME) + assert serialize_tool_definition_canonical(ccr_1) == serialize_tool_definition_canonical(ccr_2) + + +# ─── SessionCcrTracker unit coverage ──────────────────────────────────── + + +def test_session_ccr_tracker_monotonic_has_done_ccr(): + """``has_done_ccr`` is monotonic — never flips back to False.""" + tracker = SessionCcrTracker(max_sessions=10) + assert tracker.has_done_ccr("anthropic", "s1") is False + + golden = serialize_tool_definition_canonical(create_ccr_tool_definition("anthropic")) + tracker.record_ccr_done("anthropic", "s1", golden) + assert tracker.has_done_ccr("anthropic", "s1") is True + + # Re-record with a different golden_bytes — original bytes win + # (first-write wins) and flag stays True. + new_golden = b'{"name":"different","input_schema":{}}' + tracker.record_ccr_done("anthropic", "s1", new_golden) + assert tracker.has_done_ccr("anthropic", "s1") is True + assert tracker.get_golden_tool_bytes("anthropic", "s1") == golden + + +def test_session_ccr_tracker_lru_bound(): + """Tracker evicts oldest sessions once `max_sessions` is exceeded.""" + tracker = SessionCcrTracker(max_sessions=3) + golden = b"{}" + for i in range(5): + tracker.record_ccr_done("anthropic", f"s{i}", golden) + # Only 3 most recent should remain. + assert tracker.active_sessions == 3 + assert tracker.has_done_ccr("anthropic", "s0") is False + assert tracker.has_done_ccr("anthropic", "s1") is False + assert tracker.has_done_ccr("anthropic", "s4") is True + + +def test_session_ccr_tracker_reset_clears_state(): + tracker = SessionCcrTracker(max_sessions=10) + tracker.record_ccr_done("anthropic", "s1", b"{}") + assert tracker.active_sessions == 1 + tracker.reset() + assert tracker.active_sessions == 0 + assert tracker.has_done_ccr("anthropic", "s1") is False + + +# ─── Per-request injector legacy path (PR-B7 backwards compat) ───────── + + +def test_ccrtoolinjector_session_has_done_ccr_kwarg(): + """``CCRToolInjector.inject_tool_definition`` accepts session_has_done_ccr.""" + injector = CCRToolInjector(provider="anthropic", inject_tool=True) + # No fresh markers, no sticky flag → no inject (legacy behaviour). + tools, was = injector.inject_tool_definition(None) + assert was is False + assert tools == [] + + # No fresh markers, sticky flag set → inject (PR-B7 path). + tools, was = injector.inject_tool_definition(None, session_has_done_ccr=True) + assert was is True + assert _has_ccr_tool(tools)