perf(compression): take large cold-start contexts off the synchronous kompress path (#1171) (#1298)

## Description

On a cold-start large context, kompress (ModernBERT ONNX) runs
**synchronously on the request thread** — ~200–300s for ~1M tokens. It
blows the 30s compression budget, leaks a non-preemptible worker, and
cascades (executor saturation → queue timeouts on healthy requests); on
timeout the request is forwarded **uncompressed** after eating 30s. This
adds four layered, **default-off, fail-open** mitigations so the request
path is never blocked on ML compression.

Closes #1171

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- **Phase 0 — size gate** (`HEADROOM_KOMPRESS_MAX_TOKENS`, default
50000): route oversized text away from ModernBERT (→ LogCompressor /
TextCrusher / passthrough) at the single `_try_ml_compressor` boundary.
- **Phase 1 — cooperative deadline**
(`HEADROOM_COMPRESSION_DEADLINE_MS`, default 20000): any kompress run
self-terminates at the next chunk boundary past the budget, keeping the
unprocessed tail verbatim.
- **Phase 2 — TextCrusher** (`HEADROOM_TEXT_CRUSHER`): a new **native
Rust** extractive prose compressor in
`crates/headroom-core/src/transforms/text_crusher/`, exposed via PyO3 as
`headroom._core.TextCrusher` with a thin Python wrapper. It **reuses the
shared `crate::relevance::BM25Scorer`** rather than reimplementing BM25,
and ships record/replay parity fixtures (mirroring the SmartCrusher
Rust-core + Python-shim pattern).
- **Phase 3 — off-path compression**
(`HEADROOM_BACKGROUND_COMPRESSION`): forward uncompressed immediately
and compress in a per-process background drain; a byte-identical cache
hit on a later turn means the request never blocks on ML.
- Benchmark (`benchmarks/text_crusher_quality_eval.py`), CHANGELOG
entry, and docstrings documenting the fail-open limits.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`, new modules)
- [x] New tests added for new functionality
- [ ] Manual testing performed (Phase 0/1 gate-fire + deadline observed
on real traffic in earlier iterations; Phase 3 off-path is unit- +
byte-identity-tested, not yet live-validated)

### Test Output

```text
$ pytest tests/test_transforms/ tests/test_cache/ \
    tests/test_proxy/test_background_compression.py tests/test_proxy/test_phase3_byte_identity.py -q
501 passed, 37 skipped in 40.33s

$ cargo test -p headroom-core --lib text_crusher
test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 834 filtered out

$ ruff check <changed files>
All checks passed!

$ mypy headroom/proxy/background_compression.py headroom/transforms/text_crusher.py
Success: no issues found in 2 source files
```

New coverage: size-gate incl. the strategy-dispatch funnel (KOMPRESS +
TEXT); partial-run deadline (chunk-0 compressed + chunk-1 verbatim
tail); BackgroundCompressor (dedup / queue-full / fail-open); Phase 3
byte-identity round-trip; TextCrusher unit + parity.

## Real Behavior Proof

- Environment: macOS, local dev — `uv` venv, Rust `_core` built via `uv
pip install -e .`.
- Exact command / steps: the `pytest` / `cargo test` / `ruff` / `mypy`
commands shown under Test Output; quality eval `python
benchmarks/text_crusher_quality_eval.py /tmp/squad_dev.json`.
- Observed result: 501 Python + 3 Rust tests pass; ruff + mypy clean on
changed/new modules. Quality eval: TextCrusher keeps ~94% of buried
SQuAD answers at 30% size vs ~36% truncate/random; self-contained speed
run ~333k words in ~76ms (one O(n) pass) — sub-second where ModernBERT
takes minutes (fast-vs-slow contrast, not a same-input run).
- Not tested: Phase 3 off-path on live traffic; multi-worker
(per-process by design — see Additional Notes).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

- **All four features are off by default and fail-open** — with the env
flags unset the paths are no-ops for realistic inputs; on any error the
request is forwarded (compressed if possible, else verbatim), never
dropped. A full background queue / duplicate key surfaces as
`deferred:dropped`.
- **Known limits (documented in `background_compression.py`):** Phase 3
is per-process, in-memory, and token-mode-only — these are
**lost-savings, never lost-correctness**, and consistent with the
project's existing per-process compression cache + sticky-session
multi-worker model. The startup multi-worker warning now names off-path
background compression.
- Phase 2 reuses the existing BM25 scorer; reuse did not improve
answer-retention over a Python prototype (query-awareness dominates) —
its value is the Rust speed + repo-conventional Rust-core/Python-shim
shape.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Zhenjia ZHOU 2026-06-23 23:48:06 +08:00 committed by GitHub
parent c7f75b27e9
commit 6c68ff4e9f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
28 changed files with 1855 additions and 14 deletions

16
.codegraph/.gitignore vendored Normal file
View file

@ -0,0 +1,16 @@
# CodeGraph data files
# These are local to each machine and should not be committed
# Database
*.db
*.db-wal
*.db-shm
# Cache
cache/
# Logs
*.log
# Hook markers
.dirty

View file

@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
* **learn:** weight loops in `headroom learn`. A new loop detector (`headroom/learn/loops.py`) recognizes repeated tool-call patterns — including RTK re-fetch loops, where RTK's output truncation makes the agent re-run larger-limit variants of a *successful* command — collapses output-limit variants to one signature, measures the wasted tokens, surfaces loops as a highest-priority digest section, and weights loop guardrails above one-off rules by their measured waste. Previously loops had no special weight and a no-failure re-fetch loop was skipped entirely. Adds an RTK-loop eval (`benchmarks/rtk_loop_learn_eval.py`) that reproduces a loop, runs it through Learn, and asserts the generated guardrail ranks first and prevents re-triggering. * **learn:** weight loops in `headroom learn`. A new loop detector (`headroom/learn/loops.py`) recognizes repeated tool-call patterns — including RTK re-fetch loops, where RTK's output truncation makes the agent re-run larger-limit variants of a *successful* command — collapses output-limit variants to one signature, measures the wasted tokens, surfaces loops as a highest-priority digest section, and weights loop guardrails above one-off rules by their measured waste. Previously loops had no special weight and a no-failure re-fetch loop was skipped entirely. Adds an RTK-loop eval (`benchmarks/rtk_loop_learn_eval.py`) that reproduces a loop, runs it through Learn, and asserts the generated guardrail ranks first and prevents re-triggering.
* **learn:** write per-project learnings to the personal, gitignored `CLAUDE.local.md` by default instead of the team-shared `CLAUDE.md`, matching Claude Code's memory convention so machine-specific paths and tool-discovery byproducts no longer pollute the shared file. Adds a `--target` flag to override the destination (e.g. `--target CLAUDE.md` to opt back into the shared file, or any custom path), and auto-migrates a stale learned-patterns block out of an existing `CLAUDE.md` into `CLAUDE.local.md` with a warning ([#1072](https://github.com/chopratejas/headroom/issues/1072)). * **learn:** write per-project learnings to the personal, gitignored `CLAUDE.local.md` by default instead of the team-shared `CLAUDE.md`, matching Claude Code's memory convention so machine-specific paths and tool-discovery byproducts no longer pollute the shared file. Adds a `--target` flag to override the destination (e.g. `--target CLAUDE.md` to opt back into the shared file, or any custom path), and auto-migrates a stale learned-patterns block out of an existing `CLAUDE.md` into `CLAUDE.local.md` with a warning ([#1072](https://github.com/chopratejas/headroom/issues/1072)).
* **proxy/transforms:** take large cold-start contexts off the synchronous kompress path — the root cause behind the `compression_first_stage` 30s-timeout + leaked-thread → executor-saturation cascade ([#1171](https://github.com/chopratejas/headroom/issues/1171)). A token size-gate inside the ML boundary routes oversized text away from ModernBERT (`HEADROOM_KOMPRESS_MAX_TOKENS`); a cooperative chunk-deadline bounds any kompress run that does proceed (`HEADROOM_COMPRESSION_DEADLINE_MS`); an opt-in off-path mode forwards uncompressed immediately and compresses in a single per-process background drain so the request never blocks on ML (`HEADROOM_BACKGROUND_COMPRESSION`); and a new native `TextCrusher` — a fast deterministic extractive prose compressor in `headroom._core` that reuses the shared BM25 relevance scorer — is the fast alternative to ModernBERT for large plain text (`HEADROOM_TEXT_CRUSHER`). All default off and fail-open. On a SQuAD answer-retention eval (requires the SQuAD dev set) TextCrusher keeps ~94% of buried answers at 30% size vs ~36% for truncate/random, and runs in one O(n) pass -- sub-second where ModernBERT takes minutes (self-contained speed benchmark in `benchmarks/text_crusher_quality_eval.py`).
* **proxy:** measure and surface rolling and current token throughput metrics (active/wall-clock input, compression, effective forward, and streamed generation) in `headroom perf` CLI and the dashboard ([#959](https://github.com/chopratejas/headroom/issues/959)). * **proxy:** measure and surface rolling and current token throughput metrics (active/wall-clock input, compression, effective forward, and streamed generation) in `headroom perf` CLI and the dashboard ([#959](https://github.com/chopratejas/headroom/issues/959)).
* **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`. * **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`.
* **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/<name>` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table. * **proxy:** per-project savings breakdown on the dashboard for all wrapped agents — Claude Code, Codex, aider, Copilot, and Cursor ([#802](https://github.com/chopratejas/headroom/issues/802)). `headroom wrap claude`/`codex` tag requests with an `X-Headroom-Project` header (launch-directory name); `wrap aider`/`copilot`/`cursor` — whose clients cannot send custom headers — use a `/p/<name>` base-URL prefix the proxy strips. Savings are aggregated per project (persisted, schema v3 with transparent v2 migration), exposed as `savings.per_project` in `/stats` and `projects` in `/stats-history`, and shown in a Per-Project Savings dashboard table.

View file

@ -0,0 +1,221 @@
#!/usr/bin/env python3
"""Quality eval for TextCrusher (Phase 2, #1171): does extractive compression
preserve the answer-bearing content? No LLM/API calls -- fully local.
Part A -- SQuAD answer-retention (the strong, labeled metric): bury a real QA
answer in a haystack of distractor paragraphs, compress to a target ratio, and
measure whether the gold answer SURVIVES. TextCrusher (query-aware) vs truncate
(keep-recent) vs random baselines. Mirrors kompress's published
must_keep_recall (0.977 on its own labeled set).
Part B -- real-transcript fidelity: compress large text blocks from a real
Claude Code transcript (ANONYMIZED), measuring ratio, speed, and salient-token
retention (identifiers/numbers/errors -- the must-keep info in coding contexts).
Only aggregate metrics are printed; raw content is never echoed.
Usage: python benchmarks/text_crusher_quality_eval.py [squad_dev.json] [transcript.jsonl]
"""
from __future__ import annotations
import glob
import json
import os
import random
import re
import sys
import time
from headroom.transforms.text_crusher import TextCrusher
_SEG = re.compile(r"(?<=[.!?])\s+|\n+")
_SALIENT = re.compile(
r"\b(?:error|exception|fail(?:ed|ure)?|warning|traceback|assert|todo|fixme)\b"
r"|\b[A-Z]{2,}\b|\b[A-Za-z_][A-Za-z0-9_]*\.[A-Za-z_][A-Za-z0-9_]*\b|\b\d+\b"
)
# --- anonymization (脱敏): scrub before any processing; never echo raw content ---
_REDACT = [
(re.compile(r"/Users/[^/\s]+"), "/Users/USER"),
(re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), "EMAIL"),
(re.compile(r"\b(?:sk|pk|ghp|gho|xox[baprs])-[A-Za-z0-9_-]{10,}\b"), "TOKEN"),
(re.compile(r"\bBearer\s+[A-Za-z0-9._-]{10,}"), "Bearer TOKEN"),
(re.compile(r"\b[A-Fa-f0-9]{40,}\b"), "HEX"),
]
def anon(t: str) -> str:
for rx, rep in _REDACT:
t = rx.sub(rep, t)
return t
def norm(s: str) -> str:
return re.sub(r"\s+", " ", s.lower()).strip()
def _segs(text: str) -> list[str]:
return [s for s in _SEG.split(text) if s.strip()]
def truncate_keep_last(text: str, ratio: float) -> str:
segs = _segs(text)
budget = int(sum(len(s) for s in segs) * ratio)
kept: list[str] = []
c = 0
for s in reversed(segs):
if c >= budget:
break
kept.append(s)
c += len(s)
return "\n".join(reversed(kept))
def random_keep(text: str, ratio: float, seed: int) -> str:
segs = _segs(text)
idx = list(range(len(segs)))
random.Random(seed).shuffle(idx)
budget = int(sum(len(s) for s in segs) * ratio)
kept: set[int] = set()
c = 0
for i in idx:
if c >= budget:
break
kept.add(i)
c += len(segs[i])
return "\n".join(segs[i] for i in sorted(kept))
def eval_squad(path: str, n: int = 200, n_distract: int = 40, ratio: float = 0.3, seed: int = 0):
data = json.load(open(path))
paras = [(p["context"], p["qas"]) for a in data["data"] for p in a["paragraphs"]]
all_ctx = [c for c, _ in paras]
examples = [
(ctx, qas[0]["question"], qas[0]["answers"][0]["text"])
for ctx, qas in paras
if qas and qas[0]["answers"]
]
rnd = random.Random(seed)
rnd.shuffle(examples)
examples = examples[:n]
tc = TextCrusher()
hit = {"text_crusher": 0, "truncate": 0, "random": 0}
tc_ratios: list[float] = []
for gold_ctx, q, ans in examples:
docs = rnd.sample(all_ctx, n_distract) + [gold_ctx]
rnd.shuffle(docs)
haystack = "\n\n".join(docs)
a = norm(ans)
out_tc = tc.compress(haystack, context=q, target_ratio=ratio).compressed
hit["text_crusher"] += a in norm(out_tc)
hit["truncate"] += a in norm(truncate_keep_last(haystack, ratio))
hit["random"] += a in norm(random_keep(haystack, ratio, seed))
tc_ratios.append(len(out_tc) / max(1, len(haystack)))
nn = len(examples)
print(
f"\n=== Part A: SQuAD answer-retention (n={nn}, distractors={n_distract}, target_ratio={ratio}) ==="
)
print(
f" TextCrusher (query-aware): {hit['text_crusher'] / nn:6.1%} answer survives compression"
)
print(f" Truncate (keep recent): {hit['truncate'] / nn:6.1%}")
print(f" Random keep: {hit['random'] / nn:6.1%}")
print(
f" TextCrusher mean char-ratio: {sum(tc_ratios) / nn:.2f} (kept ~{sum(tc_ratios) / nn:.0%} of bytes)"
)
print(" reference: kompress published must_keep_recall = 0.977 (its own labeled set)")
def _block_texts(jsonl_path: str, min_words: int, limit: int) -> list[str]:
out: list[str] = []
with open(jsonl_path) as fh:
for line in fh:
try:
o = json.loads(line)
except json.JSONDecodeError:
continue
m = o.get("message") or {}
c = m.get("content")
parts = (
[c]
if isinstance(c, str)
else [
p["text"] for p in c if isinstance(p, dict) and isinstance(p.get("text"), str)
]
if isinstance(c, list)
else []
)
for t in parts:
if len(t.split()) >= min_words:
out.append(anon(t))
if len(out) >= limit:
break
return out[:limit]
def eval_transcript(jsonl_path: str, ratio: float = 0.4, min_words: int = 1500, limit: int = 40):
blocks = _block_texts(jsonl_path, min_words, limit)
if not blocks:
print(
f"\n=== Part B: no text blocks >= {min_words} words in {os.path.basename(jsonl_path)} ==="
)
return
tc = TextCrusher()
ratios: list[float] = []
times: list[float] = []
retentions: list[float] = []
for b in blocks:
sal_before = set(_SALIENT.findall(b))
t0 = time.perf_counter()
out = tc.compress(b, target_ratio=ratio).compressed
times.append((time.perf_counter() - t0) * 1000)
sal_after = set(_SALIENT.findall(out))
retentions.append(len(sal_before & sal_after) / max(1, len(sal_before)))
ratios.append(len(out.split()) / max(1, len(b.split())))
n = len(blocks)
print(
f"\n=== Part B: real transcript fidelity (n={n} large blocks, anonymized, target_ratio={ratio}) ==="
)
print(f" mean token-ratio kept: {sum(ratios) / n:.2f}")
print(f" mean speed: {sum(times) / n:.1f} ms/block")
print(
f" salient-token retention: {sum(retentions) / n:6.1%} (identifiers/numbers/errors kept)"
)
print(
f" -> keeps salient info at {sum(retentions) / n:.0%} while dropping to {sum(ratios) / n:.0%} of tokens"
)
def eval_speed(scale_words: int = 250_000):
# Reproducible throughput on a large synthetic prose block (no external data).
text = " ".join(
f"Sentence {i} discusses subsystem {i} and its failure mode {i % 7} in detail."
for i in range(scale_words // 9)
)
nwords = len(text.split())
tc = TextCrusher()
t0 = time.perf_counter()
out = tc.compress(text, target_ratio=0.3)
ms = (time.perf_counter() - t0) * 1000
print(f"\n=== Part C: speed (synthetic, {nwords:,} words, fully reproducible) ===")
print(f" TextCrusher compress: {ms:.0f} ms ({nwords / max(ms / 1000, 1e-6):,.0f} words/sec)")
print(f" kept ratio: {out.compressed_tokens / max(1, out.original_tokens):.2f}")
print(" reference: kompress (ModernBERT ONNX) ~272s for ~1M tokens (measured, query-blind)")
print(" -> fast-vs-slow CONTRAST, not a same-input side-by-side run")
if __name__ == "__main__":
eval_speed()
squad = sys.argv[1] if len(sys.argv) > 1 else "/tmp/squad_dev.json"
tx = sys.argv[2] if len(sys.argv) > 2 else None
if os.path.exists(squad):
eval_squad(squad)
else:
print(f"SQuAD not found at {squad}; skipping Part A")
if tx is None:
found = glob.glob(os.path.expanduser("~/.claude/projects/*headroom*/*.jsonl"))
tx = max(found, key=os.path.getsize) if found else None
if tx and os.path.exists(tx):
eval_transcript(tx)
else:
print("no transcript jsonl found; skipping Part B")

View file

@ -29,6 +29,7 @@ pub mod safety;
pub mod search_compressor; pub mod search_compressor;
pub mod smart_crusher; pub mod smart_crusher;
pub mod tag_protector; pub mod tag_protector;
pub mod text_crusher;
pub mod unidiff_detector; pub mod unidiff_detector;
pub use content_detector::{ pub use content_detector::{
@ -61,4 +62,5 @@ pub use search_compressor::{
SearchCompressorStats, SearchMatch, SearchCompressorStats, SearchMatch,
}; };
pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats}; pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats};
pub use text_crusher::{TextCrusher, TextCrusherConfig, TextCrusherResult};
pub use unidiff_detector::{detect_diff, is_diff}; pub use unidiff_detector::{detect_diff, is_diff};

View file

@ -0,0 +1,34 @@
//! TextCrusher configuration (Phase 2, #1171).
//!
//! Mirrors the Python `TextCrusherConfig`. Weights and thresholds are tuning
//! knobs for the recency + relevance + salience scoring.
#[derive(Debug, Clone)]
pub struct TextCrusherConfig {
/// Keep roughly this fraction of characters.
pub target_ratio: f64,
pub w_recency: f64,
pub w_relevance: f64,
pub w_salience: f64,
/// Segments shorter than this are de-prioritized (× 0.25).
pub min_segment_chars: usize,
/// Skip a candidate when this fraction of its word-shingles is already
/// covered by kept segments (near-duplicate suppression).
pub near_dup_threshold: f64,
/// Below this many segments, pass through unchanged (nothing to gain).
pub min_segments_for_crush: usize,
}
impl Default for TextCrusherConfig {
fn default() -> Self {
TextCrusherConfig {
target_ratio: 0.5,
w_recency: 1.0,
w_relevance: 2.0,
w_salience: 1.5,
min_segment_chars: 12,
near_dup_threshold: 0.85,
min_segments_for_crush: 6,
}
}
}

View file

@ -0,0 +1,317 @@
//! TextCrusher: fast deterministic extractive prose compressor (Phase 2, #1171).
//!
//! Splits prose into sentence segments, scores each by recency + query
//! relevance + structural salience, suppresses near-duplicates via a global
//! word-shingle index, and keeps the top segments (in original order) up to a
//! target ratio. Output is extractive: the kept sentences are verbatim words
//! (each segment trimmed, re-joined with `\n`) -- no invented words, no rewrite.
//!
//! The relevance term REUSES the shared [`BM25Scorer`](crate::relevance) rather
//! than reimplementing BM25 -- only the prose-specific splitting + selection
//! lives here.
use std::cmp::Ordering;
use std::collections::HashSet;
use super::config::TextCrusherConfig;
use crate::relevance::{BM25Scorer, RelevanceScorer};
const KEYWORDS: [&str; 10] = [
"error",
"exception",
"failed",
"failure",
"fail",
"warning",
"traceback",
"assert",
"todo",
"fixme",
];
#[derive(Debug, Clone)]
pub struct TextCrusherResult {
pub compressed: String,
pub original_tokens: usize,
pub compressed_tokens: usize,
pub compression_ratio: f64,
pub kept_segments: usize,
pub total_segments: usize,
}
pub struct TextCrusher {
config: TextCrusherConfig,
scorer: BM25Scorer,
}
impl Default for TextCrusher {
fn default() -> Self {
TextCrusher::new(TextCrusherConfig::default())
}
}
impl TextCrusher {
pub fn new(config: TextCrusherConfig) -> Self {
TextCrusher {
config,
scorer: BM25Scorer::default(),
}
}
fn passthrough(content: &str, n_segments: usize) -> TextCrusherResult {
let toks = content.split_whitespace().count();
TextCrusherResult {
compressed: content.to_string(),
original_tokens: toks,
compressed_tokens: toks,
compression_ratio: 1.0,
kept_segments: n_segments,
total_segments: n_segments,
}
}
pub fn compress(
&self,
content: &str,
context: &str,
target_ratio: Option<f64>,
) -> TextCrusherResult {
let cfg = &self.config;
let ratio = target_ratio.unwrap_or(cfg.target_ratio).clamp(0.05, 1.0);
let segments = split_segments(content);
if segments.len() < cfg.min_segments_for_crush {
return Self::passthrough(content, segments.len());
}
let n = segments.len();
let total_chars: usize = segments.iter().map(|s| s.len()).sum();
// .max(1) so a tiny input never truncates the budget to 0 (which would
// admit nothing and silently fall back to a 100% passthrough).
let target_chars = ((total_chars as f64 * ratio) as usize).max(1);
// Relevance via the shared BM25 scorer (already [0, 1]).
let seg_refs: Vec<&str> = segments.iter().map(|s| s.as_str()).collect();
let relevance = self.scorer.score_batch(&seg_refs, context);
let seg_tokens: Vec<Vec<String>> = segments.iter().map(|s| tokens(s)).collect();
let mut scores = vec![0.0f64; n];
for i in 0..n {
let recency = (i as f64 + 1.0) / n as f64;
let rel = relevance.get(i).map(|r| r.score).unwrap_or(0.0);
let words: Vec<&str> = segments[i].split_whitespace().collect();
let salient = words.iter().filter(|w| is_salient(w)).count();
let salience = salient as f64 / (words.len() as f64 + 1.0);
let mut score =
cfg.w_recency * recency + cfg.w_relevance * rel + cfg.w_salience * salience;
if segments[i].len() < cfg.min_segment_chars {
score *= 0.25;
}
scores[i] = score;
}
// Highest score first; stable tiebreak by index for determinism.
let mut order: Vec<usize> = (0..n).collect();
order.sort_by(|&a, &b| {
scores[b]
.partial_cmp(&scores[a])
.unwrap_or(Ordering::Equal)
.then(a.cmp(&b))
});
let mut kept = vec![false; n];
let mut seen: HashSet<String> = HashSet::new();
let mut kept_chars = 0usize;
let mut kept_count = 0usize;
for &i in &order {
if kept_chars >= target_chars {
break;
}
let sh = shingles(&seg_tokens[i], 3);
if !sh.is_empty() {
let covered =
sh.iter().filter(|s| seen.contains(*s)).count() as f64 / sh.len() as f64;
if covered >= cfg.near_dup_threshold {
continue; // near-duplicate: most shingles already kept
}
}
kept[i] = true;
kept_count += 1;
for s in sh {
seen.insert(s);
}
kept_chars += segments[i].len();
}
if kept_count == 0 {
return Self::passthrough(content, n);
}
let compressed = (0..n)
.filter(|&i| kept[i])
.map(|i| segments[i].as_str())
.collect::<Vec<_>>()
.join("\n");
let orig_tok = content.split_whitespace().count();
let comp_tok = compressed.split_whitespace().count();
TextCrusherResult {
compression_ratio: if orig_tok > 0 {
comp_tok as f64 / orig_tok as f64
} else {
1.0
},
compressed,
original_tokens: orig_tok,
compressed_tokens: comp_tok,
kept_segments: kept_count,
total_segments: n,
}
}
}
/// Split into sentence/line segments: on newlines, and after `.`/`!`/`?`
/// followed by whitespace. Byte-faithful (kept segments are joined verbatim).
fn split_segments(text: &str) -> Vec<String> {
let mut segs = Vec::new();
for line in text.split('\n') {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let mut cur = String::new();
let mut prev_term = false;
for c in trimmed.chars() {
if prev_term && c.is_whitespace() {
let s = cur.trim();
if !s.is_empty() {
segs.push(s.to_string());
}
cur.clear();
prev_term = false;
continue;
}
cur.push(c);
prev_term = matches!(c, '.' | '!' | '?');
}
let s = cur.trim();
if !s.is_empty() {
segs.push(s.to_string());
}
}
segs
}
fn tokens(text: &str) -> Vec<String> {
let mut out = Vec::new();
let mut cur = String::new();
for c in text.chars() {
if c.is_alphanumeric() || c == '_' {
for lc in c.to_lowercase() {
cur.push(lc);
}
} else if !cur.is_empty() {
out.push(std::mem::take(&mut cur));
}
}
if !cur.is_empty() {
out.push(cur);
}
out
}
fn shingles(words: &[String], k: usize) -> HashSet<String> {
let mut set = HashSet::new();
if words.is_empty() {
return set;
}
if words.len() < k {
// Short segment: emit every sub-window (1..=len) so identical/overlapping
// short segments still near-dup-match each other. (They can't match a
// longer segment's k-grams, but short segments are score-penalized and
// rarely survive selection anyway.)
for size in 1..=words.len() {
for w in words.windows(size) {
set.insert(w.join("\u{1}"));
}
}
return set;
}
for w in words.windows(k) {
set.insert(w.join("\u{1}"));
}
set
}
/// A word carries specific, hard-to-reconstruct information if it has a digit,
/// is an error/status keyword, is ALLCAPS (2+ letters), or is a dotted
/// identifier (`foo.bar`).
fn is_salient(word: &str) -> bool {
if word.chars().any(|c| c.is_ascii_digit()) {
return true;
}
let lower = word
.trim_matches(|c: char| !c.is_alphanumeric())
.to_lowercase();
if KEYWORDS.contains(&lower.as_str()) {
return true;
}
let alpha: Vec<char> = word.chars().filter(|c| c.is_alphabetic()).collect();
if alpha.len() >= 2 && alpha.iter().all(|c| c.is_uppercase()) {
return true;
}
if let Some(dot) = word.find('.') {
let a = &word[..dot];
let b = &word[dot + 1..];
if !a.is_empty()
&& !b.is_empty()
&& a.chars()
.next()
.is_some_and(|c| c.is_alphabetic() || c == '_')
&& b.chars()
.next()
.is_some_and(|c| c.is_alphabetic() || c == '_')
{
return true;
}
}
false
}
#[cfg(test)]
mod tests {
use super::*;
fn doc(n: usize) -> String {
(0..n)
.map(|i| format!("Sentence number {i} describes a distinct topic {i} in some detail."))
.collect::<Vec<_>>()
.join(" ")
}
#[test]
fn extractive_and_compresses() {
let content = doc(40);
let r = TextCrusher::default().compress(&content, "", Some(0.3));
assert!(r.compressed_tokens < r.original_tokens);
// extractive: every output word appears in the input
let orig: HashSet<&str> = content.split_whitespace().collect();
assert!(r.compressed.split_whitespace().all(|w| orig.contains(w)));
}
#[test]
fn deterministic() {
let content = doc(40);
let tc = TextCrusher::default();
assert_eq!(
tc.compress(&content, "", Some(0.4)).compressed,
tc.compress(&content, "", Some(0.4)).compressed
);
}
#[test]
fn passthrough_when_small() {
let r = TextCrusher::default().compress("one. two. three.", "", None);
assert_eq!(r.compression_ratio, 1.0);
}
}

View file

@ -0,0 +1,11 @@
//! TextCrusher — fast deterministic extractive prose compressor (Phase 2, #1171).
//!
//! The request-path-safe alternative to ModernBERT (kompress) for large plain
//! text: heuristic sentence scoring (recency + reused BM25 relevance +
//! salience) with near-duplicate suppression, in one O(n) pass.
mod config;
mod crusher;
pub use config::TextCrusherConfig;
pub use crusher::{TextCrusher, TextCrusherResult};

View file

@ -40,6 +40,10 @@ use headroom_core::transforms::{
SearchCompressionResult as RustSearchResult, SearchCompressor as RustSearchCompressor, SearchCompressionResult as RustSearchResult, SearchCompressor as RustSearchCompressor,
SearchCompressorConfig as RustSearchConfig, SearchCompressorStats as RustSearchStats, SearchCompressorConfig as RustSearchConfig, SearchCompressorStats as RustSearchStats,
}; };
use headroom_core::transforms::{
TextCrusher as RustTextCrusher, TextCrusherConfig as RustTextCrusherConfig,
TextCrusherResult as RustTextCrusherResult,
};
use pyo3::prelude::*; use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict}; use pyo3::types::{PyBytes, PyDict};
@ -1622,6 +1626,145 @@ fn compress_openai_responses_live_zone(
} }
} }
// --- TextCrusher (Phase 2, #1171): fast extractive prose compressor ---
#[pyclass(name = "TextCrusherConfig", module = "headroom._core")]
#[derive(Clone)]
struct PyTextCrusherConfig {
inner: RustTextCrusherConfig,
}
#[pymethods]
impl PyTextCrusherConfig {
#[new]
#[pyo3(signature = (
target_ratio = 0.5,
w_recency = 1.0,
w_relevance = 2.0,
w_salience = 1.5,
min_segment_chars = 12,
near_dup_threshold = 0.85,
min_segments_for_crush = 6,
))]
#[allow(clippy::too_many_arguments)]
fn new(
target_ratio: f64,
w_recency: f64,
w_relevance: f64,
w_salience: f64,
min_segment_chars: usize,
near_dup_threshold: f64,
min_segments_for_crush: usize,
) -> Self {
Self {
inner: RustTextCrusherConfig {
target_ratio,
w_recency,
w_relevance,
w_salience,
min_segment_chars,
near_dup_threshold,
min_segments_for_crush,
},
}
}
#[getter]
fn target_ratio(&self) -> f64 {
self.inner.target_ratio
}
#[getter]
fn near_dup_threshold(&self) -> f64 {
self.inner.near_dup_threshold
}
#[getter]
fn min_segments_for_crush(&self) -> usize {
self.inner.min_segments_for_crush
}
#[getter]
fn w_recency(&self) -> f64 {
self.inner.w_recency
}
#[getter]
fn w_relevance(&self) -> f64 {
self.inner.w_relevance
}
#[getter]
fn w_salience(&self) -> f64 {
self.inner.w_salience
}
#[getter]
fn min_segment_chars(&self) -> usize {
self.inner.min_segment_chars
}
}
#[pyclass(name = "TextCrusherResult", module = "headroom._core")]
struct PyTextCrusherResult {
inner: RustTextCrusherResult,
}
#[pymethods]
impl PyTextCrusherResult {
#[getter]
fn compressed(&self) -> String {
self.inner.compressed.clone()
}
#[getter]
fn original_tokens(&self) -> usize {
self.inner.original_tokens
}
#[getter]
fn compressed_tokens(&self) -> usize {
self.inner.compressed_tokens
}
#[getter]
fn compression_ratio(&self) -> f64 {
self.inner.compression_ratio
}
#[getter]
fn kept_segments(&self) -> usize {
self.inner.kept_segments
}
#[getter]
fn total_segments(&self) -> usize {
self.inner.total_segments
}
}
#[pyclass(name = "TextCrusher", module = "headroom._core")]
struct PyTextCrusher {
inner: RustTextCrusher,
}
#[pymethods]
impl PyTextCrusher {
#[new]
#[pyo3(signature = (config = None))]
fn new(config: Option<&PyTextCrusherConfig>) -> Self {
let cfg = config.map(|c| c.inner.clone()).unwrap_or_default();
Self {
inner: RustTextCrusher::new(cfg),
}
}
/// `compress(content, context="", target_ratio=None) -> TextCrusherResult`.
/// Releases the GIL across the Rust compress call.
#[pyo3(signature = (content, context = "", target_ratio = None))]
fn compress(
&self,
py: Python<'_>,
content: &str,
context: &str,
target_ratio: Option<f64>,
) -> PyTextCrusherResult {
let content = content.to_string();
let context = context.to_string();
let inner = py.allow_threads(|| self.inner.compress(&content, &context, target_ratio));
PyTextCrusherResult { inner }
}
}
#[pymodule] #[pymodule]
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
// Bridge Rust diagnostics into Python's `logging`. headroom-core emits // Bridge Rust diagnostics into Python's `logging`. headroom-core emits
@ -1647,6 +1790,9 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PySmartCrusherConfig>()?; m.add_class::<PySmartCrusherConfig>()?;
m.add_class::<PyCrushResult>()?; m.add_class::<PyCrushResult>()?;
m.add_class::<PySmartCrusher>()?; m.add_class::<PySmartCrusher>()?;
m.add_class::<PyTextCrusherConfig>()?;
m.add_class::<PyTextCrusherResult>()?;
m.add_class::<PyTextCrusher>()?;
m.add_class::<PyDetectionResult>()?; m.add_class::<PyDetectionResult>()?;
m.add_class::<PyLogCompressorConfig>()?; m.add_class::<PyLogCompressorConfig>()?;
m.add_class::<PyLogCompressionResult>()?; m.add_class::<PyLogCompressionResult>()?;

View file

@ -0,0 +1,144 @@
"""Off-path background compression (Phase 3, #1171).
The request path must never block on ML compression. When a cold-start-large
request would otherwise run kompress synchronously under the 30s budget (and
leak a non-preemptible worker on timeout -> executor saturation -> cascade),
it instead forwards the already-cached/uncompressed messages immediately and
enqueues the compression here. A single per-process drain runs it with NO
request-coupled deadline and stores the result in the session
``CompressionCache``, so the next turn is a cache hit and the forwarded bytes
become (and stay) the compressed form.
This is per-process by design: ``CompressionCache`` is already per-process
(``HeadroomProxy._compression_caches``), and multi-worker deployments are
already warned to use ``--workers 1`` or sticky sessions, so a per-process
drain matches the existing cache semantics without any new cross-process lock.
Limitations, all fail-open (lost savings, never lost correctness): only the
token-mode cold-start path defers here -- other modes compress synchronously;
the queue is in-memory, so a restart mid-drain drops queued jobs (they re-defer
on a later turn); and a full queue or a duplicate in-flight key drops the job,
surfaced to telemetry as ``deferred:dropped``. Background work is bounded by the
Phase 1 kompress deadline (a non-terminating compressor would pin the single
drain thread, but the compressors terminate).
"""
from __future__ import annotations
import asyncio
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any
logger = logging.getLogger(__name__)
@dataclass
class _Job:
key: str
compress: Callable[[], Any] # sync callable, runs in the executor (no timeout)
store: Callable[[Any], None] # sync callable, stores the result into the cache
class BackgroundCompressor:
"""Single per-process async drain that compresses enqueued work off the
request path, with no request-coupled deadline.
``run_in_executor`` is injected so the drain reuses the proxy's compression
ThreadPoolExecutor (without the request-path ``asyncio.wait_for`` timeout),
and so tests can supply a trivial runner.
"""
def __init__(
self,
run_in_executor: Callable[[Callable[[], Any]], Awaitable[Any]],
*,
max_queue: int = 256,
) -> None:
self._run_in_executor = run_in_executor
self._queue: asyncio.Queue[_Job] = asyncio.Queue(maxsize=max_queue)
self._pending: set[str] = set()
self._task: asyncio.Task[None] | None = None
self._processed = 0
self._dropped = 0
self._errors = 0
def enqueue(
self,
key: str,
compress: Callable[[], Any],
store: Callable[[Any], None],
) -> bool:
"""Queue a compression job. Returns False (and drops) if the key is
already in flight or the queue is full -- both are safe: the request
has already been forwarded uncompressed, so a drop just defers the
savings to a later turn."""
if key in self._pending:
return False # already queued / in flight -- dedup
# Claim the slot BEFORE the job is observable in the queue so dedup is
# atomic against another enqueue of the same key.
self._pending.add(key)
try:
self._queue.put_nowait(_Job(key, compress, store))
except asyncio.QueueFull:
self._pending.discard(key)
self._dropped += 1
logger.warning(
"background compression queue full (%d); dropping %s "
"(request already forwarded uncompressed)",
self._queue.maxsize,
key,
)
return False
return True
async def _process_one(self, job: _Job) -> None:
try:
result = await self._run_in_executor(job.compress)
job.store(result)
self._processed += 1
except Exception as e: # noqa: BLE001 -- fail-open: request already went out uncompressed
self._errors += 1
logger.warning("background compression failed for %s: %s", job.key, e)
finally:
self._pending.discard(job.key)
async def _drain(self) -> None:
while True:
job = await self._queue.get()
try:
await self._process_one(job)
finally:
self._queue.task_done()
async def start(self) -> None:
if self._task is None or self._task.done():
self._task = asyncio.create_task(self._drain(), name="headroom-bg-compress")
async def stop(self, *, drain: bool = True, timeout: float = 5.0) -> None:
if self._task is None:
return
if drain:
try:
await asyncio.wait_for(self._queue.join(), timeout=timeout)
except asyncio.TimeoutError:
logger.warning(
"background compression drain timed out with %d queued",
self._queue.qsize(),
)
self._task.cancel()
try:
await self._task
except asyncio.CancelledError:
pass
self._task = None
def stats(self) -> dict[str, int]:
return {
"queued": self._queue.qsize(),
"pending": len(self._pending),
"processed": self._processed,
"dropped": self._dropped,
"errors": self._errors,
}

