diff --git a/crates/headroom-core/Cargo.toml b/crates/headroom-core/Cargo.toml index 3deb63a58..2bf5034ee 100644 --- a/crates/headroom-core/Cargo.toml +++ b/crates/headroom-core/Cargo.toml @@ -203,6 +203,10 @@ tempfile = "3" name = "tokenizer" harness = false +[[bench]] +name = "live_zone_dispatch" +harness = false + [[bench]] name = "ccr_store" harness = false diff --git a/crates/headroom-core/benches/live_zone_dispatch.rs b/crates/headroom-core/benches/live_zone_dispatch.rs new file mode 100644 index 000000000..3afbeb537 --- /dev/null +++ b/crates/headroom-core/benches/live_zone_dispatch.rs @@ -0,0 +1,248 @@ +//! Latency benchmark for the PR-B4 live-zone dispatch arms. +//! +//! Measures `compress_anthropic_live_zone` end to end — body parse, +//! live-zone walk, content-type detection, byte-threshold gate, and the +//! per-arm compressor — over representative payload shapes: +//! +//! - `source_code_2kb` / `source_code_20kb`: generated Python routed to +//! the CodeAwareCompressor arm (the tree-sitter parse dominates). +//! - `plain_text_8kb`: prose routed to the Kompress arm. What this case +//! actually measures depends on the machine's Hugging Face cache: +//! with the Kompress model cache-resident it benchmarks a real ONNX +//! inference; on a cold cache `Kompress::from_cache` resolves to a +//! deterministic NoOp and the number is detection + gate cost only. +//! Both are legitimate measurements — record which state applied +//! alongside any published number. +//! - `below_threshold_prose`: a payload under every byte threshold — +//! the cost of the gate itself, i.e. the overhead every small block +//! pays whether or not compression ever fires. +//! +//! Deliberately NOT wired into CI — benchmarks in shared CI runners +//! generate noise, not signal. Run manually: +//! +//! ```text +//! cargo bench -p headroom-core --bench live_zone_dispatch +//! ``` +//! +//! # Windows: `ORT_DYLIB_PATH` +//! +//! On Windows with a warm model cache, ONNX Runtime's shared-library +//! resolution can deadlock inside `ort` init unless `ORT_DYLIB_PATH` +//! points at the onnxruntime library (see +//! `docs/content/docs/troubleshooting.mdx`, "Windows ML DLL" entry). A +//! hung benchmark is strictly worse than a refused one, so on Windows +//! this harness exits early with instructions when the variable is +//! missing AND the Kompress model cache is warm — the only state where +//! the hang is reachable. Cache-cold it proceeds (no ONNX session is +//! ever created; the plain_text case then measures the deterministic +//! no-op path, a valid mode in its own right). Other platforms resolve +//! the library normally and are not gated. + +use std::hint::black_box; + +use criterion::{criterion_group, Criterion, Throughput}; +use headroom_core::transforms::live_zone::DEFAULT_MODEL; +use headroom_core::transforms::{compress_anthropic_live_zone, AuthMode}; +use serde_json::json; + +/// Build the standard single-`tool_result` Anthropic body around `text` +/// — the same shape the dispatch integration tests use (duplicated, not +/// shared: benches and integration tests are independent targets). +fn body_with_tool_result(text: &str) -> Vec { + serde_json::to_vec(&json!({ + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "system": "you are a helpful assistant", + "messages": [{ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "toolu_dispatch_bench", + "content": text, + }], + }], + })) + .expect("bench body serializes") +} + +/// Syntactically valid Python that detects as `SourceCode`, sized to at +/// least `target_bytes`. Same generator style as the dispatch tests' +/// `python_module_source`, parameterized by byte target instead of +/// function count so the two size cases are explicit at the call site. +fn python_source(target_bytes: usize) -> String { + let mut code = String::from( + "\"\"\"Example data-processing module used by the dispatch bench.\"\"\"\n\n\ + import json\nimport os\nfrom typing import Any, Optional\n\n\n", + ); + let mut i = 0usize; + while code.len() < target_bytes { + code.push_str(&format!( + "def process_record_{i}(record: dict) -> dict:\n \ + \"\"\"Normalize record {i} and compute its derived fields.\"\"\"\n \ + result = dict(record)\n \ + result[\"index\"] = {i}\n \ + result[\"doubled\"] = record.get(\"value\", 0) * 2\n \ + result[\"source\"] = \"batch\"\n \ + if result[\"doubled\"] > 100:\n \ + result[\"flag\"] = \"high\"\n \ + else:\n \ + result[\"flag\"] = \"low\"\n \ + return result\n\n\n" + )); + i += 1; + } + code +} + +/// Varied natural-language prose that detects as `PlainText`, sized to +/// at least `target_bytes`. Sentence shape varies so this is prose to +/// the detector, not a single repeated token run. +fn prose(target_bytes: usize) -> String { + const WORDS: &[&str] = &[ + "the", + "release", + "shipped", + "after", + "review", + "and", + "every", + "service", + "reported", + "healthy", + "metrics", + "while", + "operators", + "watched", + "dashboards", + "during", + "rollout", + "windows", + "before", + "traffic", + "returned", + "to", + "baseline", + "levels", + "overnight", + ]; + let mut text = String::with_capacity(target_bytes + 64); + let mut in_sentence = 0u32; + for (n, word) in WORDS.iter().cycle().enumerate() { + if text.len() >= target_bytes { + break; + } + text.push_str(word); + in_sentence += 1; + // Vary sentence length between 7 and 12 words. + if in_sentence >= 7 + (n as u32 % 6) { + text.push_str(".\n"); + in_sentence = 0; + } else { + text.push(' '); + } + } + text +} + +fn dispatch(body: &[u8]) { + let outcome = compress_anthropic_live_zone(body, 0, AuthMode::Payg, DEFAULT_MODEL) + .expect("dispatcher returns Ok on valid bodies"); + black_box(outcome); +} + +fn bench_dispatch(c: &mut Criterion) { + let cases: &[(&str, Vec)] = &[ + ( + "source_code_2kb", + body_with_tool_result(&python_source(2_200)), + ), + ( + "source_code_20kb", + body_with_tool_result(&python_source(20_000)), + ), + ("plain_text_8kb", body_with_tool_result(&prose(8_192))), + ("below_threshold_prose", body_with_tool_result(&prose(400))), + ]; + + // One warm-up dispatch per payload BEFORE any timed group, so + // process-latched one-time costs (compressor singletons, the + // Kompress `OnceLock` init and — cache-warm — its model load) land + // here rather than skewing the first timed case. + for (_, body) in cases { + dispatch(body); + } + + let mut group = c.benchmark_group("live_zone/dispatch"); + // The Kompress arm runs a real ONNX inference per iteration when + // the model is cache-resident; keep sampling time bounded. + group.sample_size(30); + for (name, body) in cases { + group.throughput(Throughput::Bytes(body.len() as u64)); + group.bench_function(*name, |b| b.iter(|| dispatch(black_box(body)))); + } + group.finish(); +} + +/// Refuse to run on Windows without `ORT_DYLIB_PATH` — but only when the +/// hang it guards against is actually reachable. The `ort`-init deadlock +/// needs a warm Kompress model cache: cache-cold, `Kompress::from_cache` +/// resolves `Ok(None)` before any ONNX session exists, no `ort` code +/// runs, and the cache-cold measurement mode the module doc advertises +/// is perfectly safe — refusing it would block a legitimate +/// configuration to guard against a hang it cannot have. A hung +/// benchmark process gives no diagnostic; this message does. +#[cfg(windows)] +fn check_ort_dylib_path() { + match std::env::var("ORT_DYLIB_PATH") { + Ok(v) if !v.trim().is_empty() => {} + _ if !kompress_model_cached() => { + eprintln!( + "live_zone_dispatch bench: ORT_DYLIB_PATH is not set, but the \ + Kompress model cache is cold — no ONNX session will be \ + created, so the Windows ort-init hang cannot occur. \ + Proceeding in cache-cold mode (the plain_text case measures \ + the deterministic no-op path)." + ); + } + _ => { + eprintln!( + "live_zone_dispatch bench: ORT_DYLIB_PATH is not set.\n\n\ + On Windows with a warm Kompress model cache, ONNX Runtime's\n\ + shared-library resolution can deadlock during `ort` init (see\n\ + docs/content/docs/troubleshooting.mdx), which would hang the\n\ + plain_text bench case indefinitely. Set ORT_DYLIB_PATH to the\n\ + onnxruntime shared library and re-run, e.g.:\n\n \ + ORT_DYLIB_PATH=C:\\path\\to\\onnxruntime.dll cargo bench \ + -p headroom-core --bench live_zone_dispatch\n\n\ + Refusing to start rather than risk a silent hang." + ); + std::process::exit(2); + } + } +} + +/// Whether the Kompress model is cache-resident, asked of the loader +/// itself. A hand-rolled probe would have to mirror +/// `Kompress::from_cache`'s root and artifact resolution and would go +/// stale silently; this cannot. +#[cfg(all(windows, feature = "ml"))] +fn kompress_model_cached() -> bool { + use headroom_core::transforms::kompress::{Kompress, KompressConfig}; + + matches!(Kompress::from_cache(KompressConfig::default()), Ok(Some(_))) +} + +/// Without the `ml` feature there is no ONNX session to deadlock on. +#[cfg(all(windows, not(feature = "ml")))] +fn kompress_model_cached() -> bool { + false +} + +criterion_group!(benches, bench_dispatch); + +fn main() { + #[cfg(windows)] + check_ort_dylib_path(); + benches(); + Criterion::default().configure_from_args().final_summary(); +} diff --git a/crates/headroom-core/src/transforms/content_detector.rs b/crates/headroom-core/src/transforms/content_detector.rs index 61aacd425..f3bef5ea3 100644 --- a/crates/headroom-core/src/transforms/content_detector.rs +++ b/crates/headroom-core/src/transforms/content_detector.rs @@ -41,6 +41,25 @@ pub enum ContentType { } impl ContentType { + /// Every variant, in declaration order. + /// + /// Rust has no stable reflection over enum variants, so callers that + /// need to enumerate the set (operator-facing name tables, exhaustive + /// round-trip tests) would otherwise hand-maintain their own copy and + /// silently miss a variant added later. Adding a variant without + /// extending this array is caught by the length annotation, and the + /// exhaustive `match` in [`ContentType::as_str`] forces the author + /// into this file in the first place. + pub const ALL: [ContentType; 7] = [ + ContentType::JsonArray, + ContentType::SourceCode, + ContentType::SearchResults, + ContentType::BuildOutput, + ContentType::GitDiff, + ContentType::Html, + ContentType::PlainText, + ]; + /// Stable string tag — matches Python's `ContentType..value`. pub fn as_str(&self) -> &'static str { match self { @@ -53,6 +72,46 @@ impl ContentType { ContentType::PlainText => "text", } } + + /// Operator-facing spelling of this variant: the name a human writes + /// in configuration. Equal to [`ContentType::as_str`] except where + /// that tag is abbreviated for Python parity (`search`, `build`, + /// `diff`, `text`), which are the spellings least likely to be + /// guessed correctly. Both forms parse — see the `FromStr` impl. + pub fn natural_name(&self) -> &'static str { + match self { + ContentType::JsonArray => "json_array", + ContentType::SourceCode => "source_code", + ContentType::SearchResults => "search_results", + ContentType::BuildOutput => "build_output", + ContentType::GitDiff => "git_diff", + ContentType::Html => "html", + ContentType::PlainText => "plain_text", + } + } +} + +/// Error returned when a string names no [`ContentType`]. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("unknown content type {0:?}")] +pub struct ParseContentTypeError(String); + +/// Parse a [`ContentType`] from either its [`as_str`](ContentType::as_str) +/// tag or its [`natural_name`](ContentType::natural_name). +/// +/// Accepting both matters for operator-facing configuration: the +/// abbreviated Python-parity tags (`text`, `search`, `build`, `diff`) are +/// not what a human writes, so a config naming `plain_text` would +/// otherwise parse as nothing and silently do nothing. +impl std::str::FromStr for ContentType { + type Err = ParseContentTypeError; + + fn from_str(name: &str) -> Result { + Self::ALL + .into_iter() + .find(|ct| ct.as_str() == name || ct.natural_name() == name) + .ok_or_else(|| ParseContentTypeError(name.to_owned())) + } } /// Result of `detect_content_type`. `metadata` is per-type free-form key/ diff --git a/crates/headroom-core/src/transforms/live_zone.rs b/crates/headroom-core/src/transforms/live_zone.rs index 6105dd346..ac51519d4 100644 --- a/crates/headroom-core/src/transforms/live_zone.rs +++ b/crates/headroom-core/src/transforms/live_zone.rs @@ -49,15 +49,20 @@ //! //! - **PR-B2** shipped the dispatcher *skeleton*: identify live-zone //! blocks, route to no-op compressors, always return `NoChange`. -//! - **PR-B3** (this PR) wires per-content-type compressors: +//! - **PR-B3** wired the first per-content-type compressors: //! `JsonArray` → SmartCrusher; `BuildOutput` → LogCompressor; -//! `SearchResults` → SearchCompressor; `GitDiff` → DiffCompressor; -//! `SourceCode` / `PlainText` / `Html` → no-op (B4 + a Rust -//! code-compressor port follow-up). -//! - **PR-B4** adds the tokenizer-validation gate (per-block +//! `SearchResults` → SearchCompressor; `GitDiff` → DiffCompressor — +//! leaving `SourceCode` / `PlainText` / `Html` as documented no-ops. +//! - **PR-B4** specifies the tokenizer-validation gate (per-block //! `compressed.tokens >= original.tokens` → fall back) and the -//! per-content-type byte threshold below which compression is -//! skipped. +//! per-content-type byte thresholds. The gate and the threshold +//! plumbing already shipped; **this PR** completes the item by wiring +//! the two remaining compressor arms — `SourceCode` → +//! CodeAwareCompressor (falling back to cached Kompress when the +//! compressor returns its input unchanged) and `PlainText` → cached +//! Kompress — and by moving the `SourceCode` / `PlainText` thresholds +//! off their placeholder 512 to the specified 2048 / 5120. `Html` +//! stays no-op. //! - **PR-B7** wires CCR retrieval-marker injection. //! //! # Cache safety invariant @@ -96,13 +101,19 @@ use std::{collections::HashSet, sync::OnceLock}; +#[cfg(feature = "ml")] +use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering}; + use serde::Deserialize; use serde_json::value::RawValue; use serde_json::Value; use thiserror::Error; +use super::code_compressor::{CodeAwareCompressor, CodeCompressorConfig}; use super::content_detector::{detect_content_type, ContentType}; use super::diff_compressor::{DiffCompressor, DiffCompressorConfig}; +#[cfg(feature = "ml")] +use super::kompress::{Kompress, KompressConfig}; use super::log_compressor::{LogCompressor, LogCompressorConfig}; use super::search_compressor::{SearchCompressor, SearchCompressorConfig}; use super::smart_crusher::{SmartCrusher, SmartCrusherConfig}; @@ -119,6 +130,14 @@ const STRATEGY_LOG_COMPRESSOR: &str = "log_compressor"; const STRATEGY_SEARCH_COMPRESSOR: &str = "search_compressor"; /// Strategy tag emitted when DiffCompressor rewrote a unified-diff block. const STRATEGY_DIFF_COMPRESSOR: &str = "diff_compressor"; +/// Strategy tag emitted when CodeAwareCompressor rewrote a source-code block. +const STRATEGY_CODE_COMPRESSOR: &str = "code_compressor"; +/// Strategy tag emitted when Kompress rewrote a plain-text block (also +/// used by the SourceCode arm's passthrough fallback — see +/// [`dispatch_compressor`]). Only the `ml` build has a Kompress arm to +/// emit it. +#[cfg(feature = "ml")] +const STRATEGY_KOMPRESS: &str = "kompress"; /// Empty query context passed to compressors that take a relevance /// query string. PR-B3 dispatcher does not yet plumb the user's last @@ -158,22 +177,33 @@ const THRESHOLD_BUILD_OUTPUT: usize = 512; const THRESHOLD_SEARCH_RESULTS: usize = 512; /// Git-diff blocks below this size route to no-op. const THRESHOLD_GIT_DIFF: usize = 512; -/// Source-code blocks below this size route to no-op. Pinned -/// for the future Rust code-compressor port — currently unused -/// because `ContentType::SourceCode` short-circuits to no-op above -/// the dispatch (see `dispatch_compressor`). -const THRESHOLD_SOURCE_CODE: usize = 512; -/// Plain-text blocks below this size route to no-op. Pinned -/// for the future Kompress wiring (PR-B7 follow-up); currently unused. -const THRESHOLD_PLAIN_TEXT: usize = 512; -/// HTML blocks have no compressor; threshold matches plain text so -/// when an HTML compressor lands the value is already pinned. +/// Source-code blocks below this size route to no-op. Sourced from +/// `REALIGNMENT/04-phase-B-live-zone.md::PR-B4` (2 KiB) — the +/// CodeAwareCompressor's own `min_tokens_for_compression` floor +/// already guards tiny snippets, but the byte gate keeps the dispatcher +/// from spinning up tree-sitter parsing for content that can't possibly +/// clear it. +const THRESHOLD_SOURCE_CODE: usize = 2048; +/// Plain-text blocks below this size route to no-op. Sourced from +/// `REALIGNMENT/04-phase-B-live-zone.md::PR-B4` (5 KiB). The threshold +/// gate fires before dispatch, so sub-threshold prose never reaches the +/// Kompress arm wired below. +const THRESHOLD_PLAIN_TEXT: usize = 5120; +/// HTML blocks have no compressor. 512 was PlainText's pre-PR-B4 value; +/// HTML deliberately keeps it rather than following PlainText to 5120, +/// because `REALIGNMENT/04` pins no HTML threshold — when an HTML +/// compressor lands, its threshold is that PR's own decision. const THRESHOLD_HTML: usize = 512; /// Map a content type to its byte threshold. Returning `usize` rather /// than an `Option` because every variant has a sensible default; /// `Html` is a no-op anyway so the threshold check never fires. -fn threshold_for(content_type: ContentType) -> usize { +/// +/// `pub` so tests and benches can ask "would this block have cleared the +/// gate?" against the real constants instead of mirroring them — a +/// mirrored copy drifts silently the moment a threshold moves. +#[must_use] +pub const fn threshold_for(content_type: ContentType) -> usize { match content_type { ContentType::JsonArray => THRESHOLD_JSON_ARRAY, ContentType::BuildOutput => THRESHOLD_BUILD_OUTPUT, @@ -549,6 +579,274 @@ fn diff_compressor() -> &'static DiffCompressor { INSTANCE.get_or_init(|| DiffCompressor::new(DiffCompressorConfig::default())) } +fn code_compressor() -> &'static CodeAwareCompressor { + static INSTANCE: OnceLock = OnceLock::new(); + INSTANCE.get_or_init(|| CodeAwareCompressor::new(CodeCompressorConfig::default())) +} + +// ─── Kompress model slot: never initialized on the request path ──────── +// +// `Kompress::from_cache` is cache-only — it never downloads — but it is +// NOT cheap: it loads the tokenizer and commits a ~261 MB ONNX session, +// whole seconds of wall time on CPU, and on Windows without +// `ORT_DYLIB_PATH` the ONNX Runtime init can block indefinitely. A +// request thread must never pay, or wait on, that construction: the +// first qualifying dispatch CASes UNINIT → INITIALIZING, hands the build +// to a background thread, and returns a NoOp; every dispatch during +// INITIALIZING is a NoOp; READY is terminal and serves the slot — or a +// deterministic NoOp forever when the cache was cold or the load failed, +// which is never retried. `warm_live_zone_compressors` runs the same +// construction synchronously off the request path (proxy startup, +// tests), so steady-state traffic normally never meets the lazy path. + +#[cfg(feature = "ml")] +mod kompress_slot { + pub const UNINIT: u8 = 0; + pub const INITIALIZING: u8 = 1; + pub const READY: u8 = 2; +} + +#[cfg(feature = "ml")] +static KOMPRESS_STATE: AtomicU8 = AtomicU8::new(kompress_slot::UNINIT); + +/// Written exactly once, by whichever initializer won the CAS; the +/// `Release` store of `READY` publishes it to `Acquire` readers. +#[cfg(feature = "ml")] +static KOMPRESS_INSTANCE: OnceLock> = OnceLock::new(); + +/// How many times the expensive construction has actually run in this +/// process. Exactly-once observability; asserted by +/// `tests/live_zone_kompress_async_init.rs`. +#[cfg(feature = "ml")] +static KOMPRESS_INIT_RUNS: AtomicUsize = AtomicUsize::new(0); + +#[doc(hidden)] +#[cfg(feature = "ml")] +pub fn kompress_init_runs() -> usize { + KOMPRESS_INIT_RUNS.load(Ordering::Relaxed) +} + +/// Without the `ml` feature no construction can ever run. +#[doc(hidden)] +#[cfg(not(feature = "ml"))] +pub fn kompress_init_runs() -> usize { + 0 +} + +/// The one construction path. `Ok(None)` on a cold cache and hard errors +/// both settle the slot as `None` — a deterministic NoOp for the life of +/// the process, never retried — and the outcome is observable via the +/// `kompress_init_complete` event's `kompress_ready` field. +#[cfg(feature = "ml")] +fn run_kompress_init() { + KOMPRESS_INIT_RUNS.fetch_add(1, Ordering::Relaxed); + let slot = match Kompress::from_cache(KompressConfig::default()) { + Ok(k) => k, // None => HF cache cold: deterministic NoOp downstream + Err(e) => { + tracing::warn!(event = "kompress_init_failed", error = %e); + None + } + }; + let ready = slot.is_some(); + let _ = KOMPRESS_INSTANCE.set(slot); + KOMPRESS_STATE.store(kompress_slot::READY, Ordering::Release); + tracing::info!(event = "kompress_init_complete", kompress_ready = ready); +} + +/// Non-blocking request-path accessor for the Kompress slot. +/// +/// Returns the model only when a completed initialization loaded it; +/// otherwise `None`, so the caller passes text through untouched. This +/// mirrors the Python router's not-ready → passthrough behavior, and it +/// NEVER waits: a virgin slot starts the build on a background thread +/// and answers `None` for this request. +#[cfg(feature = "ml")] +fn kompress_cached() -> Option<&'static Kompress> { + match KOMPRESS_STATE.load(Ordering::Acquire) { + kompress_slot::READY => KOMPRESS_INSTANCE.get().and_then(|slot| slot.as_ref()), + kompress_slot::UNINIT => { + if KOMPRESS_STATE + .compare_exchange( + kompress_slot::UNINIT, + kompress_slot::INITIALIZING, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + if let Err(e) = std::thread::Builder::new() + .name("kompress-init".into()) + .spawn(run_kompress_init) + { + // The expensive construction never started; give the + // (cheap) spawn another chance on a later dispatch. + KOMPRESS_STATE.store(kompress_slot::UNINIT, Ordering::Release); + tracing::warn!(event = "kompress_init_spawn_failed", error = %e); + } + } + None + } + // INITIALIZING (or an impossible value): fail open, never wait. + _ => None, + } +} + +/// Route `text` through Kompress when the model is cache-resident; +/// otherwise fall back to a deterministic no-op. Shared by the +/// `PlainText` arm and the `SourceCode` arm's passthrough fallback. +/// +/// Not a full port of the Python router's fallback chain: that one also +/// re-runs Kompress when the code compressor returns *changed but not +/// smaller* output, whereas this arm sends only an exact passthrough +/// here and leaves the changed-but-not-smaller case to the +/// tokenizer-rejection gate in `compress_one_block`. +#[cfg(feature = "ml")] +fn kompress_or_noop(text: &str, content_type: ContentType) -> DispatchResult { + if let Some(k) = kompress_cached() { + let result = k.compress(text); + if !result.is_passthrough() && result.compressed != text { + return DispatchResult::Compressed { + strategy: STRATEGY_KOMPRESS, + compressed: result.compressed, + }; + } + } + DispatchResult::NoOp { + content_type: content_type.as_str(), + } +} + +/// Kompress is an `ml`-feature transform; without it the arm is a +/// deterministic no-op. Separate definition rather than a lint-silenced +/// unused parameter, so the no-ml build has no dead argument to explain. +#[cfg(not(feature = "ml"))] +fn kompress_or_noop(_text: &str, content_type: ContentType) -> DispatchResult { + DispatchResult::NoOp { + content_type: content_type.as_str(), + } +} + +/// Eagerly construct every live-zone compressor singleton, off the +/// request path, and block until the Kompress slot has settled. Returns +/// whether Kompress is loaded and ready to compress — `false` on a cold +/// cache, a failed load, or a build without the `ml` feature. +/// +/// Call it at process startup (the proxy does, from a background thread) +/// or from tests that need the arm live before dispatching; the request +/// path itself never blocks either way. Deliberately ignores +/// `HEADROOM_LIVE_ZONE_DISABLE_ARMS`: the kill switch gates dispatch, +/// not model residency, so both switch configurations run the identical +/// warmed process and differ only in the dispatch decision. +pub fn warm_live_zone_compressors() -> bool { + let _ = smart_crusher(); + let _ = log_compressor(); + let _ = search_compressor(); + let _ = diff_compressor(); + let _ = code_compressor(); + warm_kompress_blocking() +} + +/// Blocking Kompress warmup: run the construction inline if this thread +/// wins the slot, otherwise wait for whichever initializer did. A poll +/// loop rather than `OnceLock::wait`, deliberately: a request thread +/// that won the CAS but failed to spawn its init thread resets the +/// state to UNINIT without ever setting the instance, and a waiter +/// parked on the `OnceLock` would sleep forever through that reset. +#[cfg(feature = "ml")] +fn warm_kompress_blocking() -> bool { + loop { + match KOMPRESS_STATE.load(Ordering::Acquire) { + kompress_slot::READY => { + return KOMPRESS_INSTANCE.get().is_some_and(|slot| slot.is_some()); + } + kompress_slot::UNINIT => { + if KOMPRESS_STATE + .compare_exchange( + kompress_slot::UNINIT, + kompress_slot::INITIALIZING, + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + run_kompress_init(); + } + // Lost the race to another initializer: re-read the state. + } + _ => std::thread::sleep(std::time::Duration::from_millis(10)), + } + } +} + +/// Without the `ml` feature there is no Kompress slot to warm. +#[cfg(not(feature = "ml"))] +fn warm_kompress_blocking() -> bool { + false +} + +// ─── Kill switch (env-var arm disable) ───────────────────────────────── +// +// Experiment control and rollback affordance: set +// `HEADROOM_LIVE_ZONE_DISABLE_ARMS=source_code,plain_text` (a +// comma-separated list of content-type names, in either the +// `ContentType::as_str()` or `natural_name()` spelling) to force those +// arms to a deterministic no-op, bypassing their compressor entirely. +// +// One operator-facing caveat: granularity is per ARM, not per +// compressor. `plain_text` alone does not stop Kompress runs reached +// through the SourceCode arm's passthrough fallback; disabling +// `source_code` as well is what removes Kompress from traffic entirely +// (at the cost of also disabling the code compressor). + +/// Parse the comma-separated arm list. Blank entries are skipped; +/// unknown names are ignored with a warning rather than failing the +/// request, because a typo in an operator's rollback switch must not +/// take the proxy down. +/// +/// Pure and `pub` so the parsing contract can be tested directly. The +/// alternative — asserting on it through the process environment — +/// requires `std::env::set_var`, which is unsound in a multi-threaded +/// process and becomes `unsafe` in edition 2024. +#[must_use] +pub fn parse_disabled_arms(raw: &str) -> HashSet { + let mut set = HashSet::new(); + for token in raw.split(',') { + let token = token.trim(); + if token.is_empty() { + continue; + } + match token.parse::() { + Ok(ct) => { + set.insert(ct); + } + Err(_) => tracing::warn!(event = "disable_arms_unknown_token", token = %token), + } + } + set +} + +/// Arms disabled via `HEADROOM_LIVE_ZONE_DISABLE_ARMS`. +/// +/// Latched on **first dispatch**, not at process start: the value is read +/// the first time a block is routed and cached for the process lifetime, +/// so a run's arm-disable set cannot change mid-flight (determinism +/// invariant). A non-UTF-8 value is read lossily rather than silently +/// treated as unset. +fn disabled_arms() -> &'static HashSet { + static INSTANCE: OnceLock> = OnceLock::new(); + INSTANCE.get_or_init(|| { + std::env::var_os("HEADROOM_LIVE_ZONE_DISABLE_ARMS") + .map(|raw| parse_disabled_arms(&raw.to_string_lossy())) + .unwrap_or_default() + }) +} + +/// True when `content_type`'s dispatch arm has been disabled via +/// `HEADROOM_LIVE_ZONE_DISABLE_ARMS`. +fn arm_disabled(content_type: ContentType) -> bool { + disabled_arms().contains(&content_type) +} + // ─── Public entry point ──────────────────────────────────────────────── /// Inspect a buffered Anthropic `/v1/messages` body and decide which @@ -1318,15 +1616,20 @@ enum DispatchResult { /// Map `(text, content_type)` to the compressor result. /// -/// Per spec PR-B3: +/// Per spec PR-B3 (arms marked PR-B4 wired by that PR): /// /// - `JsonArray` (with `is_dict_array=true`) → SmartCrusher /// - `BuildOutput` → LogCompressor /// - `SearchResults` → SearchCompressor /// - `GitDiff` → DiffCompressor -/// - `SourceCode` → no-op (Rust port pending; see TODO below) -/// - `PlainText` → no-op (PR-B4 wires Kompress) +/// - `SourceCode` → CodeAwareCompressor, falling back to cached +/// Kompress on exact passthrough (PR-B4) +/// - `PlainText` → cached Kompress, cache-cold → no-op (PR-B4) /// - `Html` → no-op (no compressor) +/// +/// Any arm whose content type is named in +/// `HEADROOM_LIVE_ZONE_DISABLE_ARMS` short-circuits to a no-op before +/// the table below is consulted. fn dispatch_compressor(text: &str, content_type: ContentType) -> DispatchResult { if text.is_empty() { return DispatchResult::NoOp { @@ -1334,6 +1637,16 @@ fn dispatch_compressor(text: &str, content_type: ContentType) -> DispatchResult }; } + // Kill switch, checked once for every arm rather than per arm: an + // operator naming any content type gets that arm disabled, instead + // of the switch silently covering only the two arms whose bodies + // happened to consult it. + if arm_disabled(content_type) { + return DispatchResult::NoOp { + content_type: content_type.as_str(), + }; + } + match content_type { ContentType::JsonArray => { // The detector classifies arrays-of-scalars as JsonArray @@ -1387,18 +1700,39 @@ fn dispatch_compressor(text: &str, content_type: ContentType) -> DispatchResult compressed: result.compressed, } } - // TODO(PR-B4 / Rust code-compressor port): Python has a - // CodeAwareCompressor; the Rust port is not yet shipped. Once - // that crate lands, `ContentType::SourceCode` routes here - // exactly as the others above. - ContentType::SourceCode => DispatchResult::NoOp { - content_type: content_type.as_str(), - }, - // TODO(PR-B4): wire Kompress (lossless prose compressor) for - // PlainText. For now, leave untouched. - ContentType::PlainText => DispatchResult::NoOp { - content_type: content_type.as_str(), - }, + // Routes to the tree-sitter-backed CodeAwareCompressor (PR-B4). + // The compressor fails open internally (re-parse-and-revert on + // syntax errors, a min-token floor, and a compression-ratio + // guard), so this arm only has to compare its output against + // the input; the tokenizer-rejection gate in + // `compress_one_block` applies on top of that. When the code + // compressor doesn't shrink the block (e.g. it's already + // dense, or tree-sitter reverted on a parse error), fall back + // to Kompress rather than giving up — mirrors the Python + // router's no-shrink → Kompress fallback. + ContentType::SourceCode => { + let result = code_compressor().compress(text); + if result.compressed == text { + // Exact passthrough only: the compressor declined + // (min-token floor, unknown language, parse failure, + // ratio guard), so hand the block to Kompress. A + // rewrite that changed bytes without shrinking tokens + // is NOT retried here — `compress_one_block`'s + // tokenizer gate rejects it and forwards the original. + return kompress_or_noop(text, content_type); + } + DispatchResult::Compressed { + strategy: STRATEGY_CODE_COMPRESSOR, + compressed: result.compressed, + } + } + // Cache-only EXTRACTIVE prose compressor (PR-B4): Kompress drops + // low-salience words and rejoins on single spaces — lossy by + // design. `kompress_or_noop` degrades to a deterministic no-op + // when the model isn't cache-resident, hasn't finished its + // background initialization yet, or the `ml` feature is off — + // no network call and no model build ever happens on this path. + ContentType::PlainText => kompress_or_noop(text, content_type), // No HTML compressor on the Rust side; pages are handled by // upstream extractors, not the proxy. ContentType::Html => DispatchResult::NoOp { diff --git a/crates/headroom-core/tests/common/mod.rs b/crates/headroom-core/tests/common/mod.rs new file mode 100644 index 000000000..0b51c4898 --- /dev/null +++ b/crates/headroom-core/tests/common/mod.rs @@ -0,0 +1,171 @@ +//! Shared fixtures for the live-zone integration tests and their +//! dedicated single-test binaries. +//! +//! Integration test files are separate crates, so each `mod common;` +//! gets its own copy of this module — the standard Rust layout for +//! sharing test helpers (The Book, ch. 11.3). Before this module the +//! same body builders, Python/prose generators and HuggingFace-cache +//! probes were pasted into four targets, where the cache probe in +//! particular had to be kept byte-compatible with the production +//! loader by hand. +//! +//! Not every target uses every helper, so unused items here are +//! expected rather than a defect. +#![allow(dead_code)] + +use headroom_core::transforms::live_zone::DEFAULT_MODEL; +use headroom_core::transforms::{ + compress_anthropic_live_zone, AuthMode, BlockAction, CompressionManifest, LiveZoneOutcome, +}; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; + +/// Serialize a JSON value to a request body. +pub fn body_of(value: &Value) -> Vec { + serde_json::to_vec(value).expect("fixture body serializes") +} + +/// Run the public dispatcher entry point with no frozen prefix. +pub fn dispatch(body: &[u8]) -> LiveZoneOutcome { + compress_anthropic_live_zone(body, 0, AuthMode::Payg, DEFAULT_MODEL) + .expect("dispatcher returns Ok on valid bodies") +} + +pub fn sha256(bytes: &[u8]) -> [u8; 32] { + let mut h = Sha256::new(); + h.update(bytes); + h.finalize().into() +} + +/// Byte range of the first occurrence of `needle` in `haystack`, +/// half-open. Used to locate the JSON-encoded `content` slot for +/// byte-fidelity assertions. +pub fn find_byte_range(haystack: &[u8], needle: &[u8]) -> (usize, usize) { + let pos = haystack + .windows(needle.len()) + .position(|w| w == needle) + .unwrap_or_else(|| { + panic!( + "needle of {} bytes not found in haystack of {} bytes", + needle.len(), + haystack.len() + ) + }); + (pos, pos + needle.len()) +} + +/// A body with one user message holding one `tool_result` whose +/// `content` is `text`. Returns the body and the byte range of the +/// JSON-encoded content slot (quotes included) inside it. +pub fn body_with_tool_result(text: &str) -> (Vec, (usize, usize)) { + let body = body_of(&json!({ + "model": "claude-sonnet-4-6", + "max_tokens": 64, + "system": "you are a helpful assistant", + "messages": [{ + "role": "user", + "content": [{ + "type": "tool_result", + "tool_use_id": "toolu_live_zone_test", + "content": text, + }], + }], + })); + // The encoded slot is exactly `to_vec(&text)`: serde uses the same + // string encoding for the embedded value. + let needle = serde_json::to_vec(&text).expect("text serializes"); + let range = find_byte_range(&body, &needle); + (body, range) +} + +/// The `tool_result` block's action from a manifest, cloned out. +pub fn tool_result_action(manifest: &CompressionManifest) -> BlockAction { + manifest + .block_outcomes + .iter() + .find(|b| b.block_type == "tool_result") + .expect("tool_result block present in manifest") + .action + .clone() +} + +/// Syntactically valid Python with `n` small functions, each with a +/// docstring and a body longer than `CodeCompressorConfig`'s default +/// `max_body_lines` (5) so body elision has something to trim. Real +/// Python: the compressor re-parses and reverts on syntax errors, so a +/// fixture that merely looks like code would silently pass through. +pub fn python_module_source(n: usize) -> String { + use std::fmt::Write as _; + + let mut code = String::from( + "\"\"\"Example data-processing module used by the live-zone tests.\"\"\"\n\n\ + import json\n\ + import os\n\ + from typing import Any, Optional\n\n\n", + ); + for i in 0..n { + let _ = write!( + code, + "def process_record_{i}(record: dict) -> dict:\n \ + \"\"\"Normalize record {i} and compute its derived fields.\"\"\"\n \ + result = dict(record)\n \ + result[\"index\"] = {i}\n \ + result[\"doubled\"] = record.get(\"value\", 0) * 2\n \ + result[\"source\"] = \"batch\"\n \ + if result[\"doubled\"] > 100:\n \ + result[\"flag\"] = \"high\"\n \ + else:\n \ + result[\"flag\"] = \"low\"\n \ + return result\n\n\n" + ); + } + code +} + +/// Repetitive plain prose of at least `min_bytes`, varied enough to +/// classify as `PlainText`. +pub fn plain_prose(min_bytes: usize) -> String { + use std::fmt::Write as _; + + let mut text = String::with_capacity(min_bytes + 256); + let mut i = 0usize; + while text.len() < min_bytes { + let _ = write!( + text, + "City officials announced today that the downtown revitalization \ + project will proceed as planned despite budget concerns raised \ + during round {i} of public comment. " + ); + i += 1; + } + text +} + +/// A JSON array of dicts, `n` entries — SmartCrusher's shape. +pub fn json_array_of_dicts(n: usize) -> String { + let rows: Vec = (0..n) + .map(|i| { + json!({ + "id": i, + "status": "ok", + "value": format!("repeat-pattern-{}", i % 3), + }) + }) + .collect(); + serde_json::to_string(&rows).expect("fixture array serializes") +} + +/// Whether the Kompress model is cache-resident and loaded, asked of the +/// production slot itself: this runs `live_zone`'s blocking warmup — the +/// proxy's own startup path — and reports what it found, rather than +/// re-deriving HuggingFace cache paths in a probe that could desync. +/// +/// Call it BEFORE dispatching content that should reach the arm. +/// Dispatch never waits on model construction, so an unwarmed first +/// dispatch is a NoOp by design — the contract pinned by +/// `live_zone_kompress_async_init.rs`. Warming here both settles the +/// slot the dispatch under test will read and prices the probe at one +/// shared construction instead of a throwaway one. +pub fn kompress_available() -> bool { + headroom_core::transforms::live_zone::warm_live_zone_compressors() +} diff --git a/crates/headroom-core/tests/live_zone_arm_properties.rs b/crates/headroom-core/tests/live_zone_arm_properties.rs new file mode 100644 index 000000000..9885077d5 --- /dev/null +++ b/crates/headroom-core/tests/live_zone_arm_properties.rs @@ -0,0 +1,823 @@ +//! Property coverage for the PR-B4 dispatch arms (`SourceCode` → +//! `CodeAwareCompressor`, `PlainText` → Kompress, cache-only) plus the +//! round-trip test that closes the kill-switch alias bug class. +//! +//! Background: `ContentType`'s string parsing used to accept only each +//! variant's `as_str()` tag, so an `HEADROOM_LIVE_ZONE_DISABLE_ARMS` +//! token spelled the "natural" way a human is more likely to write +//! (e.g. `plain_text`, `search_results`) silently failed to match. The +//! token fell through to the unknown-token branch, the arm was never +//! disabled, and there was no error — a misconfiguration that looked +//! like a no-op. Both spellings now parse (`ContentType`'s `FromStr`), +//! and the round-trip test below pins that for every variant by +//! iterating `ContentType::ALL` rather than a hand-maintained copy of +//! the variant list. +//! +//! Property-test coverage for the dispatch arms themselves follows +//! below: no-panic and determinism properties (proptest), plus a +//! byte-fidelity test for the SourceCode arm cloned from +//! `live_zone_dispatch.rs::byte_fidelity_outside_compressed_block`, plus +//! an instrumentation test that measures — rather than +//! asserts — what fraction of generated cases actually reach a dispatch +//! arm, broken down by content type. House style for the proptest +//! blocks mirrors `live_zone_token_validation.rs:201-254` — dispatch +//! only through the public `compress_anthropic_live_zone` entry point +//! over generated `tool_result` bodies, never the private +//! `dispatch_compressor` — and the no-panic parser fuzz tests in +//! `headroom-proxy/tests/sse_framing.rs:156-200` for case-count order +//! of magnitude and the comment style explaining the choice. + +mod common; + +use common::{body_with_tool_result, dispatch, python_module_source, sha256}; +use headroom_core::transforms::live_zone::threshold_for; +use headroom_core::transforms::{detect_content_type, BlockAction, ContentType, LiveZoneOutcome}; +use proptest::prelude::*; +use proptest::strategy::ValueTree; +use proptest::test_runner::{RngAlgorithm, TestRng, TestRunner}; +use serde_json::Value; + +// ─── Part 0: kill-switch alias round trip (bug class extinction) ────── + +/// The bug class this test makes extinct: a valid natural-name spelling +/// of a `ContentType` failing to parse, so +/// `HEADROOM_LIVE_ZONE_DISABLE_ARMS` looks like it disabled an arm but +/// didn't. +/// +/// Iterates `ContentType::ALL`, so a variant added later is covered +/// here automatically — the earlier version of this test carried its +/// own hand-maintained variant array plus a mirrored count and a +/// compile-time length assertion, none of which could actually tell +/// that a new variant had gone missing from the array. +#[test] +fn content_type_parses_from_both_spellings_for_all_variants() { + for content_type in ContentType::ALL { + let tag = content_type.as_str(); + assert_eq!( + tag.parse::(), + Ok(content_type), + "as_str() tag {tag:?} must round-trip back to {content_type:?}" + ); + + let natural = content_type.natural_name(); + assert_eq!( + natural.parse::(), + Ok(content_type), + "natural-name alias {natural:?} must parse to {content_type:?}" + ); + } +} + +/// An unknown name must be a parse error, not a silently-wrong variant. +#[test] +fn content_type_rejects_unknown_names() { + for name in ["", "bogus_type", "Source_Code", "SOURCE_CODE", "plaintext"] { + assert!( + name.parse::().is_err(), + "{name:?} must not parse to a ContentType" + ); + } +} + +// ─── Part 1: pathological-text generators (properties 1 & 2) ────────── + +/// Pure-ASCII source fragment used by the "half-truncated code +/// snippet" arm of [`pathological_text`]. ASCII-only so every byte +/// index is also a valid `char` boundary — slicing at an arbitrary +/// length can never panic on a UTF-8 boundary violation. +const CODE_FRAGMENT: &str = "def handler(event, context):\n \ +payload = json.loads(event[\"body\"])\n \ +if payload.get(\"kind\") == \"ping\":\n \ +return {\"statusCode\": 200, \"body\": \"pong\"}\n \ +result = process(payload)\n \ +return {\"statusCode\": 200, \"body\": json.dumps(result)}\n"; + +/// Function-name pool for the "code-like" generator bucket's varied +/// identifiers — deliberately NOT the +/// single `process_record_{i}` pattern `python_module_source` uses for +/// its fixed byte-fidelity fixture. +const CODE_IDENTIFIERS: &[&str] = &[ + "process_record", + "normalize_entry", + "compute_score", + "parse_payload", + "build_summary", + "validate_input", + "apply_filter", + "merge_results", + "extract_fields", + "transform_row", + "classify_item", + "enrich_context", + "dedupe_values", + "flatten_tree", + "sanitize_text", + "batch_update", + "load_config", + "fetch_metadata", + "index_documents", + "score_candidate", + "resolve_alias", + "collect_stats", + "prune_stale", + "rank_matches", + "join_segments", + "split_batch", + "verify_checksum", + "encode_payload", + "decode_payload", + "queue_task", +]; + +/// Render one function: a docstring, then `body_lines` body statements +/// (floored at 2: `result = ...` plus `return result`), then a blank +/// line. `body_lines` deliberately spans both sides of +/// `CodeCompressorConfig::default().max_body_lines` (5) — a whole +/// function node of `<= max_body_lines + 2 == 7` lines passes through +/// `compress_function_ast` untouched, longer ones get their body +/// collapsed — so callers can generate a mix of both shapes in one +/// module instead of the "always long enough to collapse" shape +/// `python_module_source` uses. +fn render_function(name: &str, index: usize, body_lines: usize) -> String { + let body_lines = body_lines.max(2); + let filler = body_lines - 2; + let mut out = String::new(); + out.push_str(&format!("def {name}_{index}(record: dict) -> dict:\n")); + out.push_str(&format!( + " \"\"\"Derive fields for {name} #{index}.\"\"\"\n" + )); + out.push_str(" result = dict(record)\n"); + for i in 0..filler { + out.push_str(&format!(" result[\"{name}_{i}\"] = {i}\n")); + } + out.push_str(" return result\n"); + out.push('\n'); + out +} + +/// Syntactically-plausible Python, 2048-6000 bytes, with VARIED +/// structure: differing function counts and body lengths straddling +/// the CodeAwareCompressor's 5-line collapse floor, varied identifiers +/// drawn from `CODE_IDENTIFIERS`. Added after a review measured that +/// before this bucket existed, +/// 0% of `pathological_text()`'s generated cases ever reached the +/// SourceCode dispatch arm (the fixed `python_module_source(10)` +/// fixture in `byte_fidelity_outside_compressed_source_block` was +/// SourceCode's only exercise anywhere in this file). +/// +/// Always classifies `SourceCode` — but NOT because of the header +/// alone: `try_detect_code`'s confidence is +/// `0.4 + (matching_lines / non_empty_lines) * 0.4 + matching_lines * 0.02`, +/// so a fixed count of header matches DILUTES below the 0.5 floor once +/// enough non-matching lines follow (~4 matches over 100 non-empty +/// lines ≈ 0.496). What actually holds classification is that every +/// `render_function` body emits `def ...:` / docstring lines that also +/// match `CODE_PATTERNS`, keeping `matching_lines` roughly proportional +/// to length. Consequence for future edits: rewriting `render_function` +/// with shapes that DON'T match the detector's patterns (e.g. +/// assignment-bound lambdas) would silently reclassify large samples as +/// PlainText — `dispatch_reach_fractions_meet_floor` fails loudly when +/// that happens; believe it, and look here first. +/// +/// Byte-length mechanism: pick a `target_bytes` first, then append +/// functions from a pool of varied specs until the buffer crosses it +/// (each function bounded to at most ~510 bytes by `body_lines`'s 1..=10 +/// range, so the overshoot past `target_bytes`'s 5000 ceiling can't +/// reach the outer 6000 bound). If the pool shrinks smaller than +/// `target_bytes` needs (proptest's shrinker actively tries this), a +/// deterministic filler function pads the buffer past the 2048 floor +/// regardless — this keeps the "always >= 2048 bytes, always +/// SourceCode" invariant true even for a minimized failing case, so a +/// shrunk counterexample stays inside the region that made it +/// interesting in the first place. +fn code_like_source() -> impl Strategy { + ( + 2_048usize..=5_000, + proptest::collection::vec( + (proptest::sample::select(CODE_IDENTIFIERS), 1usize..=10), + 1..=200, + ), + ) + .prop_map(|(target_bytes, specs)| { + let mut code = String::from( + "\"\"\"Generated module for the code-like dispatch-reach bucket.\"\"\"\n\n\ + import json\n\ + import os\n\ + from typing import Any, Optional\n\n\n", + ); + for (index, (name, body_lines)) in specs.into_iter().enumerate() { + if code.len() >= target_bytes { + break; + } + code.push_str(&render_function(name, index, body_lines)); + } + // Guaranteed floor (see doc comment above): pad with a + // deterministic filler function so this bucket always + // clears THRESHOLD_SOURCE_CODE (2048), even for a pool too + // small (or too shrunk) to reach `target_bytes` on its own. + let mut filler_index = 10_000usize; + while code.len() < 2_048 { + code.push_str(&render_function("filler", filler_index, 8)); + filler_index += 1; + } + code + }) +} + +/// Vocabulary for the "prose-like" generator bucket: ordinary lowercase +/// English words chosen to have NO overlap with any content_detector.rs +/// keyword or symbol (`def`, `class`, `import`, `ERROR`, `diff`, `<`, +/// `[`, `:`, `@`, ...), so generated text is guaranteed to fall through +/// every specialized detector (JSON / diff / HTML / search / log / +/// code) and land on `PlainText`'s default branch. +const PROSE_WORDS: &[&str] = &[ + "the", + "team", + "reviewed", + "quarterly", + "report", + "and", + "found", + "several", + "interesting", + "trends", + "worth", + "discussing", + "further", + "during", + "next", + "weeks", + "meeting", + "customers", + "have", + "been", + "asking", + "about", + "new", + "features", + "that", + "would", + "improve", + "their", + "daily", + "workflow", + "without", + "adding", + "unnecessary", + "complexity", + "to", + "existing", + "processes", + "engineers", + "spent", + "most", + "of", + "afternoon", + "debugging", + "a", + "tricky", + "issue", + "related", + "caching", + "behavior", + "under", + "heavy", + "load", + "documentation", + "was", + "updated", + "reflect", + "recent", + "changes", + "in", + "policy", + "procedure", + "across", + "departments", + "stakeholders", + "expressed", + "cautious", + "optimism", + "regarding", + "timeline", + "for", + "launch", + "despite", + "lingering", + "concerns", + "budget", + "allocation", + "remains", + "topic", + "ongoing", + "discussion", + "among", + "leadership", + "training", + "materials", + "were", + "distributed", + "all", + "staff", + "ahead", + "upcoming", + "transition", + "period", + "feedback", + "collected", + "from", + "survey", + "suggests", + "broad", + "satisfaction", + "with", + "current", + "support", + "channels", + "seasonal", + "demand", + "typically", + "increases", + "toward", + "end", + "fiscal", + "year", + "requiring", + "additional", + "planning", + "resources", + "roadmap", + "priorities", + "shifted", + "slightly", + "after", + "user", + "research", + "revealed", + "unexpected", + "usage", + "patterns", + "warrant", + "deeper", + "study", + "onboarding", + "experience", + "still", + "feels", + "clunky", + "several", + "early", + "testers", + "noted", + "confusion", + "around", + "default", + "settings", + "release", + "notes", + "circulated", + "internally", + "before", + "public", + "announcement", + "went", + "out", + "later", + "same", + "week", +]; + +/// Prose-like text, 5200-11800 bytes, built from `PROSE_WORDS` with +/// varied sentence/paragraph shape — deliberately NOT a single repeated +/// character, unlike the pre-existing `pathological_text()` bucket that +/// repeats `'x'` 1000-8000 times. Added by the same review as +/// `code_like_source`: the problem with the old repeated-`'x'` +/// bucket was that it only reaches the PlainText arm by luck of length +/// crossing 5120 bytes, not because it looks like prose — this bucket +/// looks like prose AND always clears the threshold. +/// +/// Same target-then-fill-then-pad mechanism as `code_like_source` +/// (see its doc comment): the length floor (5200 bytes, comfortably +/// above `THRESHOLD_PLAIN_TEXT`'s 5120) holds even for a shrunk word +/// pool, so a minimized counterexample from this bucket stays inside +/// the region that made it interesting. +fn prose_like_text() -> impl Strategy { + ( + 5_200usize..=11_800, + proptest::collection::vec(proptest::sample::select(PROSE_WORDS), 1..=2_200), + ) + .prop_map(|(target_bytes, words)| { + let mut text = String::with_capacity(target_bytes + 64); + let mut words_in_sentence = 0u32; + for word in words { + if text.len() >= target_bytes { + break; + } + text.push_str(word); + words_in_sentence += 1; + // Vary sentence/paragraph shape instead of one giant + // run-on line: end a "sentence" every ~9 words. + if words_in_sentence >= 9 { + text.push_str(".\n"); + words_in_sentence = 0; + } else { + text.push(' '); + } + } + while text.len() < 5_200 { + text.push_str("padding "); + } + text + }) +} + +/// Strategy generating "pathological" text for the dispatcher's +/// no-panic and determinism properties below. Mixes: +/// +/// - plain arbitrary strings (the common case, still worth covering; +/// proptest 1.11.0's `any::()` caps at 32 chars, so this +/// bucket never clears any dispatch threshold on its own), +/// - control characters (0x00-0x1F, capped ~256 bytes) — the sort of +/// byte a shell or log scraper can hand a tool_result without ever +/// going through a terminal's escaping, +/// - unpaired UTF-16 surrogates (capped ~384 bytes), repaired via +/// `String::from_utf16_lossy` (a Rust `String` can't hold an actual +/// lone surrogate — this is the closest in-process torture test: what +/// a naive UTF-16 → UTF-8 bridge upstream would hand us after +/// "fixing" a bad pair), +/// - half-truncated code snippets (capped ~272 bytes; a coding agent's +/// tool output cut off mid-token is a realistic production shape, not +/// just a fuzz artifact), +/// - very long single lines with no newlines, 1000-8000 chars (minified +/// JS, a base64 blob, ...) — before the two rich buckets below were +/// added, the only bucket that +/// could ever clear a threshold (PlainText's 5120, when the sampled +/// length lands >= 5120, ~41% of this bucket's own range), +/// - `code_like_source()` — syntactically +/// plausible, varied-structure Python that always clears the +/// SourceCode threshold (2048) and always detects as SourceCode, +/// - `prose_like_text()` — varied +/// natural-language-shaped text that always clears the PlainText +/// threshold (5120) and always detects as PlainText. +/// +/// Weights are set so the two rich buckets carry +/// enough weight to matter (see `dispatch_reach_fractions_meet_floor` +/// for the measured, instrumented result) while keeping every small +/// pathological bucket from before — the no-panic envelope over +/// encoding/detection/threshold-comparison logic on genuinely +/// pathological BYTES (not just large-but-tame text) is still valuable +/// in its own right. +fn pathological_text() -> impl Strategy { + prop_oneof![ + 2 => any::(), + 1 => proptest::collection::vec(0u8..0x20u8, 0..256) + .prop_map(|bytes| bytes.into_iter().map(|b| b as char).collect()), + 1 => proptest::collection::vec(any::(), 0..128) + .prop_map(|units| String::from_utf16_lossy(&units)), + 1 => (0..=CODE_FRAGMENT.len()).prop_map(|n| CODE_FRAGMENT[..n].to_string()), + 1 => (1_000usize..8_000).prop_map(|n| "x".repeat(n)), + 2 => code_like_source(), + 2 => prose_like_text(), + ] +} + +// The dispatcher must never panic on arbitrary text, however +// pathological, and must be deterministic. +// +// The two richer generator buckets (real, +// varied-structure source code and prose, both large enough to reach a +// dispatch arm) make each case noticeably more expensive than the old +// all-small-or-threshold-capped mix — a meaningful share of cases now +// actually run CodeAwareCompressor's tree-sitter parse or the Kompress +// path instead of bottoming out at `BelowByteThreshold` immediately. +// Case counts were cut from the earlier 2048/1024 (order of magnitude +// of `sse_framing.rs`'s parser fuzz tests) down to 40/20 after two +// larger settings both measured over the ~60s guidance with the richer +// generators in place: 128/64 measured ~84s combined, 64/32 measured +// ~60s combined (right at the boundary, not comfortably under it). +// 40/20 is what actually landed with margin. If a future +// change to these generators pushes wall time back over budget, the +// guidance is to cut case counts further, not shrink the payload +// ranges back down to threshold-capped sizes (that would silently +// regress dispatch reach). The no-panic property keeps the larger count +// since it's the property most likely to catch a real crash, the +// determinism property the smaller since each case dispatches twice. +// +// Kompress may be cache-cold on this machine (the HF cache is +// per-machine, not part of the repo) — both properties below must +// hold regardless: `kompress_or_noop` degrades to a deterministic +// NoOp when the model isn't cache-resident, which is a valid outcome +// for both "never panics" and "deterministic", not a test dependency +// on the model being loaded. + +proptest! { + #![proptest_config(ProptestConfig { + cases: 40, + // Give the shrinker room to minimize any panic it finds down + // to a small repro instead of giving up early. + max_shrink_iters: 1024, + ..ProptestConfig::default() + })] + + /// Property 1: for arbitrary (including pathological) `String`s + /// embedded as a `tool_result` body, dispatching through the + /// public `compress_anthropic_live_zone` entry point must never + /// panic. + #[test] + fn dispatch_no_panic_on_arbitrary_text(text in pathological_text()) { + let (body, _) = body_with_tool_result(&text); + // `dispatch` itself panics (via `.expect`) only on a + // dispatcher `Err`, which a well-formed JSON body constructed + // above can't produce. Reaching the end of this closure + // without unwinding IS the property. + let _ = dispatch(&body); + } +} + +proptest! { + #![proptest_config(ProptestConfig { + cases: 20, + ..ProptestConfig::default() + })] + + /// Property 2: determinism. The dispatcher's only process-global + /// state is the `HEADROOM_LIVE_ZONE_DISABLE_ARMS` kill-switch set + /// (unset in this file, and latched once regardless), so the same + /// input bytes must always produce the same output bytes AND the + /// same manifest. `CompressionManifest` / `BlockAction` are + /// `Debug`-only (no `PartialEq` — they're observability types, not + /// meant for equality comparisons in production code), so this + /// compares their `Debug` renderings as a structural-equality + /// proxy, the standard workaround for that situation. + #[test] + fn dispatch_deterministic_same_bytes(text in pathological_text()) { + let (body, _) = body_with_tool_result(&text); + + let out1 = dispatch(&body); + let out2 = dispatch(&body); + + let (bytes1, manifest1) = match &out1 { + LiveZoneOutcome::NoChange { manifest } => (body.clone(), format!("{manifest:?}")), + LiveZoneOutcome::Modified { new_body, manifest } => { + (new_body.get().as_bytes().to_vec(), format!("{manifest:?}")) + } + }; + let (bytes2, manifest2) = match &out2 { + LiveZoneOutcome::NoChange { manifest } => (body.clone(), format!("{manifest:?}")), + LiveZoneOutcome::Modified { new_body, manifest } => { + (new_body.get().as_bytes().to_vec(), format!("{manifest:?}")) + } + }; + + prop_assert_eq!( + &bytes1, &bytes2, + "same input bytes must yield the same output bytes (bytes in -> bytes out)" + ); + prop_assert_eq!( + manifest1, manifest2, + "same input bytes must yield the same manifest" + ); + } +} + +// ─── Part 1, property 3: byte fidelity around the SourceCode arm ────── + +#[test] +fn byte_fidelity_outside_compressed_source_block() { + // Same central invariant as `live_zone_dispatch.rs`'s + // `byte_fidelity_outside_compressed_block` (the B3 SmartCrusher + // pin), cloned onto the PR-B4 SourceCode/CodeAwareCompressor arm: + // bytes OUTSIDE the rewritten block must hash byte-identical to + // the input, regardless of which compressor did the rewriting. + let code = python_module_source(10); + assert!( + code.len() > 2048, + "fixture must clear the SourceCode byte threshold (2048); got {} bytes", + code.len() + ); + + let (body_in, content_range) = body_with_tool_result(&code); + let (block_start, block_end) = content_range; + + let out = dispatch(&body_in); + let (new_body, strategy) = match &out { + LiveZoneOutcome::Modified { new_body, manifest } => { + let action = manifest + .block_outcomes + .iter() + .find(|b| b.block_type == "tool_result") + .expect("tool_result block present in manifest") + .action + .clone(); + let strategy = match action { + BlockAction::Compressed { strategy, .. } => strategy, + other => panic!( + "expected Compressed action for a 10-function Python module, got {other:?}" + ), + }; + (new_body.get().as_bytes().to_vec(), strategy) + } + LiveZoneOutcome::NoChange { manifest } => panic!( + "expected CodeAwareCompressor to shrink a 10-function Python module; \ + got NoChange. manifest: {manifest:?}" + ), + }; + assert_eq!( + strategy, "code_compressor", + "expected code_compressor dispatch for SourceCode content" + ); + + // Prefix bytes (before the content slot) must be byte-identical. + let in_prefix = &body_in[..block_start]; + let out_prefix = &new_body[..block_start]; + assert_eq!( + sha256(in_prefix), + sha256(out_prefix), + "prefix bytes outside the compressed block must be byte-equal" + ); + + // Suffix length will differ by the compression delta, so locate + // the suffix in the output by length: it's the trailing + // (in.len() - block_end) bytes. + let in_suffix_len = body_in.len() - block_end; + let in_suffix = &body_in[block_end..]; + let out_suffix = &new_body[new_body.len() - in_suffix_len..]; + assert_eq!( + sha256(in_suffix), + sha256(out_suffix), + "suffix bytes outside the compressed block must be byte-equal" + ); + + // Output must still be valid JSON with the untouched top-level + // fields intact. + let parsed: Value = serde_json::from_slice(&new_body).expect("output is valid JSON"); + assert_eq!(parsed["model"], "claude-sonnet-4-6"); + assert_eq!(parsed["system"], "you are a helpful assistant"); +} + +/// The same byte-fidelity invariant for the PlainText/Kompress arm. +/// +/// Kompress is cache-only: on a host with no cached model the arm falls +/// open to a no-op by design, so a test that merely tolerates both +/// outcomes asserts nothing on a cold host. This one establishes ground +/// truth FIRST — asking the loader directly whether the model is +/// cache-resident — and then requires the dispatcher to agree: +/// available means the arm MUST compress and preserve the bytes outside +/// the block; unavailable means it MUST be a no-op. Either way the +/// assertion is real, and a regression that silently disabled the arm +/// on a warm host fails here rather than passing as "cold cache". +#[test] +fn plain_text_arm_preserves_bytes_outside_the_compressed_block() { + let prose = common::plain_prose(8_000); + assert!( + prose.len() > 5_120, + "fixture must clear the PlainText byte threshold (5120); got {} bytes", + prose.len() + ); + let (body, (start, end)) = body_with_tool_result(&prose); + let prefix_hash = sha256(&body[..start]); + let suffix_hash = sha256(&body[end..]); + + let kompress_available = common::kompress_available(); + + match dispatch(&body) { + LiveZoneOutcome::NoChange { manifest } => { + assert!( + !kompress_available, + "Kompress IS cache-resident on this host, so the PlainText arm should have compressed a {}-byte prose block above the 5120-byte threshold -- a NoChange here means the arm is unwired or silently disabled, which is exactly the regression this test exists to catch. manifest: {manifest:?}", + prose.len() + ); + } + LiveZoneOutcome::Modified { new_body, .. } => { + assert!( + kompress_available, + "the PlainText arm compressed a block while the loader reports no cached model" + ); + let new_bytes = new_body.get().as_bytes().to_vec(); + let suffix_len = body.len() - end; + assert!( + new_bytes.len() > start + suffix_len, + "rewritten body is too short to contain the untouched prefix and suffix" + ); + let new_end = new_bytes.len() - suffix_len; + assert_eq!( + sha256(&new_bytes[..start]), + prefix_hash, + "PlainText arm rewrote bytes BEFORE the compressed block" + ); + assert_eq!( + sha256(&new_bytes[new_end..]), + suffix_hash, + "PlainText arm rewrote bytes AFTER the compressed block" + ); + assert!( + new_bytes.len() < body.len(), + "arm reported Modified without shrinking the body" + ); + } + } +} + +// ─── Part 1, instrumentation: measured dispatch-reach fractions ─────── +// +// A review quantified (from pinned +// proptest 1.11.0 source and the earlier `pathological_text()` +// weights) that ~4.5% of generated cases reached ANY dispatch arm and +// 0% ever reached SourceCode. This section re-measures the same +// quantity against the CURRENT generator empirically, rather than just +// asserting the fix worked. + +/// Sample count for the reach-fraction measurement below. Classification +/// runs through `detect_content_type` only — no compressor is ever +/// invoked — so this stays well under a second even at this size. +const REACH_SAMPLE_COUNT: u32 = 5_000; + +/// Measures, over `REACH_SAMPLE_COUNT` samples of `pathological_text()`, +/// what fraction actually clear their content type's byte threshold and +/// would reach a `dispatch_compressor` arm (as opposed to being filtered +/// out at `BelowByteThreshold` before any compressor runs), broken down +/// by content type. Classification uses the SAME public function the +/// dispatcher itself calls (`headroom_core::transforms::detect_content_type`, +/// invoked from `live_zone.rs` right before `compress_one_block`'s +/// threshold gate — see that call site), on the same raw content-text +/// string this file's own `body_with_tool_result` embeds, so this +/// measures the real thing rather than a guess. +/// +/// Uses a fixed-seed `TestRunner` (not the `proptest!` macro) so the +/// sample is reproducible run to run — this test's job is to produce a +/// stable, reportable number, not to hunt for new failing cases (that's +/// what the two properties above are for). +/// +/// Floors are set well below the closed-form expectation from the +/// current bucket weights (SourceCode ~20%, PlainText ~24%, overall +/// ~44% — from each bucket's weight times its probability of clearing +/// its threshold) so ordinary sampling noise +/// at N=5000 can't flake this, while a future regression that collapses +/// a bucket back to threshold-capped output (the original +/// failure mode this instrumentation exists to catch) fails loudly +/// instead of silently. +#[test] +fn dispatch_reach_fractions_meet_floor() { + let mut runner = TestRunner::new_with_rng( + ProptestConfig::default(), + TestRng::from_seed(RngAlgorithm::ChaCha, &[0x5A; 32]), + ); + let strategy = pathological_text(); + + let mut reached_total = 0u32; + let mut reached_source_code = 0u32; + let mut reached_plain_text = 0u32; + let mut sampled = 0u32; + + for _ in 0..REACH_SAMPLE_COUNT { + let tree = strategy + .new_tree(&mut runner) + .expect("pathological_text() strategy never rejects a case"); + let text = tree.current(); + if text.is_empty() { + // `dispatch_compressor` special-cases empty content to an + // unconditional NoOp before the byte-threshold gate even + // runs (see live_zone.rs) — "reach" is undefined for it. + // Vanishingly rare from these generators; excluded from + // both numerator and denominator rather than counted + // either way. + continue; + } + sampled += 1; + let detected = detect_content_type(&text); + let reached = text.len() >= threshold_for(detected.content_type); + if reached { + reached_total += 1; + match detected.content_type { + ContentType::SourceCode => reached_source_code += 1, + ContentType::PlainText => reached_plain_text += 1, + _ => {} + } + } + } + + let denom = f64::from(sampled.max(1)); + let overall_frac = f64::from(reached_total) / denom; + let source_frac = f64::from(reached_source_code) / denom; + let plain_frac = f64::from(reached_plain_text) / denom; + + println!( + "dispatch-reach over {sampled} sampled cases (of {REACH_SAMPLE_COUNT} drawn): \ + overall {overall_frac:.4} ({reached_total}), \ + source_code {source_frac:.4} ({reached_source_code}), \ + plain_text {plain_frac:.4} ({reached_plain_text})" + ); + + assert!( + source_frac > 0.10, + "SourceCode dispatch-reach fraction regressed: {source_frac:.4} (want > 0.10)" + ); + assert!( + plain_frac > 0.10, + "PlainText dispatch-reach fraction regressed: {plain_frac:.4} (want > 0.10)" + ); + assert!( + overall_frac > 0.25, + "overall dispatch-reach fraction regressed: {overall_frac:.4} (want > 0.25)" + ); +} diff --git a/crates/headroom-core/tests/live_zone_disable_arms.rs b/crates/headroom-core/tests/live_zone_disable_arms.rs new file mode 100644 index 000000000..c9f137e4b --- /dev/null +++ b/crates/headroom-core/tests/live_zone_disable_arms.rs @@ -0,0 +1,140 @@ +//! Env-var kill-switch for the live-zone dispatch arms — dedicated +//! integration test file. +//! +//! `disabled_arms` (`crates/headroom-core/src/transforms/live_zone.rs`) +//! reads `HEADROOM_LIVE_ZONE_DISABLE_ARMS` on first dispatch and latches +//! the parsed set behind a process-global `OnceLock` (the determinism +//! invariant — a run's arm-disable set cannot change mid-flight). +//! Environment variables are process-global too, so setting the variable +//! must happen in a test process that does nothing else: any other test +//! in the same binary that dispatched first would freeze the `OnceLock` +//! against whatever it saw at that point. That is why this file holds +//! exactly ONE `#[test]` fn — a dedicated file is a dedicated test +//! binary, so no other test can race the initialization. +//! +//! Do NOT add more tests to this file. +//! +//! **On `set_var`.** `std::env::set_var` is unsound in a multi-threaded +//! process (it races any concurrent reader of the environment, including +//! ones inside libc) and becomes `unsafe` in edition 2024. It is used +//! here because the switch's input genuinely *is* the environment and +//! this binary is single-threaded at the point of the call. The parsing +//! contract itself — aliases, trimming, unknown tokens — is covered +//! without any environment mutation by +//! `live_zone_disable_arms_parsing.rs` against the pure +//! `parse_disabled_arms`. Threading an explicit config value into the +//! dispatcher would remove the need for this file entirely; see the PR +//! description. + +mod common; + +use common::{ + body_with_tool_result, dispatch, json_array_of_dicts, kompress_available, plain_prose, + python_module_source, tool_result_action, +}; +use headroom_core::transforms::{BlockAction, LiveZoneOutcome}; + +/// Assert the single `tool_result` block was left uncompressed. +fn assert_no_compression(label: &str, body: &[u8]) { + let out = dispatch(body); + let manifest = match &out { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => { + panic!( + "disabled {label} arm must not rewrite bytes; got Modified. manifest: {manifest:?}" + ) + } + }; + let action = tool_result_action(manifest); + assert!( + matches!(action, BlockAction::NoCompressionApplied { .. }), + "disabled {label} arm must yield NoCompressionApplied, got {action:?}" + ); +} + +#[test] +fn disabled_arms_route_to_no_op_others_unaffected() { + // `source_code, plain_text, json_array, bogus_type` — the internal + // spaces exercise the trim path and `bogus_type` the unknown-token + // branch (logged and ignored, never a panic). + std::env::set_var( + "HEADROOM_LIVE_ZONE_DISABLE_ARMS", + "source_code, plain_text, json_array, bogus_type", + ); + + // (a) SourceCode disabled: a >2048-byte Python tool_result must not + // reach the CodeAwareCompressor. Deterministic on every host — the + // code arm has no model-cache dependency. + let code = python_module_source(10); + assert!( + code.len() > 2048, + "fixture must clear the SourceCode byte threshold (2048); got {} bytes", + code.len() + ); + assert_no_compression("SourceCode", &body_with_tool_result(&code).0); + + // (b) JsonArray disabled: SmartCrusher must not fire on a shape it + // would otherwise compress on any host. This is the case that + // discriminates the kill switch unconditionally — unlike (c), it + // cannot be satisfied by an absent model — and it also pins that the + // switch is generic across arms rather than special-cased to the two + // this PR wires. + let payload = json_array_of_dicts(200); + assert!( + payload.len() > 512, + "fixture must clear the JsonArray byte threshold (512); got {} bytes", + payload.len() + ); + assert_no_compression("JsonArray", &body_with_tool_result(&payload).0); + + // (c) PlainText disabled: a >5120-byte prose tool_result must not + // reach Kompress. Two vacuity traps, not one: + // + // 1. The fixture MUST clear THRESHOLD_PLAIN_TEXT (5120) or the + // byte-threshold gate short-circuits before dispatch and the + // assertion passes regardless of the switch. + // 2. On a cold HuggingFace cache the enabled arm's fall-open path + // returns the identical no-op, so this cannot discriminate there. + // `kompress_available()` says which case actually ran; (b) above + // carries the unconditional coverage. + let prose = plain_prose(5_200); + assert!( + prose.len() > 5120, + "fixture must clear the PlainText byte threshold (5120); got {} bytes", + prose.len() + ); + assert_no_compression("PlainText", &body_with_tool_result(&prose).0); + if !kompress_available() { + eprintln!( + "NOTE: cold model cache — the PlainText case passed but could not discriminate \ + the kill switch from an absent model this run. Unconditional coverage comes \ + from the JsonArray case above." + ); + } + + // (d) An arm NOT named in the list still compresses: build output + // routes to LogCompressor as usual, proving the switch is selective + // rather than a global off. + let logs = "\ +ERROR src/main.rs:42 connection refused +WARN src/pool.rs:17 retrying in 250ms +ERROR src/main.rs:42 connection refused +WARN src/pool.rs:17 retrying in 500ms +" + .repeat(40); + let out = dispatch(&body_with_tool_result(&logs).0); + let manifest = match &out { + LiveZoneOutcome::Modified { manifest, .. } => manifest, + LiveZoneOutcome::NoChange { manifest } => panic!( + "an arm absent from the disable list must be unaffected; got NoChange. \ + manifest: {manifest:?}" + ), + }; + match tool_result_action(manifest) { + BlockAction::Compressed { strategy, .. } => assert_eq!( + strategy, "log_compressor", + "expected LogCompressor dispatch, unaffected by the disabled arms" + ), + other => panic!("expected BlockAction::Compressed via log_compressor, got {other:?}"), + } +} diff --git a/crates/headroom-core/tests/live_zone_disable_arms_parsing.rs b/crates/headroom-core/tests/live_zone_disable_arms_parsing.rs new file mode 100644 index 000000000..f972beb1c --- /dev/null +++ b/crates/headroom-core/tests/live_zone_disable_arms_parsing.rs @@ -0,0 +1,64 @@ +//! Parsing contract of the `HEADROOM_LIVE_ZONE_DISABLE_ARMS` kill +//! switch, tested against the pure `parse_disabled_arms` rather than +//! through the process environment. +//! +//! The end-to-end wiring (env var → latched set → dispatcher no-op) +//! lives in `live_zone_disable_arms.rs`, which needs `std::env::set_var` +//! and therefore a dedicated single-test binary. Everything about how a +//! string becomes a set of arms is checkable here instead: no global +//! state, no isolation requirement, and these cases can run alongside +//! anything else. + +use headroom_core::transforms::live_zone::parse_disabled_arms; +use headroom_core::transforms::ContentType; + +#[test] +fn parses_both_spellings_and_trims_whitespace() { + let parsed = parse_disabled_arms("source_code, plain_text"); + assert!(parsed.contains(&ContentType::SourceCode)); + assert!(parsed.contains(&ContentType::PlainText)); + assert_eq!(parsed.len(), 2); + + // `as_str()` tags parse to the same variants as the natural names. + let tags = parse_disabled_arms("source_code,text"); + assert_eq!(parsed, tags, "both spellings must yield the same set"); +} + +#[test] +fn every_variant_can_be_named() { + for content_type in ContentType::ALL { + for spelling in [content_type.as_str(), content_type.natural_name()] { + let parsed = parse_disabled_arms(spelling); + assert!( + parsed.contains(&content_type), + "{spelling:?} must disable {content_type:?}" + ); + } + } +} + +#[test] +fn unknown_and_blank_tokens_are_ignored_not_fatal() { + // A typo in an operator's rollback switch must not take the proxy + // down, and must not silently disable something else either. + let parsed = parse_disabled_arms("bogus_type, , source_code,, "); + assert_eq!( + parsed, + std::iter::once(ContentType::SourceCode).collect(), + "unknown and empty tokens drop out; valid ones survive" + ); +} + +#[test] +fn empty_input_disables_nothing() { + assert!(parse_disabled_arms("").is_empty()); + assert!(parse_disabled_arms(" ").is_empty()); + assert!(parse_disabled_arms(",,,").is_empty()); +} + +#[test] +fn repeated_tokens_collapse() { + let parsed = parse_disabled_arms("plain_text,text,plain_text"); + assert_eq!(parsed.len(), 1); + assert!(parsed.contains(&ContentType::PlainText)); +} diff --git a/crates/headroom-core/tests/live_zone_dispatch.rs b/crates/headroom-core/tests/live_zone_dispatch.rs index c8f55cf00..3f5c84b01 100644 --- a/crates/headroom-core/tests/live_zone_dispatch.rs +++ b/crates/headroom-core/tests/live_zone_dispatch.rs @@ -6,77 +6,17 @@ //! - Build/log output → LogCompressor //! - Search-result tool_results → SearchCompressor //! - Git diff tool_results → DiffCompressor -//! - Source code → no-op (Rust port pending) +//! - Source code → CodeAwareCompressor //! - Unknown / image / html → no-op //! //! Plus the cache-safety invariant: bytes outside the rewritten //! block are byte-identical to the input (SHA-256 prefix + suffix). -use headroom_core::transforms::live_zone::DEFAULT_MODEL; -use headroom_core::transforms::{ - compress_anthropic_live_zone, AuthMode, BlockAction, LiveZoneOutcome, -}; +mod common; + +use common::{body_with_tool_result, dispatch, kompress_available, python_module_source, sha256}; +use headroom_core::transforms::{BlockAction, LiveZoneOutcome}; use serde_json::{json, Value}; -use sha2::{Digest, Sha256}; - -fn body_of(value: Value) -> Vec { - serde_json::to_vec(&value).unwrap() -} - -fn dispatch(body: &[u8]) -> LiveZoneOutcome { - compress_anthropic_live_zone(body, 0, AuthMode::Payg, DEFAULT_MODEL) - .expect("dispatcher returns Ok on valid bodies") -} - -/// Find the byte range of the FIRST occurrence of `needle` inside -/// `haystack`. Used by the byte-fidelity test below to identify the -/// JSON-encoded tool_result.content slot we expect the dispatcher to -/// rewrite. Returns `(start, end)` half-open. -fn find_byte_range(haystack: &[u8], needle: &[u8]) -> (usize, usize) { - let pos = haystack - .windows(needle.len()) - .position(|w| w == needle) - .unwrap_or_else(|| { - panic!( - "needle of {} bytes not found in haystack of {} bytes", - needle.len(), - haystack.len() - ) - }); - (pos, pos + needle.len()) -} - -fn sha256(bytes: &[u8]) -> [u8; 32] { - let mut h = Sha256::new(); - h.update(bytes); - h.finalize().into() -} - -/// Build a body with one user message containing one `tool_result` -/// whose `content` is `text`. Returns the full body and the byte -/// range of the JSON-encoded `content` slot (including the surrounding -/// quotes) within that body — useful for byte-fidelity assertions. -fn body_with_tool_result(text: &str) -> (Vec, (usize, usize)) { - let body = body_of(json!({ - "model": "claude-sonnet-4-6", - "max_tokens": 64, - "system": "you are a helpful assistant", - "messages": [{ - "role": "user", - "content": [{ - "type": "tool_result", - "tool_use_id": "toolu_dispatch_test", - "content": text, - }], - }], - })); - // The JSON-encoded `content` slot is exactly `serde_json::to_vec(&text)`, - // since text is shorter than the whole body and serde uses the same - // encoding for the embedded string. - let needle = serde_json::to_vec(&text).unwrap(); - let range = find_byte_range(&body, &needle); - (body, range) -} // ─── Routing tests ───────────────────────────────────────────────────── @@ -262,29 +202,83 @@ fn diff_tool_result_routes_to_diff_compressor() { } #[test] -fn source_code_tool_result_routes_to_no_op() { - // Detector classifies this as SourceCode. PR-B3 routes it to - // no-op (Rust code-compressor port pending). Pin the contract - // so a future "wire it up" PR can flip this assertion. - let code = " -fn main() { - let x: i32 = 42; - let y = x * 2; - println!(\"{}\", y); - if x > 0 { - println!(\"positive\"); - } else { - println!(\"non-positive\"); +fn source_code_tool_result_routes_to_code_compressor() { + // Detector classifies this as SourceCode. PR-B4 wires the arm up to + // the tree-sitter-backed CodeAwareCompressor. This flips the PR-B3 + // pin ("a future 'wire it up' PR can flip this assertion"). + let code = python_module_source(10); + assert!( + code.len() > 2048, + "fixture must clear the SourceCode byte threshold (2048); got {} bytes", + code.len() + ); + + let (body, _) = body_with_tool_result(&code); + let out = dispatch(&body); + let manifest = match &out { + LiveZoneOutcome::Modified { manifest, .. } => manifest, + LiveZoneOutcome::NoChange { manifest } => panic!( + "expected CodeAwareCompressor to shrink a 10-function Python module; got NoChange. manifest: {manifest:?}" + ), + }; + let action = manifest + .block_outcomes + .iter() + .find(|b| b.block_type == "tool_result") + .expect("tool_result block present") + .action + .clone(); + match action { + BlockAction::Compressed { + strategy, + original_tokens, + compressed_tokens, + .. + } => { + assert_eq!( + strategy, "code_compressor", + "expected code_compressor dispatch" + ); + assert!( + compressed_tokens < original_tokens, + "tokenizer-validated gate (PR-B4) must accept only token-shrinking output \ + ({compressed_tokens} < {original_tokens})" + ); + } + other => panic!("expected BlockAction::Compressed, got {other:?}"), } } -" - .repeat(20); - let (body, _) = body_with_tool_result(&code); + +#[test] +fn tiny_source_code_below_threshold_no_op() { + // Detector classifies this as SourceCode, but it's well under the + // 2048-byte SourceCode threshold, so the dispatcher must not even + // spin up the CodeAwareCompressor. + let code = "\ +import os +from typing import Any + + +def add(a: int, b: int) -> int: + \"\"\"Add two integers.\"\"\" + return a + b + + +if __name__ == \"__main__\": + print(add(1, 2)) +"; + assert!( + code.len() < 2048, + "fixture must stay below the SourceCode byte threshold (2048); got {} bytes", + code.len() + ); + + let (body, _) = body_with_tool_result(code); let out = dispatch(&body); let manifest = match &out { LiveZoneOutcome::NoChange { manifest } => manifest, LiveZoneOutcome::Modified { manifest, .. } => { - panic!("PR-B3 must NOT compress SourceCode (Rust port pending). manifest: {manifest:?}") + panic!("tiny source-code block must not be compressed. manifest: {manifest:?}") } }; let action = manifest @@ -295,26 +289,119 @@ fn main() { .action .clone(); match action { - BlockAction::NoCompressionApplied { content_type } => { - // Source-code-shaped content above the SourceCode byte - // threshold (2 KiB) but below any active compressor: - // SmartCrusher / log / search / diff don't apply, and - // the Rust code-compressor port is not yet wired. - assert!( - content_type == "source_code" || content_type == "text", - "unexpected content_type tag: {content_type}" - ); + BlockAction::BelowByteThreshold { + content_type, + threshold_bytes, + .. + } => { + assert_eq!(threshold_bytes, 2048, "expected the SourceCode threshold"); + assert_eq!(content_type, "source_code", "unexpected content_type tag"); } - BlockAction::BelowByteThreshold { content_type, .. } => { - // Detector may classify code-with-prose as PlainText - // (5 KiB threshold) — for ~2.6 KiB of mixed code/prose - // that still routes to no-op for B4. Pin the tag. - assert!( - content_type == "text" || content_type == "source_code", - "unexpected content_type tag: {content_type}" - ); + other => panic!("expected BelowByteThreshold, got {other:?}"), + } +} + +#[test] +fn plain_text_below_threshold_no_op() { + // Plain prose, well under the 5120-byte PlainText threshold, so the + // dispatcher must not even attempt Kompress. + let prose = "The quarterly report highlighted steady growth across every \ + region, with the operations team noting improved throughput on the \ + warehouse floor and customer support tickets trending downward for \ + the third month running. Leadership expects the trend to continue \ + into next quarter, barring any supply chain disruptions."; + assert!( + prose.len() < 5120, + "fixture must stay below the PlainText byte threshold (5120); got {} bytes", + prose.len() + ); + + let (body, _) = body_with_tool_result(prose); + let out = dispatch(&body); + let manifest = match &out { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => { + panic!("sub-threshold plain text must not be compressed. manifest: {manifest:?}") } - other => panic!("expected NoCompressionApplied or BelowByteThreshold, got {other:?}"), + }; + let action = manifest + .block_outcomes + .iter() + .find(|b| b.block_type == "tool_result") + .expect("tool_result block present") + .action + .clone(); + match action { + BlockAction::BelowByteThreshold { + content_type, + threshold_bytes, + .. + } => { + assert_eq!(threshold_bytes, 5120, "expected the PlainText threshold"); + assert_eq!(content_type, "text", "unexpected content_type tag"); + } + other => panic!("expected BelowByteThreshold, got {other:?}"), + } +} + +#[test] +fn plain_text_routes_to_kompress_when_model_cached() { + // RUNTIME-SKIP: ask the loader itself whether the model is + // cache-resident, rather than re-deriving cache paths here. A + // hand-rolled probe has to be kept byte-compatible with + // `Kompress::from_cache`'s own root and artifact resolution, and + // when it drifts the skip becomes a lie — the test silently stops + // running on hosts where production would load the model. + if !kompress_available() { + eprintln!( + "SKIP: kompress model/tokenizer not cache-resident; run `python scripts/record_kompress_trace.py` first" + ); + return; + } + + // Repetitive news-article-like prose: > 350 words so Kompress's + // chunk_words=350 chunking actually engages, and > 5120 bytes to + // clear the PlainText byte threshold. + let mut article = String::new(); + for i in 0..60 { + article.push_str(&format!( + "City officials announced today that the downtown revitalization \ + project will proceed as planned despite budget concerns raised \ + during round {i} of public comment. " + )); + } + assert!( + article.split_whitespace().count() > 350, + "fixture must exceed 350 words so Kompress chunking engages; got {} words", + article.split_whitespace().count() + ); + assert!( + article.len() > 5120, + "fixture must clear the PlainText byte threshold (5120); got {} bytes", + article.len() + ); + + let (body, _) = body_with_tool_result(&article); + let out = dispatch(&body); + let manifest = match &out { + LiveZoneOutcome::Modified { manifest, .. } => manifest, + LiveZoneOutcome::NoChange { manifest } => panic!( + "expected Kompress to compress a {}-word repetitive article; got NoChange. manifest: {manifest:?}", + article.split_whitespace().count() + ), + }; + let action = manifest + .block_outcomes + .iter() + .find(|b| b.block_type == "tool_result") + .expect("tool_result block present") + .action + .clone(); + match action { + BlockAction::Compressed { strategy, .. } => { + assert_eq!(strategy, "kompress", "expected kompress dispatch"); + } + other => panic!("expected BlockAction::Compressed via kompress, got {other:?}"), } } diff --git a/crates/headroom-core/tests/live_zone_kompress_absent.rs b/crates/headroom-core/tests/live_zone_kompress_absent.rs new file mode 100644 index 000000000..0235efd81 --- /dev/null +++ b/crates/headroom-core/tests/live_zone_kompress_absent.rs @@ -0,0 +1,80 @@ +//! Model-absent NoOp — dedicated integration test file. +//! +//! The Kompress slot (`crates/headroom-core/src/transforms/live_zone.rs`) +//! latches process-globally: the first PlainText dispatch starts its +//! background initialization, whose cache lookup reads `HF_HUB_CACHE` / +//! `HF_HOME` / `HOME` / `USERPROFILE`. Environment variables are +//! process-global too, so forcing those four to a cold, empty cache dir +//! must happen in a test process that does nothing else — any other test +//! in the same binary that dispatched PlainText content first would +//! settle the slot against the *ambient* environment instead. That is +//! why this file holds exactly ONE `#[test]` fn. +//! +//! Do NOT add more tests to this file, and do NOT move this test into +//! `live_zone_dispatch.rs` (see +//! `plain_text_routes_to_kompress_when_model_cached` there, which relies +//! on the loader observing the *ambient* cache). +//! +//! On `set_var`: see the note in `live_zone_disable_arms.rs`. The same +//! caveat applies, for the same reason — the input under test is the +//! environment. An injectable cache root on `Kompress::from_cache` would +//! remove the need for this file; see the PR description. + +mod common; + +use common::{body_with_tool_result, dispatch, plain_prose, tool_result_action}; +use headroom_core::transforms::{BlockAction, LiveZoneOutcome}; + +#[test] +fn plain_text_model_absent_is_deterministic_no_op() { + // Force every cache root the loader consults to a fresh, empty temp + // dir so the lookup deterministically misses, rather than picking up + // whatever happens to be in the real user cache on this machine. + let cold_dir = tempfile::tempdir().expect("create fresh temp dir for cold model cache"); + let cold_path = cold_dir.path().to_str().expect("temp dir path is UTF-8"); + for var in ["HF_HUB_CACHE", "HF_HOME", "HOME", "USERPROFILE"] { + std::env::set_var(var, cold_path); + } + + // > 5120 bytes so the PlainText byte threshold is cleared and the + // dispatcher actually attempts the arm rather than short-circuiting + // at `BelowByteThreshold`. + let prose = plain_prose(5_200); + assert!( + prose.len() > 5120, + "fixture must clear the PlainText byte threshold (5120); got {} bytes", + prose.len() + ); + + let body = body_with_tool_result(&prose).0; + let assert_no_op = |label: &str| { + let out = dispatch(&body); + let manifest = match &out { + LiveZoneOutcome::NoChange { manifest } => manifest, + LiveZoneOutcome::Modified { manifest, .. } => panic!( + "cache-cold Kompress must not rewrite bytes ({label}); got Modified. \ + manifest: {manifest:?}" + ), + }; + let action = tool_result_action(manifest); + assert!( + matches!(action, BlockAction::NoCompressionApplied { .. }), + "cache-cold Kompress must degrade to a deterministic NoOp ({label}), \ + not an error: {action:?}" + ); + }; + + // First dispatch: a NoOp by the non-blocking contract — this call is + // what starts the slot's background initialization. + assert_no_op("first dispatch, slot unsettled"); + + // Settle the slot off the request path and pin WHY it is empty: the + // loader found nothing under the forced-cold roots. + assert!( + !headroom_core::transforms::live_zone::warm_live_zone_compressors(), + "the loader reported a model ready under a forced-cold cache root" + ); + + // A post-settle dispatch is the same deterministic NoOp. + assert_no_op("post-settle dispatch"); +} diff --git a/crates/headroom-core/tests/live_zone_kompress_async_init.rs b/crates/headroom-core/tests/live_zone_kompress_async_init.rs new file mode 100644 index 000000000..afc2fef5a --- /dev/null +++ b/crates/headroom-core/tests/live_zone_kompress_async_init.rs @@ -0,0 +1,116 @@ +//! Non-blocking Kompress initialization — dedicated integration test file. +//! +//! The Kompress slot in `live_zone.rs` is process-global and latches on +//! first use, so proving anything about its *virgin* state needs a test +//! process that does nothing else first — the same reason +//! `live_zone_disable_arms.rs` and `live_zone_kompress_absent.rs` are +//! single-test binaries. Do NOT add more tests to this file. +//! +//! Contract under test (#3227 review): activating the PlainText arm +//! because model artifacts happen to be cache-resident must not make the +//! first qualifying request pay — or wait on — the ~261 MB ONNX session +//! build. Dispatches that race a virgin slot return promptly as NoOp; +//! the expensive construction runs exactly once, off the request path. +//! +//! Two vacuity notes: +//! +//! 1. The timing assertion discriminates only on a warm HuggingFace +//! cache, where the synchronous build costs whole seconds; on a cold +//! cache the loader returns `None` quickly and even a blocking +//! implementation passes the bound. The exactly-once and NoOp +//! assertions hold on any host. +//! 2. The NoOp assertion cannot be raced into `Modified` by a straggler +//! thread observing a *completed* init: every thread dispatches +//! barrier-synchronized within microseconds of the CAS that starts +//! the initializer, and the build it would have to lose to costs +//! seconds (tokenizer load + 261 MB session commit). + +mod common; + +use std::sync::{Arc, Barrier}; +use std::time::{Duration, Instant}; + +use common::{body_with_tool_result, dispatch, plain_prose, tool_result_action}; +use headroom_core::transforms::live_zone; +use headroom_core::transforms::{BlockAction, LiveZoneOutcome}; + +#[test] +fn first_dispatch_never_blocks_on_model_init() { + // > 5120 bytes so the PlainText byte threshold is cleared and the + // dispatcher actually attempts the arm rather than short-circuiting + // at `BelowByteThreshold`. + let prose = plain_prose(5_200); + assert!( + prose.len() > 5120, + "fixture must clear the PlainText byte threshold (5120); got {} bytes", + prose.len() + ); + let (body, _) = body_with_tool_result(&prose); + + // Phase 1 — virgin slot, concurrent first requests. Every dispatch + // must return promptly (never waiting on a model build another + // thread started) and must leave the body untouched. + let n_threads = 8; + let barrier = Arc::new(Barrier::new(n_threads)); + let handles: Vec<_> = (0..n_threads) + .map(|_| { + let body = body.clone(); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + let started = Instant::now(); + let out = dispatch(&body); + (started.elapsed(), out) + }) + }) + .collect(); + + for (i, handle) in handles.into_iter().enumerate() { + let (elapsed, out) = handle.join().expect("dispatch thread panicked"); + assert!( + elapsed < Duration::from_millis(500), + "thread {i}: dispatch must not wait on model initialization; took {elapsed:?}" + ); + assert!( + matches!(out, LiveZoneOutcome::NoChange { .. }), + "thread {i}: a dispatch racing a virgin model slot must be a NoOp \ + while initialization is in flight, not a rewrite" + ); + } + + // Phase 2 — however many threads raced the virgin slot, the + // expensive construction ran exactly once; a blocking warmup (the + // proxy's startup path) settles the slot and reports readiness. + let ready = live_zone::warm_live_zone_compressors(); + let expected_runs = usize::from(cfg!(feature = "ml")); + assert_eq!( + live_zone::kompress_init_runs(), + expected_runs, + "the model construction must run exactly once per process" + ); + + // Phase 3 — with the slot settled, the arm compresses iff the model + // actually loaded; readiness comes from the warmup itself, so this + // cannot silently skip on a host where production would compress. + if ready { + match dispatch(&body) { + LiveZoneOutcome::Modified { manifest, .. } => match tool_result_action(&manifest) { + BlockAction::Compressed { strategy, .. } => assert_eq!( + strategy, "kompress", + "post-warmup PlainText dispatch must compress via Kompress" + ), + other => panic!("expected Compressed via kompress, got {other:?}"), + }, + LiveZoneOutcome::NoChange { manifest } => panic!( + "warmup reported the model ready, so a settled dispatch must compress \ + this above-threshold prose block; manifest: {manifest:?}" + ), + } + } else { + eprintln!( + "NOTE: cold model cache (or a no-ml build) — the non-blocking and \ + exactly-once contracts above are covered; the post-warmup compression \ + check could not run on this host." + ); + } +} diff --git a/crates/headroom-core/tests/live_zone_thresholds.rs b/crates/headroom-core/tests/live_zone_thresholds.rs index 429db3f29..c735ed2ca 100644 --- a/crates/headroom-core/tests/live_zone_thresholds.rs +++ b/crates/headroom-core/tests/live_zone_thresholds.rs @@ -1,8 +1,9 @@ //! PR-B4 byte-threshold gate — integration tests. //! //! The dispatcher must skip compression entirely for blocks whose -//! content is below the per-content-type byte threshold (1 KiB for -//! JSON arrays). PR-B4 spec, `REALIGNMENT/04-phase-B-live-zone.md`. +//! content is below the per-content-type byte threshold (512 B for +//! JSON arrays — asserted below). PR-B4 spec, +//! `REALIGNMENT/04-phase-B-live-zone.md`. use headroom_core::transforms::live_zone::DEFAULT_MODEL; use headroom_core::transforms::{ @@ -88,7 +89,7 @@ fn below_threshold_no_compression_attempted() { #[test] fn above_threshold_compression_attempted() { - // 10 KB of homogeneous JSON dicts — comfortably above the 1 KiB + // 10 KB of homogeneous JSON dicts — comfortably above the 512 B // JsonArray threshold and SmartCrusher's bread-and-butter shape, // so the dispatcher SHOULD route through `dispatch_compressor`. // Either `Compressed` (SmartCrusher shrunk it — the typical diff --git a/crates/headroom-proxy/src/config.rs b/crates/headroom-proxy/src/config.rs index 181589754..0ea8703a2 100644 --- a/crates/headroom-proxy/src/config.rs +++ b/crates/headroom-proxy/src/config.rs @@ -10,26 +10,26 @@ use url::Url; /// Compression mode policy for the `/v1/messages` endpoint. /// -/// Drives whether `compress_anthropic_request` does any work. PR-A1 -/// (Phase A lockdown) wires the flag in but both modes currently -/// passthrough — `live_zone` parses-but-warns until Phase B PR-B2 -/// fills in the live-zone-only block dispatcher. +/// Drives whether `compress_anthropic_request` does any work. `off` is +/// byte-faithful passthrough; `live_zone` routes the request through +/// the headroom-core live-zone dispatcher, which compresses only the +/// live-zone blocks (latest user message, latest tool/function/shell/ +/// patch outputs) via the per-content-type compressor table. /// /// We do NOT add an `icm` mode (the deleted code path) or a /// `passthrough` alias for `off` — those names are misleading. The /// only legal values are `off` (compression disabled) and `live_zone` -/// (compress only the live-zone blocks; not yet implemented). +/// (compress only the live-zone blocks). #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] #[clap(rename_all = "snake_case")] pub enum CompressionMode { /// Compression disabled. Body forwards byte-equal to upstream. - /// This is the default; Phase B will switch the default to - /// `live_zone` once that mode is implemented. + /// This is the default. Off, /// Compress only live-zone blocks (latest user message, - /// latest tool/function/shell/patch outputs). NOT YET IMPLEMENTED: - /// in PR-A1 this falls through to passthrough behaviour with a - /// loud warning. Phase B PR-B2 wires in the actual dispatcher. + /// latest tool/function/shell/patch outputs) via the + /// headroom-core live-zone dispatcher's per-content-type + /// compressors. LiveZone, } @@ -122,39 +122,35 @@ impl CacheControlAutoFrozen { } } -/// Phase F PR-F2.1 c3/6: feature flag for the per-auth-mode +/// Phase F PR-F2.1: feature flag for the per-auth-mode /// `CompressionPolicy` enforcement. /// -/// `disabled` (default until c6/6): the proxy still classifies -/// `auth_mode` and derives a `CompressionPolicy` for telemetry, but -/// every dispatcher and transform behaves as if the mode were `Payg` -/// — bit-for-bit current behaviour. +/// `enabled` (default, from c5/5 onward): the policy struct's +/// per-mode values take effect. For Subscription specifically, the +/// cache aligner is skipped and the dispatcher gates on +/// `policy.live_zone_compression_enabled()` (a no-op in F2.1 since +/// that helper currently always returns `true`, but kept as a hook +/// so F2.2 can flip without touching call sites). /// -/// `enabled`: the policy struct's per-mode values take effect. For -/// Subscription specifically, the cache aligner is skipped and the -/// dispatcher gates on `policy.live_zone_compression_enabled()` (a -/// no-op in F2.1 since that helper currently always returns `true`, -/// but kept as a hook so F2.2 can flip without touching call sites). -/// -/// Why a flag at all: F2.1 lands behind a default-disabled gate so -/// commits 4 and 5 of the PR don't ship behaviour change to default -/// users. Operators can flip this on for dogfooding before commit 6 -/// flips the default. Rollback: flip the env var back to `disabled` -/// — instant if config is hot-reloaded, redeploy otherwise. +/// `disabled`: the proxy still classifies `auth_mode` and derives a +/// `CompressionPolicy` for telemetry, but every dispatcher and +/// transform behaves as if the mode were `Payg` — bit-for-bit the +/// pre-F2.1 behaviour. Operators can flip back to this for rollback +/// if F2.1 surfaces a subscription regression — instant if config is +/// hot-reloaded, redeploy otherwise. /// /// Source priority: CLI flag → /// `HEADROOM_PROXY_AUTH_MODE_POLICY_ENFORCEMENT` env var → -/// default (`disabled`). +/// default (`enabled`). #[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] #[clap(rename_all = "snake_case")] pub enum AuthModePolicyEnforcement { /// Per-mode policy IS enforced. Subscription users see no /// cache_aligner; the dispatcher reads - /// `policy.live_zone_compression_enabled()`. + /// `policy.live_zone_compression_enabled()`. Default. Enabled, /// Per-mode policy IS NOT enforced. Every mode runs the PAYG - /// pipeline, identical to pre-F2.1 behaviour. Default in F2.1 - /// commits 1–5 so the feature is dogfood-only until c6/6. + /// pipeline, identical to pre-F2.1 behaviour. Rollback opt-out. Disabled, } @@ -340,11 +336,9 @@ pub struct CliArgs { /// Compression mode policy for `/v1/messages`. /// /// `off` (default): byte-faithful passthrough on every request. - /// `live_zone`: PR-B2 wired the dispatcher; PR-B2's per-type - /// compressors are no-ops, so the body still round-trips - /// byte-equal until PR-B3+ (which fills the per-type table). - /// The flag exists so the default can flip in one config - /// change once `live_zone` is the safer choice on real traffic. + /// `live_zone`: routes the request through the headroom-core + /// live-zone dispatcher, which compresses eligible live-zone + /// blocks via its per-content-type compressor table. /// /// Source priority: CLI flag → `HEADROOM_PROXY_COMPRESSION_MODE` /// env var → default (`off`). @@ -637,19 +631,19 @@ pub struct Config { /// Inherits `max_body_bytes` when not overridden. Bodies larger /// than this still forward, just unchanged. pub compression_max_body_bytes: u64, - /// Policy mode for compression on `/v1/messages`. PR-A1 lockdown: - /// both `Off` and `LiveZone` result in byte-faithful passthrough; - /// `LiveZone` additionally emits a `tracing::warn!` per request - /// because the dispatcher isn't implemented yet (Phase B PR-B2 - /// fills this in). + /// Policy mode for compression on `/v1/messages`. `Off` is + /// byte-faithful passthrough; `LiveZone` routes the request + /// through the headroom-core live-zone dispatcher, which + /// compresses eligible live-zone blocks via its per-content-type + /// compressor table. pub compression_mode: CompressionMode, /// Whether the live-zone dispatcher derives `frozen_message_count` /// automatically from customer `cache_control` markers. PR-A4 /// adds the derivation function (`compute_frozen_count`); Phase /// B's dispatcher consumes the resolved value here. pub cache_control_auto_frozen: CacheControlAutoFrozen, - /// Phase F PR-F2.1 c3/6: gate per-auth-mode `CompressionPolicy` - /// enforcement. `Disabled` until c6/6 flips the default. + /// Phase F PR-F2.1: gate per-auth-mode `CompressionPolicy` + /// enforcement. `Enabled` by default (from c5/5 onward). pub auth_mode_policy_enforcement: AuthModePolicyEnforcement, /// Whether to strip internal `x-headroom-*` headers from /// upstream-bound requests. PR-A5 default-on guard against diff --git a/crates/headroom-proxy/src/main.rs b/crates/headroom-proxy/src/main.rs index f21d0117f..3c7954fa2 100644 --- a/crates/headroom-proxy/src/main.rs +++ b/crates/headroom-proxy/src/main.rs @@ -54,6 +54,24 @@ async fn main() -> Result<(), Box> { ); } + // The ~261 MB Kompress ONNX session must never build on the request + // path: dispatch answers NoOp until the model is ready. Warm it in + // the background at startup so steady-state traffic meets a settled + // slot; a byte-pipe proxy (--compression off) never loads a model. + if config.compression { + if let Err(e) = std::thread::Builder::new() + .name("live-zone-warmup".into()) + .spawn(|| { + let ready = headroom_core::transforms::live_zone::warm_live_zone_compressors(); + tracing::info!(event = "live_zone_warmup_complete", kompress_ready = ready); + }) + { + // Non-fatal: the dispatcher's lazy path still initializes + // off-request; only the eager warmup is lost. + tracing::warn!(event = "live_zone_warmup_spawn_failed", error = %e); + } + } + let mut state = AppState::new(config.clone())?; // PR-D1: resolve AWS credentials at startup via the `aws-config` diff --git a/docs/content/docs/configuration.mdx b/docs/content/docs/configuration.mdx index fbcc6b297..94095dc91 100644 --- a/docs/content/docs/configuration.mdx +++ b/docs/content/docs/configuration.mdx @@ -351,6 +351,7 @@ headroom proxy --learn --min-evidence 3 | `HEADROOM_BETA_HEADER_STICKY` | Controls per-session `anthropic-beta` / `OpenAI-Beta` re-echo. `enabled` (default): the proxy unions beta tokens across turns within a session — if the client sends a token in turn N and omits it in turn N+1, the proxy re-injects it to preserve prefix-cache stability. `disabled`: the client's value is forwarded verbatim with no accumulation. Any other value raises at request time. See [Session Beta Header Tracking](/docs/configuration#session-beta-header-tracking). | `enabled` | | `HEADROOM_BETA_TRACKER_MAX_SESSIONS` | LRU capacity of the in-memory session beta tracker. Once full, the oldest session entry is evicted. | `1000` | | `HEADROOM_PROXY_BETA_HEADER_STICKY` | Rust proxy: same per-conversation beta-token union as `HEADROOM_BETA_HEADER_STICKY`, applied to `anthropic-beta` / `openai-beta` on the intercepted `/v1/messages`, `/v1/chat/completions`, and `/v1/responses` routes. Requires the compression interceptor (`HEADROOM_PROXY_COMPRESSION=1`) — with it off the Rust proxy is a strict byte-pipe and this flag has no effect (startup warns). Unlike the Python tracker (keyed on model + system prompt), sessions are keyed per conversation, shared with the cache-drift detector — parallel conversations never inherit each other's tokens. `enabled` default; `disabled` forwards the client value verbatim and keeps no state. Tracker capacity is fixed at 1000 sessions. | `enabled` | +| `HEADROOM_LIVE_ZONE_DISABLE_ARMS` | Rust proxy live-zone dispatcher: comma-separated content types whose compression arm is forced to a no-op, e.g. `source_code,plain_text`. Names are accepted in either the log-tag spelling (`json_array`, `source_code`, `search`, `build`, `diff`, `html`, `text`) or the natural spelling (`search_results`, `build_output`, `git_diff`, `plain_text`). Read once, on the first dispatch of the process, so a running proxy's behaviour cannot change mid-flight; unknown names are logged and ignored rather than failing the request. Shipped as the rollback switch for the `SourceCode` → CodeAwareCompressor and `PlainText` → Kompress arms, but applies to every arm. | -- | | `HEADROOM_MODEL_ROUTER_ENABLED` | Enable cost-aware model routing. `1`/`true`/`yes`/`on`/`enabled` turns it on and requires `HEADROOM_MODEL_ROUTES`. See [Cost-aware model routing](/docs/configuration#cost-aware-model-routing). | `off` | | `HEADROOM_MODEL_ROUTES` | JSON array of ordered routing rules for cost-aware model routing (schema below). | -- | | `HEADROOM_THINKING_COMPACT` | Compact plain-text reasoning that models re-send every turn (Kimi/GLM/DeepSeek `reasoning_content` / inline ``): Kompress it on warm turns, drop it on cold turns. No-op for Claude/Codex/OpenAI (encrypted reasoning). See [Cold-prefix hook](#cold-prefix-hook--reasoning-compaction). | `off` |