headroom/tests/test_kompress_must_keep.py
Tejas Chopra 8884d87378
fix(transforms): stop compression garbling mixed subagent output (#3286)
## The report

A user's model called compressed subagent output "too garbled to use"
and burned CCR retrievals to reconstruct it — **not** because it needed
more context. One retrieval returned nothing but the Claude Code harness
sanitizer banner, reported as `original_item_count: 33,
compressed_item_count: 25`.

Root-cause chain (verified by reproduction): the harness prepends a
bracket-delimited banner (`[harness: ... you.]` — exactly 33
whitespace-delimited words) and neutralizes `<` → `<\`. Headroom's
mixed-content splitter typed the banner as JSON (bracket balance, no
validation) → SmartCrusher couldn't parse it → the fallback chain fed it
to lossy Kompress → Kompress word-dropped the banner 33→25 and stored it
behind a retrieval hash. Meanwhile tabular sections rendered as
quote-wrapped JSON-string blobs with `\n` as two-character escapes, and
`ensure_ascii=True` boundaries turned the output's unicode (`→ └ ✓`)
into `\uXXXX` soup. The model reasonably concluded the output was
garbled.

## Fixes

1. **`split_into_sections` validates JSON before typing a block
`JSON_ARRAY`** — same validation its own mixed-content gate
(`_has_valid_json_block_with_text`) has always used. Tag-protection
placeholders, which self-isolated only by accident of that bug
(`{{HEADROOM_TAG_N}}` bracket-balances), are now isolated explicitly via
a new `isolate=` parameter fed by the router; contiguous prose fragments
re-coalesce so the `\n\n` reassembly stops doubling newlines in
uncompressed prose.
2. **Kompress gets a real floor: `min_input_words = 64`**
(config-tunable, clamped at the historical 10), applied on the
in-process, batch, apply, and remote paths. Below it, lossy
word-dropping is a net loss — the retrieval marker alone is ~20 words —
and short blocks are disproportionately instruction-like.
3. **The mixed path unwraps SmartCrusher's whole-array CSV render** when
it comes back as a bare JSON string, splicing raw readable lines into
the text instead of a quoted escape blob.
4. **`ensure_ascii=False` at model-visible boundaries**: MCP
retrieve/stats responses and the audit-safe splice reserialization
(which now also matches serde_json's non-escaping behavior).
5. **Kompress honesty**: the marker says `N words compressed to M`
(shared `ccr_retrieval_marker` helper, unit-tested), and
`store_kompress_in_ccr` no longer writes word counts into the store's
*item count* fields — token counts already carry the size story.

The upstream trigger (the harness's `<` → `<\` neutralization corrupting
JSON semantics) is not Headroom's to fix, but with #1 and #2 the banner
now passes through byte-intact and nothing lossy touches it.

## Testing

- New `tests/test_garbled_compression_fixes.py` (12 tests) pins every
fix, including an end-to-end router pass over a reconstructed
harness-sanitized fixture asserting the banner survives byte-identical
and no `\uXXXX` appears.
- Existing small-fixture kompress/router tests updated to set
`min_input_words=10` explicitly (they test other mechanics; fixtures sit
under the new production floor by design).
- Affected sweep (`-k "compress or ccr or crusher or router or mixed or
kompress or hermes"`, ~2.9k tests): green apart from order-dependent
flakes that shift identity between runs (deepseek tokenizer `AutoConfig`
import, hermes/proxy-ccr) — each passes standalone and in direct
combination with the new tests; the full CI shards are the authoritative
check.
- ruff 0.16.3 `check` + `format --check` clean.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

https://claude.ai/code/session_01EWKCmcH47hvvoQ35wftXhE

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
2026-08-27 13:22:26 +05:30

137 lines
5 KiB
Python

"""Tests for the must-keep token override in kompress_compressor."""
from __future__ import annotations
import os
from headroom.transforms import kompress_compressor as kc
from headroom.transforms.kompress_compressor import (
_KOMPRESS_MUST_KEEP_ENV,
_KOMPRESS_MUST_KEEP_RE,
KompressCompressor,
KompressConfig,
)
class _Enc(dict):
def word_ids(self, batch_index=0):
return self["_word_ids"][batch_index]
class _Tok:
def __call__(self, chunk_words, **kw):
if chunk_words and isinstance(chunk_words[0], list):
batch_words = chunk_words
else:
batch_words = [chunk_words]
return _Enc(
input_ids=[[0] * len(words) for words in batch_words],
attention_mask=[[1] * len(words) for words in batch_words],
_word_ids=[list(range(len(words))) for words in batch_words],
)
class _Model:
def get_keep_mask(self, input_ids, attention_mask):
return [[idx == 0 for idx, _ in enumerate(row)] for row in input_ids]
def get_scores(self, input_ids, attention_mask):
return [[1.0 if idx == 0 else 0.0 for idx, _ in enumerate(row)] for row in input_ids]
def _install_fake_kompress(monkeypatch):
monkeypatch.setattr(kc, "_load_kompress", lambda *a, **k: (_Model(), _Tok(), "onnx"))
monkeypatch.setattr(kc, "_model_device_type", lambda *a, **k: "cpu")
class TestMustKeepRegex:
def test_numbers(self):
assert _KOMPRESS_MUST_KEEP_RE.search("42")
assert _KOMPRESS_MUST_KEEP_RE.search("3.14")
assert _KOMPRESS_MUST_KEEP_RE.search("0x7fff2038")
assert not _KOMPRESS_MUST_KEEP_RE.search("word0")
def test_allcaps(self):
assert _KOMPRESS_MUST_KEEP_RE.search("SIGILL")
assert _KOMPRESS_MUST_KEEP_RE.search("HTTP")
assert _KOMPRESS_MUST_KEEP_RE.search("EOF")
def test_dotted_paths(self):
assert _KOMPRESS_MUST_KEEP_RE.search("libsystem_kernel.dylib")
assert _KOMPRESS_MUST_KEEP_RE.search("torch.nn")
def test_unix_paths(self):
assert _KOMPRESS_MUST_KEEP_RE.search("/usr/lib/python3")
assert _KOMPRESS_MUST_KEEP_RE.search("/workspace/ultrawhale")
def test_extensions(self):
assert _KOMPRESS_MUST_KEEP_RE.search("model.py")
assert _KOMPRESS_MUST_KEEP_RE.search("weights.so")
def test_flags(self):
assert _KOMPRESS_MUST_KEEP_RE.search("--verbose")
assert _KOMPRESS_MUST_KEEP_RE.search("-n")
def test_camelcase(self):
assert _KOMPRESS_MUST_KEEP_RE.search("IndexError")
assert _KOMPRESS_MUST_KEEP_RE.search("EXC_BAD_INSTRUCTION")
def test_plain_words_not_matched(self):
assert not _KOMPRESS_MUST_KEEP_RE.search("the")
assert not _KOMPRESS_MUST_KEEP_RE.search("process")
assert not _KOMPRESS_MUST_KEEP_RE.search("raised")
class TestMustKeepEnvVar:
def test_env_var_name(self):
assert _KOMPRESS_MUST_KEEP_ENV == "HEADROOM_KOMPRESS_MUST_KEEP"
def test_env_var_default_is_enabled(self, monkeypatch):
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
assert os.environ.get(_KOMPRESS_MUST_KEEP_ENV, "1") != "0"
def test_env_var_can_disable(self, monkeypatch):
monkeypatch.setenv(_KOMPRESS_MUST_KEEP_ENV, "0")
assert os.environ.get(_KOMPRESS_MUST_KEEP_ENV, "1") == "0"
class TestMustKeepCompression:
def test_compress_keeps_must_keep_word_when_model_drops_it(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
compressor = KompressCompressor(KompressConfig(enable_ccr=False, min_input_words=10))
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *a, **k: False)
result = compressor.compress(
"alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"
)
assert result.compressed.split() == ["alpha", "0x7fff2038"]
def test_compress_can_disable_must_keep_override(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.setenv(_KOMPRESS_MUST_KEEP_ENV, "0")
compressor = KompressCompressor(KompressConfig(enable_ccr=False, min_input_words=10))
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *a, **k: False)
result = compressor.compress(
"alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"
)
assert result.compressed.split() == ["alpha"]
def test_compress_batch_keeps_must_keep_word_when_score_is_low(self, monkeypatch):
_install_fake_kompress(monkeypatch)
monkeypatch.delenv(_KOMPRESS_MUST_KEEP_ENV, raising=False)
compressor = KompressCompressor(KompressConfig(enable_ccr=False, min_input_words=10))
monkeypatch.setattr(compressor, "_should_use_sequential_fallback", lambda: False)
[result] = compressor.compress_batch(
["alpha beta gamma delta epsilon zeta eta theta iota kappa 0x7fff2038 omega"],
batch_size=8,
)
assert result.compressed.split() == ["alpha", "0x7fff2038"]