View file

@ -1064,30 +1064,86 @@ class AnthropicHandlerMixin:
# Record all tool_results in the verified frozen prefix as stable # Record all tool_results in the verified frozen prefix as stable
comp_cache.mark_stable_from_messages(messages, frozen_message_count) comp_cache.mark_stable_from_messages(messages, frozen_message_count)
async with stage_timer.measure("compression_first_stage"): # Phase 3 (#1171): off-path deferral gate. On a cold-
result = await self._run_compression_in_executor( # start-large request (frozen=0 + large live zone) the
# synchronous kompress run would blow the 30s budget and
# leak a non-preemptible worker. Forward the cache-
# swapped messages uncompressed NOW and compress off the
# request path; the result lands in the SAME
# CompressionCache and apply_cached swaps it in (live
# zone) on a later turn. Byte-identity holds — see the
# frozen/live invariant; one-time upstream cache miss
# when the compressed form first lands, then stable.
if (
getattr(self, "_background_compression_enabled", False)
and frozen_message_count == 0
and original_tokens >= self._background_compression_min_tokens
):
# Snapshot refs for the async job. The handler must
# NOT mutate these lists/dicts in-place after this
# point -- the background job reads them on a later
# turn. It doesn't today; keep it that way.
_bg_messages = messages
_bg_working = working_messages
_bg_frozen = frozen_message_count
# Dedup key: the gate only fires at frozen==0 (the
# first in-flight deferral of a session episode), so
# session_id alone is the right granularity -- one
# background job per session in flight -- and avoids
# JSON-serializing the large message list for a key.
accepted = self._background_compressor.enqueue(
session_id,
lambda: self.anthropic_pipeline.apply( lambda: self.anthropic_pipeline.apply(
messages=working_messages, messages=_bg_working,
model=model, model=model,
model_limit=context_limit, model_limit=context_limit,
context=extract_user_query(working_messages), context=extract_user_query(_bg_working),
frozen_message_count=frozen_message_count, frozen_message_count=_bg_frozen,
biases=biases, biases=biases,
request_id=request_id, request_id=request_id,
compression_policy=compression_policy, compression_policy=compression_policy,
**proxy_pipeline_kwargs(self.config), **proxy_pipeline_kwargs(self.config),
), ),
timeout=COMPRESSION_TIMEOUT_SECONDS, lambda result: comp_cache.update_from_result(
_bg_messages, result.messages
),
) )
# Forward uncompressed either way (the request can't
# wait); only CLAIM deferral when the job was actually
# queued. A full-queue drop is visible in telemetry as
# "deferred:dropped" and self-heals on a later turn.
optimized_messages = working_messages
transforms_applied = [
"deferred:background_compression"
if accepted
else "deferred:dropped"
]
pipeline_timing = {}
else:
async with stage_timer.measure("compression_first_stage"):
result = await self._run_compression_in_executor(
lambda: self.anthropic_pipeline.apply(
messages=working_messages,
model=model,
model_limit=context_limit,
context=extract_user_query(working_messages),
frozen_message_count=frozen_message_count,
biases=biases,
request_id=request_id,
compression_policy=compression_policy,
**proxy_pipeline_kwargs(self.config),
),
timeout=COMPRESSION_TIMEOUT_SECONDS,
)
# Cache newly compressed messages (index-aligned diff) # Cache newly compressed messages (index-aligned diff)
if result.messages != working_messages: if result.messages != working_messages:
comp_cache.update_from_result(messages, result.messages) comp_cache.update_from_result(messages, result.messages)
# Always use pipeline result — Zone 1 swaps are already applied # Always use pipeline result — Zone 1 swaps applied
optimized_messages = result.messages optimized_messages = result.messages
transforms_applied = result.transforms_applied transforms_applied = result.transforms_applied
pipeline_timing = result.timing pipeline_timing = result.timing
# Issue #327 / Bug 3: pipeline.apply uses the provider- # Issue #327 / Bug 3: pipeline.apply uses the provider-
# side tokenizer (AnthropicProvider tiktoken estimator), # side tokenizer (AnthropicProvider tiktoken estimator),
# which counts ~25% higher than the proxy-side # which counts ~25% higher than the proxy-side

