mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Merge 442301b121 into 8884d87378
This commit is contained in:
r3bb1t
GitHub
commit
675db1b9f4
15 changed files with 2318 additions and 178 deletions
|
|
@ -203,6 +203,10 @@ tempfile = "3"
|
|||
name = "tokenizer"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "live_zone_dispatch"
|
||||
harness = false
|
||||
|
||||
[[bench]]
|
||||
name = "ccr_store"
|
||||
harness = false
|
||||
|
|
|
|||
248
crates/headroom-core/benches/live_zone_dispatch.rs
Normal file
248
crates/headroom-core/benches/live_zone_dispatch.rs
Normal file
|
|
@ -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<u8> {
|
||||
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<u8>)] = &[
|
||||
(
|
||||
"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();
|
||||
}
|
||||
|
|
@ -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.<NAME>.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, Self::Err> {
|
||||
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/
|
||||
|
|
|
|||
|
|
@ -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<CodeAwareCompressor> = 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<Option<Kompress>> = 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<ContentType> {
|
||||
let mut set = HashSet::new();
|
||||
for token in raw.split(',') {
|
||||
let token = token.trim();
|
||||
if token.is_empty() {
|
||||
continue;
|
||||
}
|
||||
match token.parse::<ContentType>() {
|
||||
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<ContentType> {
|
||||
static INSTANCE: OnceLock<HashSet<ContentType>> = 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 {
|
||||
|
|
|
|||
171
crates/headroom-core/tests/common/mod.rs
Normal file
171
crates/headroom-core/tests/common/mod.rs
Normal file
|
|
@ -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<u8> {
|
||||
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<u8>, (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<Value> = (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()
|
||||
}
|
||||
823
crates/headroom-core/tests/live_zone_arm_properties.rs
Normal file
823
crates/headroom-core/tests/live_zone_arm_properties.rs
Normal file
|
|
@ -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::<ContentType>(),
|
||||
Ok(content_type),
|
||||
"as_str() tag {tag:?} must round-trip back to {content_type:?}"
|
||||
);
|
||||
|
||||
let natural = content_type.natural_name();
|
||||
assert_eq!(
|
||||
natural.parse::<ContentType>(),
|
||||
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::<ContentType>().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<Value = String> {
|
||||
(
|
||||
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<Value = String> {
|
||||
(
|
||||
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::<String>()` 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<Value = String> {
|
||||
prop_oneof![
|
||||
2 => any::<String>(),
|
||||
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::<u16>(), 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)"
|
||||
);
|
||||
}
|
||||
140
crates/headroom-core/tests/live_zone_disable_arms.rs
Normal file
140
crates/headroom-core/tests/live_zone_disable_arms.rs
Normal file
|
|
@ -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:?}"),
|
||||
}
|
||||
}
|
||||
64
crates/headroom-core/tests/live_zone_disable_arms_parsing.rs
Normal file
64
crates/headroom-core/tests/live_zone_disable_arms_parsing.rs
Normal file
|
|
@ -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));
|
||||
}
|
||||
|
|
@ -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<u8> {
|
||||
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<u8>, (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:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
80
crates/headroom-core/tests/live_zone_kompress_absent.rs
Normal file
80
crates/headroom-core/tests/live_zone_kompress_absent.rs
Normal file
|
|
@ -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");
|
||||
}
|
||||
116
crates/headroom-core/tests/live_zone_kompress_async_init.rs
Normal file
116
crates/headroom-core/tests/live_zone_kompress_async_init.rs
Normal file
|
|
@ -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."
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -54,6 +54,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
|
|||
);
|
||||
}
|
||||
|
||||
// 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`
|
||||
|
|
|
|||
|
|
@ -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 `<think>`): Kompress it on warm turns, drop it on cold turns. No-op for Claude/Codex/OpenAI (encrypted reasoning). See [Cold-prefix hook](#cold-prefix-hook--reasoning-compaction). | `off` |
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue