Skip to content

Commit 6c68ff4

Browse files
lifeodysseyclaude
andauthored
perf(compression): take large cold-start contexts off the synchronous kompress path (headroomlabs-ai#1171) (headroomlabs-ai#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 headroomlabs-ai#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>
1 parent c7f75b2 commit 6c68ff4

28 files changed

Lines changed: 1855 additions & 14 deletions

.codegraph/.gitignore

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# CodeGraph data files
2+
# These are local to each machine and should not be committed
3+
4+
# Database
5+
*.db
6+
*.db-wal
7+
*.db-shm
8+
9+
# Cache
10+
cache/
11+
12+
# Logs
13+
*.log
14+
15+
# Hook markers
16+
.dirty

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1616

1717
* **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.
1818
* **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)).
19+
* **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`).
1920
* **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)).
2021
* **vibe:** add Mistral Vibe CLI support with `headroom wrap vibe`.
2122
* **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.
Lines changed: 221 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
1+
#!/usr/bin/env python3
2+
"""Quality eval for TextCrusher (Phase 2, #1171): does extractive compression
3+
preserve the answer-bearing content? No LLM/API calls -- fully local.
4+
5+
Part A -- SQuAD answer-retention (the strong, labeled metric): bury a real QA
6+
answer in a haystack of distractor paragraphs, compress to a target ratio, and
7+
measure whether the gold answer SURVIVES. TextCrusher (query-aware) vs truncate
8+
(keep-recent) vs random baselines. Mirrors kompress's published
9+
must_keep_recall (0.977 on its own labeled set).
10+
11+
Part B -- real-transcript fidelity: compress large text blocks from a real
12+
Claude Code transcript (ANONYMIZED), measuring ratio, speed, and salient-token
13+
retention (identifiers/numbers/errors -- the must-keep info in coding contexts).
14+
Only aggregate metrics are printed; raw content is never echoed.
15+
16+
Usage: python benchmarks/text_crusher_quality_eval.py [squad_dev.json] [transcript.jsonl]
17+
"""
18+
19+
from __future__ import annotations
20+
21+
import glob
22+
import json
23+
import os
24+
import random
25+
import re
26+
import sys
27+
import time
28+
29+
from headroom.transforms.text_crusher import TextCrusher
30+
31+
_SEG = re.compile(r"(?<=[.!?])\s+|\n+")
32+
_SALIENT = re.compile(
33+
r"\b(?:error|exception|fail(?:ed|ure)?|warning|traceback|assert|todo|fixme)\b"
34+
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"
35+
)
36+
37+
# --- anonymization (脱敏): scrub before any processing; never echo raw content ---
38+
_REDACT = [
39+
(re.compile(r"/Users/[^/\s]+"), "/Users/USER"),
40+
(re.compile(r"\b[\w.+-]+@[\w-]+\.[\w.-]+\b"), "EMAIL"),
41+
(re.compile(r"\b(?:sk|pk|ghp|gho|xox[baprs])-[A-Za-z0-9_-]{10,}\b"), "TOKEN"),
42+
(re.compile(r"\bBearer\s+[A-Za-z0-9._-]{10,}"), "Bearer TOKEN"),
43+
(re.compile(r"\b[A-Fa-f0-9]{40,}\b"), "HEX"),
44+
]
45+
46+
47+
def anon(t: str) -> str:
48+
for rx, rep in _REDACT:
49+
t = rx.sub(rep, t)
50+
return t
51+
52+
53+
def norm(s: str) -> str:
54+
return re.sub(r"\s+", " ", s.lower()).strip()
55+
56+
57+
def _segs(text: str) -> list[str]:
58+
return [s for s in _SEG.split(text) if s.strip()]
59+
60+
61+
def truncate_keep_last(text: str, ratio: float) -> str:
62+
segs = _segs(text)
63+
budget = int(sum(len(s) for s in segs) * ratio)
64+
kept: list[str] = []
65+
c = 0
66+
for s in reversed(segs):
67+
if c >= budget:
68+
break
69+
kept.append(s)
70+
c += len(s)
71+
return "\n".join(reversed(kept))
72+
73+
74+
def random_keep(text: str, ratio: float, seed: int) -> str:
75+
segs = _segs(text)
76+
idx = list(range(len(segs)))
77+
random.Random(seed).shuffle(idx)
78+
budget = int(sum(len(s) for s in segs) * ratio)
79+
kept: set[int] = set()
80+
c = 0
81+
for i in idx:
82+
if c >= budget:
83+
break
84+
kept.add(i)
85+
c += len(segs[i])
86+
return "\n".join(segs[i] for i in sorted(kept))
87+
88+
89+
def eval_squad(path: str, n: int = 200, n_distract: int = 40, ratio: float = 0.3, seed: int = 0):
90+
data = json.load(open(path))
91+
paras = [(p["context"], p["qas"]) for a in data["data"] for p in a["paragraphs"]]
92+
all_ctx = [c for c, _ in paras]
93+
examples = [
94+
(ctx, qas[0]["question"], qas[0]["answers"][0]["text"])
95+
for ctx, qas in paras
96+
if qas and qas[0]["answers"]
97+
]
98+
rnd = random.Random(seed)
99+
rnd.shuffle(examples)
100+
examples = examples[:n]
101+
tc = TextCrusher()
102+
hit = {"text_crusher": 0, "truncate": 0, "random": 0}
103+
tc_ratios: list[float] = []
104+
for gold_ctx, q, ans in examples:
105+
docs = rnd.sample(all_ctx, n_distract) + [gold_ctx]
106+
rnd.shuffle(docs)
107+
haystack = "\n\n".join(docs)
108+
a = norm(ans)
109+
out_tc = tc.compress(haystack, context=q, target_ratio=ratio).compressed
110+
hit["text_crusher"] += a in norm(out_tc)
111+
hit["truncate"] += a in norm(truncate_keep_last(haystack, ratio))
112+
hit["random"] += a in norm(random_keep(haystack, ratio, seed))
113+
tc_ratios.append(len(out_tc) / max(1, len(haystack)))
114+
nn = len(examples)
115+
print(
116+
f"\n=== Part A: SQuAD answer-retention (n={nn}, distractors={n_distract}, target_ratio={ratio}) ==="
117+
)
118+
print(
119+
f" TextCrusher (query-aware): {hit['text_crusher'] / nn:6.1%} answer survives compression"
120+
)
121+
print(f" Truncate (keep recent): {hit['truncate'] / nn:6.1%}")
122+
print(f" Random keep: {hit['random'] / nn:6.1%}")
123+
print(
124+
f" TextCrusher mean char-ratio: {sum(tc_ratios) / nn:.2f} (kept ~{sum(tc_ratios) / nn:.0%} of bytes)"
125+
)
126+
print(" reference: kompress published must_keep_recall = 0.977 (its own labeled set)")
127+
128+
129+
def _block_texts(jsonl_path: str, min_words: int, limit: int) -> list[str]:
130+
out: list[str] = []
131+
with open(jsonl_path) as fh:
132+
for line in fh:
133+
try:
134+
o = json.loads(line)
135+
except json.JSONDecodeError:
136+
continue
137+
m = o.get("message") or {}
138+
c = m.get("content")
139+
parts = (
140+
[c]
141+
if isinstance(c, str)
142+
else [
143+
p["text"] for p in c if isinstance(p, dict) and isinstance(p.get("text"), str)
144+
]
145+
if isinstance(c, list)
146+
else []
147+
)
148+
for t in parts:
149+
if len(t.split()) >= min_words:
150+
out.append(anon(t))
151+
if len(out) >= limit:
152+
break
153+
return out[:limit]
154+
155+
156+
def eval_transcript(jsonl_path: str, ratio: float = 0.4, min_words: int = 1500, limit: int = 40):
157+
blocks = _block_texts(jsonl_path, min_words, limit)
158+
if not blocks:
159+
print(
160+
f"\n=== Part B: no text blocks >= {min_words} words in {os.path.basename(jsonl_path)} ==="
161+
)
162+
return
163+
tc = TextCrusher()
164+
ratios: list[float] = []
165+
times: list[float] = []
166+
retentions: list[float] = []
167+
for b in blocks:
168+
sal_before = set(_SALIENT.findall(b))
169+
t0 = time.perf_counter()
170+
out = tc.compress(b, target_ratio=ratio).compressed
171+
times.append((time.perf_counter() - t0) * 1000)
172+
sal_after = set(_SALIENT.findall(out))
173+
retentions.append(len(sal_before & sal_after) / max(1, len(sal_before)))
174+
ratios.append(len(out.split()) / max(1, len(b.split())))
175+
n = len(blocks)
176+
print(
177+
f"\n=== Part B: real transcript fidelity (n={n} large blocks, anonymized, target_ratio={ratio}) ==="
178+
)
179+
print(f" mean token-ratio kept: {sum(ratios) / n:.2f}")
180+
print(f" mean speed: {sum(times) / n:.1f} ms/block")
181+
print(
182+
f" salient-token retention: {sum(retentions) / n:6.1%} (identifiers/numbers/errors kept)"
183+
)
184+
print(
185+
f" -> keeps salient info at {sum(retentions) / n:.0%} while dropping to {sum(ratios) / n:.0%} of tokens"
186+
)
187+
188+
189+
def eval_speed(scale_words: int = 250_000):
190+
# Reproducible throughput on a large synthetic prose block (no external data).
191+
text = " ".join(
192+
f"Sentence {i} discusses subsystem {i} and its failure mode {i % 7} in detail."
193+
for i in range(scale_words // 9)
194+
)
195+
nwords = len(text.split())
196+
tc = TextCrusher()
197+
t0 = time.perf_counter()
198+
out = tc.compress(text, target_ratio=0.3)
199+
ms = (time.perf_counter() - t0) * 1000
200+
print(f"\n=== Part C: speed (synthetic, {nwords:,} words, fully reproducible) ===")
201+
print(f" TextCrusher compress: {ms:.0f} ms ({nwords / max(ms / 1000, 1e-6):,.0f} words/sec)")
202+
print(f" kept ratio: {out.compressed_tokens / max(1, out.original_tokens):.2f}")
203+
print(" reference: kompress (ModernBERT ONNX) ~272s for ~1M tokens (measured, query-blind)")
204+
print(" -> fast-vs-slow CONTRAST, not a same-input side-by-side run")
205+
206+
207+
if __name__ == "__main__":
208+
eval_speed()
209+
squad = sys.argv[1] if len(sys.argv) > 1 else "/tmp/squad_dev.json"
210+
tx = sys.argv[2] if len(sys.argv) > 2 else None
211+
if os.path.exists(squad):
212+
eval_squad(squad)
213+
else:
214+
print(f"SQuAD not found at {squad}; skipping Part A")
215+
if tx is None:
216+
found = glob.glob(os.path.expanduser("~/.claude/projects/*headroom*/*.jsonl"))
217+
tx = max(found, key=os.path.getsize) if found else None
218+
if tx and os.path.exists(tx):
219+
eval_transcript(tx)
220+
else:
221+
print("no transcript jsonl found; skipping Part B")

crates/headroom-core/src/transforms/mod.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ pub mod safety;
2929
pub mod search_compressor;
3030
pub mod smart_crusher;
3131
pub mod tag_protector;
32+
pub mod text_crusher;
3233
pub mod unidiff_detector;
3334

3435
pub use content_detector::{
@@ -61,4 +62,5 @@ pub use search_compressor::{
6162
SearchCompressorStats, SearchMatch,
6263
};
6364
pub use tag_protector::{is_known_html_tag, protect_tags, restore_tags, ProtectStats};
65+
pub use text_crusher::{TextCrusher, TextCrusherConfig, TextCrusherResult};
6466
pub use unidiff_detector::{detect_diff, is_diff};
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
//! TextCrusher configuration (Phase 2, #1171).
2+
//!
3+
//! Mirrors the Python `TextCrusherConfig`. Weights and thresholds are tuning
4+
//! knobs for the recency + relevance + salience scoring.
5+
6+
#[derive(Debug, Clone)]
7+
pub struct TextCrusherConfig {
8+
/// Keep roughly this fraction of characters.
9+
pub target_ratio: f64,
10+
pub w_recency: f64,
11+
pub w_relevance: f64,
12+
pub w_salience: f64,
13+
/// Segments shorter than this are de-prioritized (× 0.25).
14+
pub min_segment_chars: usize,
15+
/// Skip a candidate when this fraction of its word-shingles is already
16+
/// covered by kept segments (near-duplicate suppression).
17+
pub near_dup_threshold: f64,
18+
/// Below this many segments, pass through unchanged (nothing to gain).
19+
pub min_segments_for_crush: usize,
20+
}
21+
22+
impl Default for TextCrusherConfig {
23+
fn default() -> Self {
24+
TextCrusherConfig {
25+
target_ratio: 0.5,
26+
w_recency: 1.0,
27+
w_relevance: 2.0,
28+
w_salience: 1.5,
29+
min_segment_chars: 12,
30+
near_dup_threshold: 0.85,
31+
min_segments_for_crush: 6,
32+
}
33+
}
34+
}

0 commit comments

Comments
 (0)