View file

@ -105,6 +105,7 @@ from headroom.providers.registry import (
) )
from headroom.proxy import runtime_env from headroom.proxy import runtime_env
from headroom.proxy.auth_mode import should_stamp_codex_client from headroom.proxy.auth_mode import should_stamp_codex_client
from headroom.proxy.background_compression import BackgroundCompressor
# ============================================================================= # =============================================================================
# Extracted modules (re-exported for backward compatibility) # Extracted modules (re-exported for backward compatibility)
@ -812,6 +813,28 @@ class HeadroomProxy(
max_workers=_compression_max, max_workers=_compression_max,
thread_name_prefix="headroom-compress", thread_name_prefix="headroom-compress",
) )
# Phase 3 (#1171): off-path background compression. When enabled, a
# cold-start-large request (frozen=0 + large live zone) forwards
# uncompressed immediately and enqueues the compression here instead of
# blocking the request thread under the 30s budget (which leaks a
# non-preemptible worker -> executor saturation -> cascade). Default
# off (opt-in), fail-open. Per-process, matching _compression_caches.
self._background_compression_enabled: bool = os.environ.get(
"HEADROOM_BACKGROUND_COMPRESSION", ""
).strip().lower() in ("1", "true", "yes", "on")
try:
self._background_compression_min_tokens: int = int(
os.environ.get("HEADROOM_BACKGROUND_COMPRESSION_MIN_TOKENS", "50000")
)
except ValueError:
self._background_compression_min_tokens = 50000
# Dedicated single thread: no-timeout background jobs never contend with
# the request-path executor (Phase 3, #1171). Lazy -- no thread spawns
# until the first off-path job is submitted.
self._background_compression_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=1, thread_name_prefix="headroom-bg-compress"
)
self._background_compressor = BackgroundCompressor(self._run_compression_background)
# Gauge: currently-running compression tasks. Mutated under # Gauge: currently-running compression tasks. Mutated under
# ``_compression_metrics_lock`` from worker threads + the asyncio # ``_compression_metrics_lock`` from worker threads + the asyncio
# event loop. # event loop.
@ -1112,6 +1135,18 @@ class HeadroomProxy(
self._compression_queue_timeouts += 1 self._compression_queue_timeouts += 1
raise raise
async def _run_compression_background(self, fn): # noqa: ANN001, ANN201
"""Run a compression callable on the shared executor with NO request-
coupled deadline (Phase 3 off-path, #1171).
Unlike ``_run_compression_in_executor`` there is no ``asyncio.wait_for``
and no leaked-thread accounting: no caller is waiting, so a slow run
backs up the background queue rather than starving the request executor.
Runs on the dedicated single-thread background executor.
"""
loop = asyncio.get_running_loop()
return await loop.run_in_executor(self._background_compression_executor, fn)
def _get_compression_cache(self, session_id: str) -> CompressionCache: def _get_compression_cache(self, session_id: str) -> CompressionCache:
"""Get or create a CompressionCache for a session. """Get or create a CompressionCache for a session.
@ -1851,6 +1886,8 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
await proxy.usage_reporter.start(proxy) await proxy.usage_reporter.start(proxy)
if proxy.traffic_learner: if proxy.traffic_learner:
await proxy.traffic_learner.start() await proxy.traffic_learner.start()
if proxy._background_compression_enabled:
await proxy._background_compressor.start()
# Only start beacon if we acquire the lock (first worker wins) # Only start beacon if we acquire the lock (first worker wins)
_beacon_is_owner[0] = _try_acquire_beacon_lock() _beacon_is_owner[0] = _try_acquire_beacon_lock()
@ -1885,6 +1922,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
await proxy.usage_reporter.stop() await proxy.usage_reporter.stop()
if proxy.traffic_learner: if proxy.traffic_learner:
await proxy.traffic_learner.stop() await proxy.traffic_learner.stop()
if proxy._background_compression_enabled:
await proxy._background_compressor.stop()
proxy._background_compression_executor.shutdown(wait=False)
if proxy.code_graph_watcher: if proxy.code_graph_watcher:
proxy.code_graph_watcher.stop() proxy.code_graph_watcher.stop()
await proxy.shutdown() await proxy.shutdown()
@ -3924,7 +3964,8 @@ def run_server(
else: else:
logger.warning( logger.warning(
"Headroom is running with workers=%d. The in-memory CCR store, " "Headroom is running with workers=%d. The in-memory CCR store, "
"compression cache, prefix tracker, TOIN state, and CostTracker are all " "compression cache (incl. off-path background compression), prefix "
"tracker, TOIN state, and CostTracker are all "
"per-process; multi-worker deployments produce silent CCR retrieval " "per-process; multi-worker deployments produce silent CCR retrieval "
"failures, avoidable cache busts, and an unstable dashboard 'Proxy $ Saved' " "failures, avoidable cache busts, and an unstable dashboard 'Proxy $ Saved' "
"hero tile (each /stats poll hits a different worker's partial total) when " "hero tile (each /stats poll hits a different worker's partial total) when "

View file

@ -970,6 +970,27 @@ class ContentRouter(Transform):
self._tabular_compressor: Any = None self._tabular_compressor: Any = None
self._kompress: Any = None self._kompress: Any = None
# Phase 0 (#1171): cap the input size handed to kompress (ModernBERT
# ONNX). Its inference scales O(tokens) and runs synchronously on the
# request thread under the 30s compression budget; above this ceiling we
# route to the fast LogCompressor instead so the request path stays
# bounded. ~4 chars/token is a cheap proxy (no tokenizer needed; counts
# dense JSON/code correctly, unlike word count). 0 disables the gate.
try:
self._kompress_max_tokens: int = int(
os.environ.get("HEADROOM_KOMPRESS_MAX_TOKENS", "50000")
)
except ValueError:
self._kompress_max_tokens = 50000
self._kompress_gate_fires: int = 0
# Phase 2 (#1171): when enabled, the size-gate routes oversized text to
# the fast extractive TextCrusher (real prose savings) instead of the
# LogCompressor (~0 savings on prose). Opt-in, default off.
self._text_crusher_enabled: bool = os.environ.get(
"HEADROOM_TEXT_CRUSHER", ""
).strip().lower() in ("1", "true", "yes", "on")
self._text_crusher: Any = None
# TOIN integration for cross-strategy learning # TOIN integration for cross-strategy learning
self._toin: Any = None self._toin: Any = None
@ -1713,6 +1734,46 @@ class ContentRouter(Transform):
compressed: str | None = None compressed: str | None = None
compressed_tokens: int | None = None compressed_tokens: int | None = None
# Phase 0 (#1171): size gate. This is the single ML boundary, so gating
# here covers EVERY kompress entry point -- TEXT, KOMPRESS-direct,
# CODE_AWARE->KOMPRESS, and the strategy-fallback path all route through
# _try_ml_compressor. Kompress ONNX inference is O(tokens) and runs
# synchronously on the request thread; on a large/cold context it
# exceeds the 30s budget and leaks a non-preemptible worker (#1171).
# Above the ceiling, route to the fast LogCompressor (or pass through)
# rather than ModernBERT, keeping the request path bounded.
if self._kompress_max_tokens > 0 and len(text_to_compress) > self._kompress_max_tokens * 4:
self._kompress_gate_fires += 1
logger.info(
"kompress size-gate fired: ~%d tok (>%d) routed off ML (fire #%d)",
len(text_to_compress) // 4,
self._kompress_max_tokens,
self._kompress_gate_fires,
)
out = text_to_compress
crusher = self._get_text_crusher()
if crusher is not None:
try:
out = crusher.compress(text_to_compress, context=context or "").compressed
except Exception as e:
logger.warning(
"Kompress size-gate -> TextCrusher failed (%s); passing through", e
)
out = text_to_compress
elif self.config.enable_log_compressor:
lc = self._get_log_compressor()
if lc:
try:
out = lc.compress(text_to_compress).compressed
except Exception as e:
logger.warning(
"Kompress size-gate -> LogCompressor failed (%s); passing through", e
)
out = text_to_compress
if protected:
out = restore_tags(out, protected)
return out, len(out.split())
# Primary: Kompress. On a cold cache the model is fetched once in the # Primary: Kompress. On a cold cache the model is fetched once in the
# background (ensure_background_load) instead of blocking this request # background (ensure_background_load) instead of blocking this request
# thread on a 274MB download that races the compression timeout and # thread on a 274MB download that races the compression timeout and
@ -1844,6 +1905,22 @@ class ContentRouter(Transform):
logger.debug("LogCompressor not available") logger.debug("LogCompressor not available")
return self._log_compressor return self._log_compressor
def _get_text_crusher(self) -> Any:
"""Get TextCrusher (Phase 2, lazy load). Returns None when disabled, or
when the native ``headroom._core`` extension is not built (mirrors the
ImportError handling of the other ``_get_*`` compressor getters)."""
if not getattr(self, "_text_crusher_enabled", False):
return None
if self._text_crusher is None:
try:
from .text_crusher import TextCrusher
self._text_crusher = TextCrusher()
except ImportError:
logger.debug("TextCrusher (headroom._core) unavailable; disabling gate route")
self._text_crusher_enabled = False
return self._text_crusher
def _get_tabular_compressor(self) -> Any: def _get_tabular_compressor(self) -> Any:
"""Get TabularCompressor (lazy load).""" """Get TabularCompressor (lazy load)."""
if self._tabular_compressor is None: if self._tabular_compressor is None:

View file

@ -856,6 +856,24 @@ class KompressCompressor(Transform):
if n_words < 10: if n_words < 10:
return self._passthrough(content, n_words) return self._passthrough(content, n_words)
# Cooperative wall-clock budget (#1171): kompress ONNX inference is
# O(tokens) and non-preemptible once the request's asyncio timeout fires,
# so one large block can run for minutes holding a worker (the leak ->
# executor-saturation -> queue-timeout cascade). Bail at the next chunk
# boundary past this budget, keeping the unprocessed tail verbatim. 0
# disables. Env HEADROOM_COMPRESSION_DEADLINE_MS overrides (default 20s).
# Cached per instance: operator config, read once -- not per compress() call.
deadline_s = getattr(self, "_deadline_s", None)
if deadline_s is None:
try:
deadline_s = max(
0.0,
float(os.environ.get("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")) / 1000.0,
)
except ValueError:
deadline_s = 20.0
self._deadline_s = deadline_s
try: try:
model, tokenizer, backend = _load_kompress( model, tokenizer, backend = _load_kompress(
self.config.model_id, self.config.device, allow_download=allow_download self.config.model_id, self.config.device, allow_download=allow_download
@ -879,8 +897,23 @@ class KompressCompressor(Transform):
kept_ids: set[int] = set() kept_ids: set[int] = set()
inference_ms = 0.0 inference_ms = 0.0
chunk_count = 0 chunk_count = 0
t_deadline = time.perf_counter()
for chunk_start in range(0, n_words, max_chunk_words): for chunk_start in range(0, n_words, max_chunk_words):
if deadline_s and (time.perf_counter() - t_deadline) > deadline_s:
# Keep everything from here on verbatim and stop: a partial
# compression that returns NOW beats a full one that leaks a
# non-preemptible worker for minutes (#1171).
kept_ids.update(range(chunk_start, n_words))
logger.warning(
"Kompress hit %.1fs deadline after %d/%d words (%d chunks done); "
"kept remainder verbatim to free the request thread (#1171)",
deadline_s,
chunk_start,
n_words,
chunk_count,
)
break
chunk_count += 1 chunk_count += 1
chunk_words = words[chunk_start : chunk_start + max_chunk_words] chunk_words = words[chunk_start : chunk_start + max_chunk_words]

View file

@ -0,0 +1,59 @@
"""TextCrusher — fast deterministic extractive prose compressor (Phase 2, #1171).
Thin Python wrapper over the native ``headroom._core.TextCrusher``. The
algorithm (sentence scoring with the SHARED BM25 relevance scorer + global
word-shingle near-dup suppression) lives in Rust (``crates/headroom-core``),
reusing the same scorer SmartCrusher uses rather than reimplementing it. This
wrapper only keeps the Python-facing interface stable for ContentRouter + tests.
TextCrusher is the request-path-safe alternative to ModernBERT (kompress) for
large plain text: ~milliseconds instead of minutes. Extractive -- the kept
sentences are verbatim words (each segment trimmed, re-joined with newlines),
never paraphrased; it selects, it does not rewrite.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from headroom._core import TextCrusher as _RustTextCrusher
from headroom._core import TextCrusherConfig as _RustTextCrusherConfig
@dataclass
class TextCrusherConfig:
target_ratio: float = 0.5
w_recency: float = 1.0
w_relevance: float = 2.0
w_salience: float = 1.5
min_segment_chars: int = 12
near_dup_threshold: float = 0.85
min_segments_for_crush: int = 6
class TextCrusher:
"""Extractive prose compressor. ``compress`` returns a result whose
``compressed`` text is the kept input sentences verbatim (each trimmed,
re-joined with newlines) in original order -- selection, not rewriting.
Backed by the Rust core."""
def __init__(self, config: TextCrusherConfig | None = None) -> None:
cfg = config or TextCrusherConfig()
self._rust = _RustTextCrusher(
_RustTextCrusherConfig(
target_ratio=cfg.target_ratio,
w_recency=cfg.w_recency,
w_relevance=cfg.w_relevance,
w_salience=cfg.w_salience,
min_segment_chars=cfg.min_segment_chars,
near_dup_threshold=cfg.near_dup_threshold,
min_segments_for_crush=cfg.min_segments_for_crush,
)
)
def compress(self, content: str, context: str = "", target_ratio: float | None = None) -> Any:
"""Returns a ``TextCrusherResult`` (Rust pyclass) with ``.compressed``,
``.original_tokens``, ``.compressed_tokens``, ``.compression_ratio``,
``.kept_segments``, ``.total_segments``."""
return self._rust.compress(content, context, target_ratio)

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "plain_prose",
"input": {
"content": "Sentence number 0 explains how distributed systems reconcile state across topic 0. Sentence number 1 explains how distributed systems reconcile state across topic 1. Sentence number 2 explains how distributed systems reconcile state across topic 2. Sentence number 3 explains how distributed systems reconcile state across topic 3. Sentence number 4 explains how distributed systems reconcile state across topic 4. Sentence number 5 explains how distributed systems reconcile state across topic 5. Sentence number 6 explains how distributed systems reconcile state across topic 6. Sentence number 7 explains how distributed systems reconcile state across topic 7. Sentence number 8 explains how distributed systems reconcile state across topic 8. Sentence number 9 explains how distributed systems reconcile state across topic 9. Sentence number 10 explains how distributed systems reconcile state across topic 10. Sentence number 11 explains how distributed systems reconcile state across topic 11. Sentence number 12 explains how distributed systems reconcile state across topic 12. Sentence number 13 explains how distributed systems reconcile state across topic 13. Sentence number 14 explains how distributed systems reconcile state across topic 14. Sentence number 15 explains how distributed systems reconcile state across topic 15. Sentence number 16 explains how distributed systems reconcile state across topic 16. Sentence number 17 explains how distributed systems reconcile state across topic 17. Sentence number 18 explains how distributed systems reconcile state across topic 18. Sentence number 19 explains how distributed systems reconcile state across topic 19. Sentence number 20 explains how distributed systems reconcile state across topic 20. Sentence number 21 explains how distributed systems reconcile state across topic 21. Sentence number 22 explains how distributed systems reconcile state across topic 22. Sentence number 23 explains how distributed systems reconcile state across topic 23. Sentence number 24 explains how distributed systems reconcile state across topic 24. Sentence number 25 explains how distributed systems reconcile state across topic 25. Sentence number 26 explains how distributed systems reconcile state across topic 26. Sentence number 27 explains how distributed systems reconcile state across topic 27. Sentence number 28 explains how distributed systems reconcile state across topic 28. Sentence number 29 explains how distributed systems reconcile state across topic 29.",
"context": "how do distributed systems reconcile state",
"target_ratio": 0.3
},
"output": {
"compressed": "Sentence number 21 explains how distributed systems reconcile state across topic 21.\nSentence number 22 explains how distributed systems reconcile state across topic 22.\nSentence number 23 explains how distributed systems reconcile state across topic 23.\nSentence number 24 explains how distributed systems reconcile state across topic 24.\nSentence number 25 explains how distributed systems reconcile state across topic 25.\nSentence number 26 explains how distributed systems reconcile state across topic 26.\nSentence number 27 explains how distributed systems reconcile state across topic 27.\nSentence number 28 explains how distributed systems reconcile state across topic 28.\nSentence number 29 explains how distributed systems reconcile state across topic 29.",
"original_tokens": 360,
"compressed_tokens": 108,
"compression_ratio": 0.3,
"kept_segments": 9,
"total_segments": 30
},
"input_sha256": "197195471bb9b8be5fc4d5f555cdd0d4857ec2ad3cd23135d8df3dc81ff50c35"
}

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "plain_prose_no_query",
"input": {
"content": "Sentence number 0 explains how distributed systems reconcile state across topic 0. Sentence number 1 explains how distributed systems reconcile state across topic 1. Sentence number 2 explains how distributed systems reconcile state across topic 2. Sentence number 3 explains how distributed systems reconcile state across topic 3. Sentence number 4 explains how distributed systems reconcile state across topic 4. Sentence number 5 explains how distributed systems reconcile state across topic 5. Sentence number 6 explains how distributed systems reconcile state across topic 6. Sentence number 7 explains how distributed systems reconcile state across topic 7. Sentence number 8 explains how distributed systems reconcile state across topic 8. Sentence number 9 explains how distributed systems reconcile state across topic 9. Sentence number 10 explains how distributed systems reconcile state across topic 10. Sentence number 11 explains how distributed systems reconcile state across topic 11. Sentence number 12 explains how distributed systems reconcile state across topic 12. Sentence number 13 explains how distributed systems reconcile state across topic 13. Sentence number 14 explains how distributed systems reconcile state across topic 14. Sentence number 15 explains how distributed systems reconcile state across topic 15. Sentence number 16 explains how distributed systems reconcile state across topic 16. Sentence number 17 explains how distributed systems reconcile state across topic 17. Sentence number 18 explains how distributed systems reconcile state across topic 18. Sentence number 19 explains how distributed systems reconcile state across topic 19. Sentence number 20 explains how distributed systems reconcile state across topic 20. Sentence number 21 explains how distributed systems reconcile state across topic 21. Sentence number 22 explains how distributed systems reconcile state across topic 22. Sentence number 23 explains how distributed systems reconcile state across topic 23. Sentence number 24 explains how distributed systems reconcile state across topic 24. Sentence number 25 explains how distributed systems reconcile state across topic 25. Sentence number 26 explains how distributed systems reconcile state across topic 26. Sentence number 27 explains how distributed systems reconcile state across topic 27. Sentence number 28 explains how distributed systems reconcile state across topic 28. Sentence number 29 explains how distributed systems reconcile state across topic 29.",
"context": "",
"target_ratio": 0.5
},
"output": {
"compressed": "Sentence number 15 explains how distributed systems reconcile state across topic 15.\nSentence number 16 explains how distributed systems reconcile state across topic 16.\nSentence number 17 explains how distributed systems reconcile state across topic 17.\nSentence number 18 explains how distributed systems reconcile state across topic 18.\nSentence number 19 explains how distributed systems reconcile state across topic 19.\nSentence number 20 explains how distributed systems reconcile state across topic 20.\nSentence number 21 explains how distributed systems reconcile state across topic 21.\nSentence number 22 explains how distributed systems reconcile state across topic 22.\nSentence number 23 explains how distributed systems reconcile state across topic 23.\nSentence number 24 explains how distributed systems reconcile state across topic 24.\nSentence number 25 explains how distributed systems reconcile state across topic 25.\nSentence number 26 explains how distributed systems reconcile state across topic 26.\nSentence number 27 explains how distributed systems reconcile state across topic 27.\nSentence number 28 explains how distributed systems reconcile state across topic 28.\nSentence number 29 explains how distributed systems reconcile state across topic 29.",
"original_tokens": 360,
"compressed_tokens": 180,
"compression_ratio": 0.5,
"kept_segments": 15,
"total_segments": 30
},
"input_sha256": "197195471bb9b8be5fc4d5f555cdd0d4857ec2ad3cd23135d8df3dc81ff50c35"
}

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "redundant",
"input": {
"content": "The quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nThe quick brown fox jumps over the very lazy dog every single morning.\nA distinct fact about subsystem 0 is recorded plainly here.\nA distinct fact about subsystem 1 is recorded plainly here.\nA distinct fact about subsystem 2 is recorded plainly here.\nA distinct fact about subsystem 3 is recorded plainly here.\nA distinct fact about subsystem 4 is recorded plainly here.\nA distinct fact about subsystem 5 is recorded plainly here.\nA distinct fact about subsystem 6 is recorded plainly here.\nA distinct fact about subsystem 7 is recorded plainly here.",
"context": "",
"target_ratio": 0.9
},
"output": {
"compressed": "The quick brown fox jumps over the very lazy dog every single morning.\nA distinct fact about subsystem 0 is recorded plainly here.\nA distinct fact about subsystem 1 is recorded plainly here.\nA distinct fact about subsystem 2 is recorded plainly here.\nA distinct fact about subsystem 3 is recorded plainly here.\nA distinct fact about subsystem 4 is recorded plainly here.\nA distinct fact about subsystem 5 is recorded plainly here.\nA distinct fact about subsystem 6 is recorded plainly here.\nA distinct fact about subsystem 7 is recorded plainly here.",
"original_tokens": 210,
"compressed_tokens": 93,
"compression_ratio": 0.44285714285714284,
"kept_segments": 9,
"total_segments": 18
},
"input_sha256": "2eaa453c47191d4687669467586834d66ab14aedbbe6ee12d01487dbd0922960"
}

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "salient_heavy",
"input": {
"content": "ERROR connection refused at host 10.0.0.42 after 3 retries.\nThe authentication module validated tokens against auth.registry before forwarding.\nA traceback was logged with code 500 and request_id req-9182.\nJust some generic filler text without any specific identifiers here.\nMore plain filler describing the overall behavior in vague terms.\nWarning: cache hit ratio dropped to 71 percent during the spike.\nAnother unremarkable sentence with no salient tokens at all today.\nThe pipeline.apply call returned 42 kept rows out of 1000 total.",
"context": "authentication tokens errors",
"target_ratio": 0.4
},
"output": {
"compressed": "The authentication module validated tokens against auth.registry before forwarding.\nAnother unremarkable sentence with no salient tokens at all today.\nThe pipeline.apply call returned 42 kept rows out of 1000 total.",
"original_tokens": 80,
"compressed_tokens": 30,
"compression_ratio": 0.375,
"kept_segments": 3,
"total_segments": 8
},
"input_sha256": "57cb8ef778b5c69b199541ebd017a34883e0084343d92ebb54e91773041b4217"
}

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "short_passthrough",
"input": {
"content": "one thing. two thing. three thing.",
"context": "",
"target_ratio": null
},
"output": {
"compressed": "one thing. two thing. three thing.",
"original_tokens": 6,
"compressed_tokens": 6,
"compression_ratio": 1.0,
"kept_segments": 3,
"total_segments": 3
},
"input_sha256": "d30100c6f67fdae02d1c9312e9572a85951191111f39233b7aa8a86ea782959b"
}

View file

@ -0,0 +1,18 @@
{
"transform": "text_crusher",
"label": "unicode",
"input": {
"content": "句子 0 描述了系统在主题 0 上的行为细节。 句子 1 描述了系统在主题 1 上的行为细节。 句子 2 描述了系统在主题 2 上的行为细节。 句子 3 描述了系统在主题 3 上的行为细节。 句子 4 描述了系统在主题 4 上的行为细节。 句子 5 描述了系统在主题 5 上的行为细节。 句子 6 描述了系统在主题 6 上的行为细节。 句子 7 描述了系统在主题 7 上的行为细节。 句子 8 描述了系统在主题 8 上的行为细节。 句子 9 描述了系统在主题 9 上的行为细节。 句子 10 描述了系统在主题 10 上的行为细节。 句子 11 描述了系统在主题 11 上的行为细节。",
"context": "系统",
"target_ratio": 0.4
},
"output": {
"compressed": "句子 0 描述了系统在主题 0 上的行为细节。 句子 1 描述了系统在主题 1 上的行为细节。 句子 2 描述了系统在主题 2 上的行为细节。 句子 3 描述了系统在主题 3 上的行为细节。 句子 4 描述了系统在主题 4 上的行为细节。 句子 5 描述了系统在主题 5 上的行为细节。 句子 6 描述了系统在主题 6 上的行为细节。 句子 7 描述了系统在主题 7 上的行为细节。 句子 8 描述了系统在主题 8 上的行为细节。 句子 9 描述了系统在主题 9 上的行为细节。 句子 10 描述了系统在主题 10 上的行为细节。 句子 11 描述了系统在主题 11 上的行为细节。",
"original_tokens": 60,
"compressed_tokens": 60,
"compression_ratio": 1.0,
"kept_segments": 1,
"total_segments": 1
},
"input_sha256": "d46b3fb761b16a7771144b388721042e322655450fb9f10e80459d118aabf600"
}

View file

@ -0,0 +1,94 @@
"""Record TextCrusher parity fixtures (Phase 2, #1171).
Locks the Rust core's ``compress`` output for a fixed set of deterministic
scenarios so a future change to the Rust algorithm is caught as a regression.
The Python wrapper delegates to ``headroom._core.TextCrusher``, so these
fixtures are recorded from (and verified against) the native implementation.
Re-record after an intentional algorithm change:
python tests/parity/record_text_crusher.py
"""
from __future__ import annotations
import hashlib
import json
import os
from headroom.transforms.text_crusher import TextCrusher
FIXTURE_DIR = os.path.join(os.path.dirname(__file__), "fixtures", "text_crusher")
def _prose(n: int) -> str:
return " ".join(
f"Sentence number {i} explains how distributed systems reconcile state across topic {i}."
for i in range(n)
)
def _redundant() -> str:
dup = "The quick brown fox jumps over the very lazy dog every single morning."
uniques = [f"A distinct fact about subsystem {i} is recorded plainly here." for i in range(8)]
return "\n".join([dup] * 10 + uniques)
def _salient() -> str:
return "\n".join(
[
"ERROR connection refused at host 10.0.0.42 after 3 retries.",
"The authentication module validated tokens against auth.registry before forwarding.",
"A traceback was logged with code 500 and request_id req-9182.",
"Just some generic filler text without any specific identifiers here.",
"More plain filler describing the overall behavior in vague terms.",
"Warning: cache hit ratio dropped to 71 percent during the spike.",
"Another unremarkable sentence with no salient tokens at all today.",
"The pipeline.apply call returned 42 kept rows out of 1000 total.",
]
)
# (label, content, context, target_ratio)
SCENARIOS: list[tuple[str, str, str, float | None]] = [
("plain_prose", _prose(30), "how do distributed systems reconcile state", 0.3),
("plain_prose_no_query", _prose(30), "", 0.5),
("redundant", _redundant(), "", 0.9),
("salient_heavy", _salient(), "authentication tokens errors", 0.4),
("short_passthrough", "one thing. two thing. three thing.", "", None),
(
"unicode",
" ".join(f"句子 {i} 描述了系统在主题 {i} 上的行为细节。" for i in range(12)),
"系统",
0.4,
),
]
def record() -> None:
os.makedirs(FIXTURE_DIR, exist_ok=True)
tc = TextCrusher()
for label, content, context, ratio in SCENARIOS:
r = tc.compress(content, context, ratio)
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
fixture = {
"transform": "text_crusher",
"label": label,
"input": {"content": content, "context": context, "target_ratio": ratio},
"output": {
"compressed": r.compressed,
"original_tokens": r.original_tokens,
"compressed_tokens": r.compressed_tokens,
"compression_ratio": r.compression_ratio,
"kept_segments": r.kept_segments,
"total_segments": r.total_segments,
},
"input_sha256": digest,
}
path = os.path.join(FIXTURE_DIR, f"{label}_{digest[:12]}.json")
with open(path, "w", encoding="utf-8") as fh:
json.dump(fixture, fh, indent=2, ensure_ascii=False)
print(f"wrote {path}")
if __name__ == "__main__":
record()

View file

@ -0,0 +1,109 @@
"""Phase 3 (#1171): off-path BackgroundCompressor.
A single per-process drain compresses enqueued work with NO request-coupled
deadline and stores the result -- so the request path never blocks on ML.
Tests use asyncio.run so they do not depend on a pytest-asyncio config.
"""
from __future__ import annotations
import asyncio
from headroom.proxy.background_compression import BackgroundCompressor
async def _passthrough_executor(fn):
return fn()
def test_enqueue_compresses_and_stores():
async def main():
stored: list[str] = []
bc = BackgroundCompressor(_passthrough_executor)
await bc.start()
ok = bc.enqueue("k1", lambda: "COMPRESSED", lambda r: stored.append(r))
await asyncio.wait_for(bc._queue.join(), timeout=2)
await bc.stop()
return ok, stored, bc.stats()
ok, stored, stats = asyncio.run(main())
assert ok is True
assert stored == ["COMPRESSED"]
assert stats["processed"] == 1
assert stats["pending"] == 0
def test_dedup_skips_inflight_key():
async def main():
calls: list[int] = []
async def slow_executor(fn):
await asyncio.sleep(0.05)
return fn()
bc = BackgroundCompressor(slow_executor)
await bc.start()
first = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
second = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
await asyncio.wait_for(bc._queue.join(), timeout=2)
await bc.stop()
return first, second, len(calls)
first, second, n = asyncio.run(main())
assert first is True
assert second is False # duplicate key skipped
assert n == 1
def test_fail_open_on_compress_error():
async def main():
stored: list[str] = []
bc = BackgroundCompressor(_passthrough_executor)
await bc.start()
def boom():
raise RuntimeError("kompress exploded")
bc.enqueue("bad", boom, lambda r: stored.append(r))
bc.enqueue("good", lambda: "ok", lambda r: stored.append(r))
await asyncio.wait_for(bc._queue.join(), timeout=2)
await bc.stop()
return stored, bc.stats()
stored, stats = asyncio.run(main())
assert stored == ["ok"] # failing job did not store; later job still processed
assert stats["errors"] == 1
assert stats["processed"] == 1
assert stats["pending"] == 0 # key released even on error
def test_queue_full_drops_without_raising():
async def main():
# Never start the drain, so the queue fills and stays full.
bc = BackgroundCompressor(_passthrough_executor, max_queue=2)
results = [bc.enqueue(f"k{i}", lambda: None, lambda r: None) for i in range(5)]
return results, bc.stats()
results, stats = asyncio.run(main())
assert results[:2] == [True, True]
assert results[2:] == [False, False, False] # overflow dropped, no exception
assert stats["dropped"] == 3
def test_stop_drains_remaining_jobs():
async def main():
stored: list[str] = []
async def slow_executor(fn):
await asyncio.sleep(0.02)
return fn()
bc = BackgroundCompressor(slow_executor)
await bc.start()
for i in range(4):
bc.enqueue(f"k{i}", (lambda i=i: f"c{i}"), lambda r: stored.append(r))
await bc.stop() # drain=True by default -> waits for queue.join()
return sorted(stored)
stored = asyncio.run(main())
assert stored == ["c0", "c1", "c2", "c3"]

View file

@ -0,0 +1,73 @@
"""Phase 3 (#1171) byte-identity + off-path data flow.
The off-path design's correctness argument is: forwarding the uncompressed
messages on turn N and the compressed form on turn N+1 does NOT corrupt the
upstream prefix cache, because ``apply_cached`` swaps in the stored compressed
bytes verbatim (one-time miss, then stable). These tests pin that claim and the
end-to-end enqueue -> drain -> store -> cache-hit flow the handler gate relies
on, without standing up the full Anthropic handler.
"""
from __future__ import annotations
import asyncio
from headroom.cache.compression_cache import CompressionCache
from headroom.proxy.background_compression import BackgroundCompressor
def _tool(content: str) -> dict:
return {"role": "tool", "content": content}
def test_apply_cached_is_byte_identical_and_stable():
cache = CompressionCache()
originals = [_tool("x " * 10000)] # a large tool result
compressed = [_tool("COMPRESSED")]
cache.update_from_result(originals, compressed)
# One-time miss already paid; from here the swap is verbatim AND stable
# across repeated turns (no every-turn thrash).
out1 = cache.apply_cached(originals)
out2 = cache.apply_cached(originals)
assert out1[0]["content"] == "COMPRESSED"
assert out1 == out2
# apply_cached never mutates its input.
assert originals[0]["content"] == "x " * 10000
def test_unchanged_content_is_not_swapped():
cache = CompressionCache()
originals = [_tool("same bytes")]
# Pipeline returned identical content (nothing to compress) -> no mapping.
cache.update_from_result(originals, [_tool("same bytes")])
assert cache.apply_cached(originals)[0]["content"] == "same bytes"
def test_offpath_enqueue_drain_store_then_cache_hit():
async def main():
cache = CompressionCache()
originals = [_tool("x " * 10000)]
async def run(fn): # trivial in-loop executor stand-in
return fn()
bc = BackgroundCompressor(run)
await bc.start()
# Exactly the two lambdas the deferral gate passes: compress -> produce
# the compressed messages; store -> fold them into the session cache.
bc.enqueue(
"sess:42",
lambda: [_tool("COMPRESSED")],
lambda result: cache.update_from_result(originals, result),
)
await bc._queue.join()
await bc.stop()
return cache.apply_cached(originals), bc.stats()
out, stats = asyncio.run(main())
# After the background job ran, the next turn is a byte-identical cache hit.
assert out[0]["content"] == "COMPRESSED"
assert stats["processed"] == 1
assert stats["errors"] == 0

View file

@ -0,0 +1,82 @@
"""Phase 1 (#1171): kompress cooperative chunk-boundary deadline.
Kompress ONNX inference is O(tokens) and non-preemptible once the request's
asyncio timeout fires, so one large block can run for minutes holding a worker
(the leak -> executor-saturation -> queue-timeout cascade). compress() checks a
wall-clock budget at each chunk boundary and, when over, keeps the unprocessed
tail verbatim and returns -- a partial compression that returns now beats a full
one that leaks.
"""
from __future__ import annotations
from headroom.transforms import kompress_compressor as kc
def test_compress_bails_at_deadline_keeping_tail_verbatim(monkeypatch):
# Fake clock: the pre-loop stamp reads 0s, the first loop-top check reads
# 999s elapsed -> deadline trips on chunk 0 before any model/tokenizer use.
clock = iter([0.0] + [999.0] * 50)
monkeypatch.setattr(kc.time, "perf_counter", lambda: next(clock))
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (object(), object(), "onnx"))
monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")
comp = kc.KompressCompressor()
monkeypatch.setattr(comp, "_should_batch_single_content", lambda *a, **k: False)
content = " ".join(f"w{i}" for i in range(1000))
result = comp.compress(content)
# Deadline tripped on the first chunk -> nothing dropped, tail kept verbatim.
assert result.compressed_tokens == 1000
assert result.compressed.split() == content.split()
def test_compress_partial_run_keeps_processed_head_plus_verbatim_tail(monkeypatch):
# The real partial case: chunk 0 processes (gets compressed), chunk 1 trips
# the deadline (kept verbatim). Output must be compressed-head + verbatim-tail.
# Clock: call1=t_deadline(0); calls 2-4 are chunk-0's check+inference reads
# (under budget); call 5+ is chunk-1's check -> trips.
# Robust clock: jump past the deadline only AFTER chunk 0 is processed
# (tracked via the model mock), so adding perf_counter calls inside the chunk
# body -- e.g. sub-stage timing -- can't shift when the deadline trips.
state = {"chunks_done": 0}
def fake_clock():
return 999.0 if state["chunks_done"] >= 1 else 0.0
monkeypatch.setattr(kc.time, "perf_counter", fake_clock)
class _Enc(dict):
def word_ids(self, batch_index=0):
return self["_word_ids"]
class _Tok:
def __call__(self, chunk_words, **kw):
n = len(chunk_words)
return _Enc(input_ids=[[0] * n], attention_mask=[[1] * n], _word_ids=list(range(n)))
class _Model:
def get_keep_mask(self, input_ids, attention_mask):
n = len(input_ids[0])
mask = [[i < n // 2 for i in range(n)]] # keep first half of the chunk
state["chunks_done"] += 1 # after chunk 0, the clock trips the deadline
return mask
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (_Model(), _Tok(), "onnx"))
monkeypatch.setattr(kc, "_model_device_type", lambda *a, **k: "cpu")
monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")
comp = kc.KompressCompressor()
comp.config.chunk_words = 10 # 20 words -> 2 chunks
monkeypatch.setattr(comp, "_should_batch_single_content", lambda *a, **k: False)
words = [f"w{i}" for i in range(20)]
out = comp.compress(" ".join(words)).compressed.split()
# chunk 0 processed: first half kept (w0..w4), second half dropped (w5..w9)
assert "w0" in out and "w4" in out
assert "w5" not in out and "w9" not in out
# chunk 1 tripped the deadline -> its words kept verbatim (w10..w19 all present)
for i in range(10, 20):
assert f"w{i}" in out

View file

@ -0,0 +1,76 @@
"""Phase 0 (#1171): kompress input-size gate.
Kompress (ModernBERT ONNX) inference scales O(tokens) and runs synchronously on
the request thread under the 30s compression budget. Above a size ceiling the
router must route around the ML path (to the fast LogCompressor, else
passthrough) so a large/cold context can't blow the timeout and leak a
non-preemptible worker. The gate lives inside ``_try_ml_compressor`` so it
covers EVERY kompress entry point (TEXT, KOMPRESS-direct, CODE_AWARE, and the
strategy-fallback path), all of which funnel through that single boundary.
"""
from __future__ import annotations
import pytest
from headroom.transforms import ContentRouter
from headroom.transforms.content_router import CompressionStrategy
def test_oversized_input_is_gated_away_from_kompress(monkeypatch):
router = ContentRouter()
router._kompress_max_tokens = 50 # tiny ceiling (×4 = 200 chars)
def _boom(): # kompress must never be fetched for a gated input
raise AssertionError("kompress must not be invoked for oversized (gated) input")
monkeypatch.setattr(router, "_get_kompress", _boom)
big = "the quick brown fox jumps over the lazy dog. " * 100 # >200 chars
out, ntok = router._try_ml_compressor(big, "")
assert router._kompress_gate_fires == 1
assert isinstance(out, str) and ntok > 0
def test_small_input_still_routes_to_kompress(monkeypatch):
router = ContentRouter()
router._kompress_max_tokens = 50000 # default ceiling; small input is under it
# Return no kompressor so the ML block is a fast no-op (no model load in test).
monkeypatch.setattr(router, "_get_kompress", lambda: None)
out, ntok = router._try_ml_compressor("short text", "")
assert router._kompress_gate_fires == 0 # gate did NOT fire for small input
assert isinstance(out, str)
def test_gate_disabled_when_threshold_zero(monkeypatch):
router = ContentRouter()
router._kompress_max_tokens = 0 # disabled
monkeypatch.setattr(router, "_get_kompress", lambda: None)
big = "x " * 100000
out, ntok = router._try_ml_compressor(big, "")
assert router._kompress_gate_fires == 0 # disabled → never gates
@pytest.mark.parametrize("strategy", [CompressionStrategy.KOMPRESS, CompressionStrategy.TEXT])
def test_gate_fires_through_strategy_dispatch(monkeypatch, strategy):
# Funnel check: drive the strategy dispatch (not _try_ml_compressor directly)
# and confirm both ML strategies reach the single gate boundary.
router = ContentRouter()
router._kompress_max_tokens = 50
def _boom():
raise AssertionError("kompress must not run for gated input")
monkeypatch.setattr(router, "_get_kompress", _boom)
big = "the quick brown fox jumps over the lazy dog. " * 100
router._apply_strategy_to_content(big, strategy, "")
assert router._kompress_gate_fires == 1

View file

@ -0,0 +1,65 @@
"""Phase 2 (#1171): TextCrusher fast extractive compressor.
Validates the core contract: extractive (no invented words), deterministic,
actually compresses, suppresses near-duplicates, and preferentially keeps
query-relevant segments. End-to-end answer-quality vs kompress is validated
separately via headroom/evals before defaulting it on.
"""
from __future__ import annotations
from headroom.transforms.text_crusher import TextCrusher, TextCrusherConfig
def _doc(n: int = 40) -> str:
return " ".join(
f"Sentence number {i} describes a distinct topic {i} in some detail." for i in range(n)
)
def test_extractive_invents_no_new_words():
content = _doc()
r = TextCrusher().compress(content, target_ratio=0.5)
orig_words = set(content.split())
assert set(r.compressed.split()) <= orig_words
def test_deterministic():
content = _doc()
a = TextCrusher().compress(content, target_ratio=0.4).compressed
b = TextCrusher().compress(content, target_ratio=0.4).compressed
assert a == b
def test_actually_compresses_large_text():
r = TextCrusher().compress(_doc(60), target_ratio=0.3)
assert r.compressed_tokens < r.original_tokens
assert r.compression_ratio < 0.6
def test_passthrough_when_too_few_segments():
content = "one thing. two thing. three thing." # < min_segments_for_crush (6)
r = TextCrusher().compress(content)
assert r.compressed == content
assert r.compression_ratio == 1.0
def test_near_duplicates_suppressed():
dup = "The quick brown fox jumps over the very lazy dog today."
uniques = [f"A unique fact about item {i} stated plainly here." for i in range(8)]
content = "\n".join([dup] * 10 + uniques)
r = TextCrusher(TextCrusherConfig(near_dup_threshold=0.8)).compress(content, target_ratio=0.9)
# The duplicated sentence must not be kept 10 times.
assert r.compressed.count("quick brown fox") <= 2
def test_relevance_keeps_query_relevant_segment():
filler = [f"Filler line number {i} with generic words and padding here." for i in range(30)]
needle = "The authentication token expires after thirty minutes of inactivity."
content = "\n".join(filler[:15] + [needle] + filler[15:])
r = TextCrusher().compress(
content,
context="how long until the authentication token expires",
target_ratio=0.2,
)
assert "authentication token expires" in r.compressed

View file

@ -0,0 +1,38 @@
"""TextCrusher parity (Phase 2, #1171): the native Rust core must keep
reproducing the recorded ``compress`` output for every fixture. Catches drift
in the Rust algorithm. Re-record intentional changes with:
python tests/parity/record_text_crusher.py
"""
from __future__ import annotations
import glob
import json
import os
import pytest
from headroom.transforms.text_crusher import TextCrusher
_FIXTURE_GLOB = os.path.join(
os.path.dirname(__file__), "..", "parity", "fixtures", "text_crusher", "*.json"
)
FIXTURES = sorted(glob.glob(_FIXTURE_GLOB))
def test_fixtures_exist():
assert FIXTURES, "no text_crusher parity fixtures; run tests/parity/record_text_crusher.py"
@pytest.mark.parametrize("path", FIXTURES, ids=lambda p: os.path.basename(p))
def test_text_crusher_matches_recorded(path):
with open(path, encoding="utf-8") as fh:
fx = json.load(fh)
inp = fx["input"]
r = TextCrusher().compress(inp["content"], inp["context"], inp["target_ratio"])
exp = fx["output"]
assert r.compressed == exp["compressed"]
assert r.compressed_tokens == exp["compressed_tokens"]
assert r.kept_segments == exp["kept_segments"]
assert r.total_segments == exp["total_segments"]
assert abs(r.compression_ratio - exp["compression_ratio"]) < 1e-9

View file

@ -0,0 +1,38 @@
"""Phase 2 (#1171): the kompress size-gate routes oversized text to TextCrusher
when HEADROOM_TEXT_CRUSHER is enabled (real prose savings), instead of the
LogCompressor (which yields ~0 on prose) or ModernBERT (slow)."""
from __future__ import annotations
from headroom.transforms.content_router import ContentRouter
def _prose() -> str:
return " ".join(
f"Sentence {i} about distributed systems and authentication tokens expiring soon."
for i in range(300)
)
def test_gate_routes_to_text_crusher_when_enabled(monkeypatch):
monkeypatch.setenv("HEADROOM_TEXT_CRUSHER", "1")
router = ContentRouter()
router._kompress_max_tokens = 50 # tiny ceiling so the gate fires
def _boom():
raise AssertionError("kompress must not run for gated input")
monkeypatch.setattr(router, "_get_kompress", _boom)
prose = _prose()
out, ntok = router._try_ml_compressor(prose, "authentication tokens")
assert router._kompress_gate_fires == 1
assert ntok < len(prose.split()) # TextCrusher actually compressed (LogCompressor ~0)
assert set(out.split()) <= set(prose.split()) # extractive: no invented words
def test_text_crusher_disabled_by_default(monkeypatch):
monkeypatch.delenv("HEADROOM_TEXT_CRUSHER", raising=False)
router = ContentRouter()
assert router._get_text_crusher() is None