akadmin/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-27 14:17:10 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(compression): add audit-safe mode with protected pattern matching (#1899)

Browse source 
## Description

`SmartCrusher.crush_array_json` (`headroom/transforms/smart_crusher.py`)
selects rows to keep using statistical signals such as variance,
structural anomaly, and position. It has no concept of "this row is
audit/compliance-relevant and must stay visible in the prompt." A rare
row, such as a leakage flag, compliance marker, or non-standard failure
line, can be sampled out like any routine row, or moved behind an opaque
`<<ccr:HASH ...>>` retrieval marker the model has no reason to ask for.
In audit, SRE, and quant-falsification workloads, rare rows are
frequently the most important evidence, so silent disappearance is a
real safety issue rather than only a lossy-compression tradeoff.

This adds an opt-in `audit_safe` mode to `SmartCrusher`:

- `SmartCrusherConfig(audit_safe=True, protected_patterns=[...],
fail_closed_on_protected_loss=True)`
- Rows are scanned for pattern matches, string or regex, against each
row's canonical JSON text before compression runs.
- After compression, any protected row missing from the output is
spliced back in verbatim, whether it was dropped by the statistical
selector or left only behind a CCR marker.
- A verification pass re-counts protected-row survivors after splicing.
If the count is still short, the crusher fails closed and returns the
original, uncompressed content instead of shipping a result with fewer
protected matches than the input had. Setting
`fail_closed_on_protected_loss=False` ships the best-effort spliced
result with a logged warning instead.

Protection applies on both `crush_array_json`, the dict-shaped API used
by direct callers and the CCR retrieval flow, and
`_smart_crush_content`, the tuple-shaped API `apply()` actually calls
for every compressed tool/tool_result message. It is live on the real
tool-output compression path.

Scope: this covers JSON-array-shaped content routed through
`SmartCrusher`, the common case for tool outputs such as API results,
log lines, and DB rows returned as JSON. Raw CSV/plain-text content
compressed by other transforms, including Kompress and log/tabular
compressors, is out of scope for this PR; `protected_patterns` only has
row structure to match against when the content is or renders to a JSON
array.

Closes #1705

## 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
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- `headroom/transforms/smart_crusher.py`: added `audit_safe`,
`protected_patterns`, and `fail_closed_on_protected_loss` fields to
`SmartCrusherConfig`; these stay Python-side and do not reach the Rust
config because this is post-processing around existing Rust-backed
compression.
- Added `_compile_protected_patterns`, `_canon`,
`_row_matches_protected`, `_scan_protected_rows`, and
`_splice_missing_protected` as the shared scan/match/splice primitives.
- Added `_apply_audit_safe_protection` for dict-shaped
`crush_array_json` results and `_apply_audit_safe_protection_to_content`
for tuple-shaped `_smart_crush_content` / `apply()` results. Both splice
missing protected rows back in, then verify and fail closed or warn on
residual loss.
- Wired both `crush_array_json` and `_smart_crush_content` to scan for
protected rows before compression and apply protection after.
- `CHANGELOG.md`: added an `Unreleased / Features` entry.
- Default `audit_safe=False`, so existing callers keep current behavior.
A regression test compares a configured-but-disabled crusher's output
byte-for-byte against an unconfigured one.

## Testing

- [x] Unit tests pass (`uv run pytest
tests/test_transforms/test_smart_crusher_audit_safe.py`)
- [x] Linting passes (`uv run ruff check .`)
- [x] Type checking passes (`uv run mypy
headroom/transforms/smart_crusher.py`)
- [x] New tests added for new functionality when applicable
- [x] Manual testing performed

### Test Output

```text
$ uv run pytest tests/ -k "(smart_crusher or crush or audit) and not test_optimizer_not_called_in_audit_mode" --no-header -q
...
tests\test_transforms\test_smart_crusher_audit_safe.py ...........                                               [ 65%]
...
169 passed, 10 skipped, 8176 deselected, 1 warning in 21.84s

$ uv run ruff check .
All checks passed!

$ uv run mypy headroom/transforms/smart_crusher.py
Success: no issues found in 1 source file
```

`-k` excludes `test_optimizer_not_called_in_audit_mode`
(`tests/test_cache/test_client_integration.py`), a pre-existing,
unrelated Windows temp-path failure in SQLite storage init that
reproduces identically on a clean `origin/main` checkout with none of
this PR's changes applied; it matched the `-k audit` filter by name
coincidence only.

## Real Behavior Proof

- Environment: Windows, Python 3.12 via uv-managed venv,
`headroom._core` built locally via `maturin` / cargo 1.95.0, no LLM
provider needed because this is pure transform-layer behavior.
- Exact command / steps: Ran `uv run pytest tests/ -k "(smart_crusher or
crush or audit) and not test_optimizer_not_called_in_audit_mode"
--no-header -q`, `uv run ruff check .`, and `uv run mypy
headroom/transforms/smart_crusher.py`; also exercised the audit-safe
tests that build a 62-row JSON array with two `AUDIT_FLAG` rows, run it
through `SmartCrusher(SmartCrusherConfig(audit_safe=True,
protected_patterns=["AUDIT_FLAG"]), with_compaction=False)` via both
`crush_array_json` and `Transform.apply()` over a synthetic tool
message, parse the compressed output back to JSON, and drive the
splice/verify/fail-closed helper paths with engineered row-drop and
forced-mismatch scenarios.
- Observed result: Protected rows are present in the compressed output
in every tested scenario; the fail-closed branch returns the original
content byte-for-byte with `strategy_info == "audit_safe:fail_closed"`
when verification detects residual loss; `audit_safe=False` produces
output byte-identical to a crusher with no audit-safe configuration.
- Not tested: Raw CSV/plain-text tool output compressed via
non-SmartCrusher transforms, including Kompress and log/tabular
compressors, is out of scope. Top-level `headroom.compress()` /
`CompressConfig` wiring for `audit_safe` and `protected_patterns` is a
natural follow-up and is not included here.

## 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
- [ ] 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

No user-facing docs were updated because I did not find an existing
`SmartCrusherConfig` field reference doc to extend. The top-level
`compress()` / `CompressConfig` wiring mentioned in "Not tested" is a
reasonable immediate follow-up if this mechanism is the right shape.
This commit is contained in:
Rod Boev 2026-07-09 09:39:35 -04:00 • committed by GitHub
parent 361adcd1a0
commit bb112dd176
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 663 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -16,6 +16,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
non-durable on POSIX. Best-effort — a no-op on Windows and virtual
filesystems where directory fsync is unsupported.
### Features
* **transforms:** add opt-in audit-safe mode to `SmartCrusher` — `SmartCrusherConfig(audit_safe=True, protected_patterns=[...], fail_closed_on_protected_loss=True)`. Rows matching a protected pattern are scanned before JSON-array compression and guaranteed to survive the compressed output verbatim afterward (never dropped, never replaced by an opaque `<<ccr:...>>` marker only). Applies on both the `crush_array_json` convenience API and the `_smart_crush_content` path `apply()` uses for real tool-output compression. If a protected row still can't be preserved after the splice-back pass, the crusher fails closed by returning the original uncompressed content (or ships a best-effort result with a warning when `fail_closed_on_protected_loss=False`). Default is `audit_safe=False` — no behavior change for existing callers ([#1705](https://github.com/chopratejas/headroom/issues/1705)).
### Changed
* **telemetry:** anonymous usage telemetry is now **opt-in** (off by default) instead of opt-out. Nothing is collected or sent unless you set `HEADROOM_TELEMETRY=on` or pass `--telemetry` to `headroom proxy` / `headroom install apply`. `is_telemetry_enabled()` is fail-closed — only explicit on-values (`on`/`true`/`1`/`yes`/`enable`/`enabled`) enable it; unset, empty, or unrecognized values stay disabled. The existing `--no-telemetry` flag and `HEADROOM_TELEMETRY=off` remain accepted for back-compat, and install manifests now write the `HEADROOM_TELEMETRY` value explicitly so generated deployments are unambiguous.

258
headroom/transforms/smart_crusher.py
View file

@ -46,6 +46,8 @@ from __future__ import annotations
import json
import logging
import os
import re
from collections import Counter
from dataclasses import dataclass
from typing import Any
@ -210,6 +212,29 @@ class SmartCrusherConfig:
compaction_min_buckets: int = 2
compaction_max_buckets: int = 8
# ─── Audit-safe mode (#1705) ───────────────────────────────────────
# Opt-in. `crush_array_json`'s row selection (Rust-side statistical
# sampling) has no concept of "this row must not disappear from the
# prompt" — a rare compliance/audit-trail row can be sampled out or
# replaced by a `<<ccr:...>>` retrieval marker like any other row.
# When `audit_safe=True` and `protected_patterns` is non-empty,
# `crush_array_json` scans rows for pattern matches before
# compression, then guarantees matched rows survive in the
# compressed output verbatim — not dropped, not marker-only. This
# field never reaches the Rust config (`_rust_cfg_kwargs` excludes
# it); it's pure Python post-processing around the Rust call.
audit_safe: bool = False
# Strings or regexes. A row is "protected" if any pattern matches
# its canonical JSON text (`json.dumps(row, sort_keys=True)`).
protected_patterns: list[str] | None = None
# If protected rows still can't be fully preserved after the
# splice-back pass (defensive — should only trip on internal
# bugs), fail closed by returning the original, uncompressed array
# instead of a result with fewer protected-row matches than the
# input had. When False, ship the best-effort result with a
# logged warning instead of refusing to compress.
fail_closed_on_protected_loss: bool = True
# ─── Rust-backed SmartCrusher ─────────────────────────────────────────────
@ -252,6 +277,18 @@ class SmartCrusher(Transform):
cfg = config or SmartCrusherConfig()
self.config = cfg
self._with_compaction = with_compaction
# Audit-safe mode (#1705). getattr fallbacks: callers may pass
# the SDK-side `headroom.config.SmartCrusherConfig`, which
# doesn't carry these fields — defaults to disabled, the safe
# choice (no behavior change for callers who don't opt in).
self._audit_safe = bool(getattr(cfg, "audit_safe", False))
self._fail_closed_on_protected_loss = bool(
getattr(cfg, "fail_closed_on_protected_loss", True)
)
self._protected_patterns = self._compile_protected_patterns(
getattr(cfg, "protected_patterns", None)
)
# Strict lossless mode. An explicit `lossless_only=` kwarg wins
# over the config field, so callers can flip it without rebuilding
# a whole config. `crush(..., lossless_only=...)` overrides again
@ -496,6 +533,206 @@ class SmartCrusher(Transform):
strategy=r.strategy,
)
# ─── Audit-safe protection (#1705) ─────────────────────────────────
#
# `crush_array_json`'s Rust-side row selection is purely statistical
# (variance, anomaly, position) — it has no notion of "this row is
# legally/compliance-significant and must stay visible in the
# prompt." Audit-safe mode bolts that on in Python: scan for
# pattern matches before compression, then guarantee matched rows
# survive the compressed output (never dropped, never marker-only).
@staticmethod
def _compile_protected_patterns(patterns: list[str] | None) -> list[re.Pattern[str]]:
"""Compile `protected_patterns` once at construction time.
A pattern that fails to compile is a caller bug, not something
to swallow — silently treating an invalid regex as "no rows
protected" would defeat the entire point of audit-safe mode
(rows the caller believes are protected wouldn't be).
"""
if not patterns:
return []
compiled = []
for p in patterns:
try:
compiled.append(re.compile(p))
except re.error as e:
raise ValueError(
f"SmartCrusher: invalid protected_patterns regex {p!r}: {e}"
) from e
return compiled
@staticmethod
def _canon(item: Any) -> str:
"""Canonical JSON text for a row.
Used both for protected-pattern matching and for identity
comparison across the crush boundary — kept rows are
re-serialized by Rust, so rows are matched by content, not by
Python object identity.
"""
return json.dumps(item, sort_keys=True, default=str)
def _row_matches_protected(self, item: Any) -> bool:
text = self._canon(item)
return any(p.search(text) for p in self._protected_patterns)
def _scan_protected_rows(self, items_json_or_content: str) -> list[Any]:
"""Rows matching any `protected_patterns` entry, or `[]` when
audit-safe mode is off, no patterns are configured, or the
input doesn't parse as a JSON array (nothing row-shaped to
protect — e.g. raw CSV/log text, out of scope for this mode)."""
if not (self._audit_safe and self._protected_patterns):
return []
try:
parsed = json.loads(items_json_or_content)
except (json.JSONDecodeError, ValueError):
return []
if not isinstance(parsed, list):
return []
return [item for item in parsed if self._row_matches_protected(item)]
def _splice_missing_protected(
self, protected: list[Any], kept: list[Any]
) -> tuple[list[Any], int]:
"""Append any `protected` row missing from `kept` (identity by
canonical JSON, multiplicity-aware via `Counter` so duplicate
protected rows are each accounted for individually).
Returns `(kept_with_splice, lost_count)`. `lost_count` is
almost always 0 after splicing — it stays non-zero only when
something structural prevents the appended row from being
recognized as a survivor (defensive; see call sites).
"""
available = Counter(self._canon(item) for item in kept)
missing = []
for item in protected:
key = self._canon(item)
if available[key] > 0:
available[key] -= 1
else:
missing.append(item)
if missing:
kept = kept + missing
surviving = Counter(self._canon(item) for item in kept)
needed = Counter(self._canon(item) for item in protected)
lost = sum(max(0, count - surviving[key]) for key, count in needed.items())
return kept, lost
def _apply_audit_safe_protection(
self,
protected: list[Any],
original_items_json: str,
result: dict[str, Any],
) -> dict[str, Any]:
"""Guarantee every row in `protected` survives in `result["items"]`.
Two phases:
1. Splice — any protected row missing from the compressed
output (statistically sampled out, or moved behind an
opaque `<<ccr:...>>` retrieval marker) is appended back
into `items` verbatim.
2. Verify — re-count protected-row survivors after splicing.
If the count is still short (defensive: should only trip
on an internal bug, e.g. the lossless table path rendering
rows into a non-addressable CSV blob), fail closed by
returning the original uncompressed array, or ship the
spliced result with a logged warning, per
`fail_closed_on_protected_loss`.
"""
kept_json = result.get("items")
try:
kept = json.loads(kept_json) if isinstance(kept_json, str) else list(kept_json or [])
except (json.JSONDecodeError, ValueError):
kept = []
before_count = len(kept)
kept, lost = self._splice_missing_protected(protected, kept)
if len(kept) != before_count:
# Only reserialize when something was actually spliced in —
# an unmodified `kept` stays byte-identical to Rust's output
# (Python's `json.dumps` and serde_json don't necessarily
# agree on e.g. non-ASCII escaping).
result = dict(result)
result["items"] = json.dumps(kept)
if not lost:
return result
msg = (
f"SmartCrusher audit_safe: {lost} protected row(s) could not be "
f"preserved through compression (pattern match count decreased "
f"even after splicing back missing rows)."
)
if self._fail_closed_on_protected_loss:
logger.warning("%s Failing closed: returning original uncompressed.", msg)
return {
"items": original_items_json,
"ccr_hash": None,
"dropped_summary": "",
"strategy_info": "audit_safe:fail_closed",
"compacted": None,
"compaction_kind": None,
}
logger.warning(
"%s fail_closed_on_protected_loss=False — shipping best-effort result.",
msg,
)
return result
def _apply_audit_safe_protection_to_content(
self,
protected: list[Any],
original_content: str,
crushed: str,
was_modified: bool,
info: str,
) -> tuple[str, bool, str]:
"""Guarantee protected rows survive `_smart_crush_content`'s
output — the tuple-shaped API `apply()` uses for real
tool-output compression (`crush_array_json` is the dict-shaped
API used by direct/test callers and the CCR retrieval flow;
`apply()` never calls it).
`crushed` may be a JSON array string (the common shape for the
lossy row-drop and passthrough paths) — spliced exactly like
`crush_array_json`. Anything else (lossless CSV/table
rendering, an opaque marker string) has no row structure left
to splice into, so verification falls back to counting
protected-pattern matches in the raw text before vs. after.
"""
try:
parsed = json.loads(crushed)
except (json.JSONDecodeError, ValueError):
parsed = None
if isinstance(parsed, list):
kept, lost = self._splice_missing_protected(protected, parsed)
# Only reserialize when something was actually spliced in —
# see the matching comment in `_apply_audit_safe_protection`.
candidate = json.dumps(kept) if len(kept) != len(parsed) else crushed
else:
lost = sum(
max(0, len(p.findall(original_content)) - len(p.findall(crushed)))
for p in self._protected_patterns
)
candidate = crushed
if not lost:
return candidate, was_modified, info
msg = f"SmartCrusher audit_safe: {lost} protected pattern match(es) lost in compression."
if self._fail_closed_on_protected_loss:
logger.warning("%s Failing closed: returning original content uncompressed.", msg)
return original_content, False, "audit_safe:fail_closed"
logger.warning(
"%s fail_closed_on_protected_loss=False — shipping best-effort result.",
msg,
)
return candidate, was_modified, info
def crush_array_json(
self,
items_json: str,
@ -511,8 +748,18 @@ class SmartCrusher(Transform):
Used by tests and by the proxy's CCR retrieval flow when it needs
the hash directly rather than parsing it out of a prompt marker.
When this instance is configured with `audit_safe=True` and a
non-empty `protected_patterns`, rows matching any pattern are
scanned *before* compression and guaranteed to survive in the
returned `items` — see `_apply_audit_safe_protection`.
"""
protected = self._scan_protected_rows(items_json)
result: dict[str, Any] = self._rust.crush_array_json(items_json, query, bias)
if protected:
result = self._apply_audit_safe_protection(protected, items_json, result)
# Row-drop case: Rust returns the structured `ccr_hash` and has
# already stashed the canonical in its own store. Mirror that
# entry into the Python compression_store keyed by the same
@ -600,8 +847,19 @@ class SmartCrusher(Transform):
threaded through to TOIN's per-tool learning records; if no
tool name is available (e.g. the legacy pipeline doesn't have
one in scope) the recording uses content-based signature only.
This is the path `apply()` actually calls for every compressed
tool/tool_result message — so it's also where audit-safe mode
(`audit_safe=True` + `protected_patterns`, #1705) has to hook
in to matter in production, not just via the `crush_array_json`
convenience API. See `_apply_audit_safe_protection_to_content`.
"""
protected = self._scan_protected_rows(content)
crushed, was_modified, info = self._rust.smart_crush_content(content, query_context, bias)
if protected:
crushed, was_modified, info = self._apply_audit_safe_protection_to_content(
protected, content, crushed, was_modified, info
)
# Same passthrough filter as `crush()` — re-canonicalization of
# JSON whitespace can flip `was_modified=True` even when the
# `info` field reports `passthrough` and no compression happened.

401
tests/test_transforms/test_smart_crusher_audit_safe.py Normal file
View file

@ -0,0 +1,401 @@
"""Audit-safe mode for SmartCrusher.crush_array_json (#1705).
`crush_array_json`'s row selection is purely statistical (variance,
anomaly, position) — it has no concept of "this row is a rare,
audit-relevant record that must stay visible in the prompt." A
compliance-significant row can be sampled out, or replaced by an
opaque `<<ccr:...>>` retrieval marker, exactly like any other row.
Audit-safe mode (`audit_safe=True` + `protected_patterns`) bolts
protection onto the existing Rust-backed compression without touching
the Rust selection logic: scan rows for pattern matches before
compression, then guarantee matched rows survive the compressed
output verbatim afterward — never dropped, never marker-only.
These tests exercise:
- Default (audit_safe=False) — zero behavior change, even if
`protected_patterns` happens to be set.
- The splice-back mechanism directly (deterministic, doesn't depend on
Rust's row-selection outcome for a given input).
- End-to-end through `crush_array_json` with a real lossy compression.
- The fail-closed / warn-and-ship-best-effort fork when verification
still finds a shortfall after splicing.
- Loud failure on an invalid regex in `protected_patterns`.
"""
from __future__ import annotations
import json
import pytest
def _build_extension() -> None:
try:
from headroom._core import SmartCrusher # noqa: F401
except ImportError:
pytest.skip(
"headroom._core not built — run `bash scripts/build_rust_extension.sh`",
allow_module_level=True,
)
_build_extension()
def test_audit_safe_disabled_by_default_no_behavior_change() -> None:
"""`protected_patterns` set but `audit_safe` left at its False
default → identical output to a crusher with no audit-safe config
at all. The flag gates the whole feature, not just the presence of
patterns."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
items = [{"id": i, "status": "ok"} for i in range(50)]
items_json = json.dumps(items)
baseline = SmartCrusher(SmartCrusherConfig(), with_compaction=False)
configured_but_off = SmartCrusher(
SmartCrusherConfig(protected_patterns=["ok"]), with_compaction=False
)
r1 = baseline.crush_array_json(items_json)
r2 = configured_but_off.crush_array_json(items_json)
assert r1["items"] == r2["items"]
assert r1["ccr_hash"] == r2["ccr_hash"]
assert r1["strategy_info"] == r2["strategy_info"]
def test_audit_safe_with_no_patterns_is_a_no_op() -> None:
"""`audit_safe=True` but `protected_patterns` empty/None → nothing
to protect, output unchanged from the unguarded crusher."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
items = [{"id": i, "status": "ok"} for i in range(50)]
items_json = json.dumps(items)
baseline = SmartCrusher(SmartCrusherConfig(), with_compaction=False)
audit_safe_no_patterns = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=None),
with_compaction=False,
)
r1 = baseline.crush_array_json(items_json)
r2 = audit_safe_no_patterns.crush_array_json(items_json)
assert r1["items"] == r2["items"]
assert r1["ccr_hash"] == r2["ccr_hash"]
def test_splice_back_restores_a_dropped_protected_row() -> None:
"""Direct unit test of `_apply_audit_safe_protection`: given a
`result["items"]` that's missing a protected row (as if the Rust
row-drop path had sampled it out), the method appends it back and
reports no loss.
Deterministic — doesn't depend on Rust's actual sampling decision
for a given input, only on the splice/verify logic this PR adds.
"""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
protected_row = {"id": 7, "note": "AUDIT_FLAG: rare compliance event"}
kept_without_protected = [{"id": i, "status": "ok"} for i in range(5)]
fake_result = {
"items": json.dumps(kept_without_protected),
"ccr_hash": "deadbeefcafe",
"dropped_summary": "<<ccr:deadbeefcafe 10_rows_offloaded>>",
"strategy_info": "smart_sample",
"compacted": None,
"compaction_kind": None,
}
out = crusher._apply_audit_safe_protection(
[protected_row], json.dumps(kept_without_protected + [protected_row]), fake_result
)
kept = json.loads(out["items"])
assert protected_row in kept
assert len(kept) == len(kept_without_protected) + 1
# Everything else about the result (ccr_hash, marker) is untouched —
# splicing doesn't erase the CCR pointer for the rows that really
# were dropped, it only guarantees the protected one is inline too.
assert out["ccr_hash"] == "deadbeefcafe"
def test_audit_safe_preserves_protected_rows_end_to_end() -> None:
"""Full path through `crush_array_json`: a real lossy compression
runs, and every protected row is present in the output afterward,
regardless of what the statistical row-selection decided on its
own."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
background = [{"id": i, "status": "ok"} for i in range(60)]
protected_rows = [
{"id": 25, "status": "ok", "note": "AUDIT_FLAG: rare compliance event A"},
{"id": 35, "status": "ok", "note": "AUDIT_FLAG: rare compliance event B"},
]
items = (
background[:25]
+ [protected_rows[0]]
+ background[25:35]
+ [protected_rows[1]]
+ background[35:]
)
items_json = json.dumps(items)
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
result = crusher.crush_array_json(items_json)
kept = json.loads(result["items"])
for row in protected_rows:
assert row in kept, f"protected row missing from compressed output: {row!r}"
# The array as a whole still compressed — audit-safe protection
# isn't a blanket opt-out of compression, only a guarantee for the
# specific protected rows.
assert len(kept) < len(items)
def test_audit_safe_fails_closed_when_verification_still_finds_loss(monkeypatch) -> None:
"""Defensive path: if the post-splice verification still finds a
shortfall (simulated here — normal splicing always succeeds, so we
force the mismatch by making `_canon` non-idempotent), and
`fail_closed_on_protected_loss=True` (the default), the whole
array is returned unmodified instead of shipping a result with
fewer protected-row matches than the input had."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
# Force every `_canon` call to return a fresh, never-repeating
# value so the Counter-based matching in both the splice and the
# verify phase can never line up — modeling an internal
# inconsistency the verification step exists to catch.
counter = iter(range(10_000))
monkeypatch.setattr(
SmartCrusher, "_canon", staticmethod(lambda item: f"unique-{next(counter)}")
)
protected_row = {"id": 1, "note": "AUDIT_FLAG"}
items_json = json.dumps([protected_row])
fake_result = {
"items": json.dumps([]),
"ccr_hash": "aaaa",
"dropped_summary": "<<ccr:aaaa 1_rows_offloaded>>",
"strategy_info": "smart_sample",
"compacted": None,
"compaction_kind": None,
}
out = crusher._apply_audit_safe_protection([protected_row], items_json, fake_result)
assert out["items"] == items_json
assert out["ccr_hash"] is None
assert out["strategy_info"] == "audit_safe:fail_closed"
def test_audit_safe_ships_best_effort_when_fail_closed_disabled(monkeypatch) -> None:
"""Same forced-mismatch scenario, but `fail_closed_on_protected_loss
=False` — ship the spliced best-effort result (with a logged
warning) instead of refusing to compress."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
crusher = SmartCrusher(
SmartCrusherConfig(
audit_safe=True,
protected_patterns=["AUDIT_FLAG"],
fail_closed_on_protected_loss=False,
),
with_compaction=False,
)
counter = iter(range(10_000))
monkeypatch.setattr(
SmartCrusher, "_canon", staticmethod(lambda item: f"unique-{next(counter)}")
)
protected_row = {"id": 1, "note": "AUDIT_FLAG"}
items_json = json.dumps([protected_row])
fake_result = {
"items": json.dumps([]),
"ccr_hash": "aaaa",
"dropped_summary": "<<ccr:aaaa 1_rows_offloaded>>",
"strategy_info": "smart_sample",
"compacted": None,
"compaction_kind": None,
}
out = crusher._apply_audit_safe_protection([protected_row], items_json, fake_result)
# Best-effort: the splice phase still ran and appended the
# protected row (splicing itself doesn't depend on `_canon` being
# idempotent across calls — only the *verification* mismatch is
# forced), so the row is present even though the strategy wasn't
# replaced with the fail-closed sentinel.
assert out["strategy_info"] != "audit_safe:fail_closed"
kept = json.loads(out["items"])
assert protected_row in kept
def test_invalid_protected_pattern_raises() -> None:
"""A regex that fails to compile is a caller bug and must raise
loudly at construction time — silently treating it as "nothing
protected" would defeat the point of audit-safe mode."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
with pytest.raises(ValueError, match="invalid protected_patterns regex"):
SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["("]),
with_compaction=False,
)
# ─── Production path: _smart_crush_content / apply() ───────────────────────
#
# `crush_array_json` is a convenience API used by tests and the CCR
# retrieval flow. The path `apply()` actually calls for every compressed
# tool/tool_result message is `_smart_crush_content`. Audit-safe mode has
# to hold on that path too, or it would only ever protect a code path
# real traffic never exercises.
def test_smart_crush_content_preserves_protected_rows_end_to_end() -> None:
"""Real lossy compression via `_smart_crush_content` (the method
`apply()` calls) still surfaces every protected row afterward."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
background = [{"id": i, "status": "ok"} for i in range(60)]
protected_rows = [
{"id": 25, "status": "ok", "note": "AUDIT_FLAG: rare compliance event A"},
{"id": 35, "status": "ok", "note": "AUDIT_FLAG: rare compliance event B"},
]
items = (
background[:25]
+ [protected_rows[0]]
+ background[25:35]
+ [protected_rows[1]]
+ background[35:]
)
content = json.dumps(items)
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
crushed, was_modified, info = crusher._smart_crush_content(content)
assert was_modified
kept = json.loads(crushed)
for row in protected_rows:
assert row in kept, f"protected row missing from _smart_crush_content output: {row!r}"
def test_apply_preserves_protected_rows_in_tool_message() -> None:
"""Full `Transform.apply()` path: a tool message with a large JSON
array containing protected rows gets compressed, and the resulting
message content (before the digest marker) still contains every
protected row. This is the exact code path the real proxy runs for
every tool output — proof audit-safe mode isn't test-only plumbing."""
from headroom import OpenAIProvider, Tokenizer
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
background = [{"id": i, "status": "ok"} for i in range(60)]
protected_rows = [
{"id": 25, "status": "ok", "note": "AUDIT_FLAG: rare compliance event A"},
{"id": 35, "status": "ok", "note": "AUDIT_FLAG: rare compliance event B"},
]
items = (
background[:25]
+ [protected_rows[0]]
+ background[25:35]
+ [protected_rows[1]]
+ background[35:]
)
content = json.dumps(items)
crusher = SmartCrusher(
SmartCrusherConfig(
audit_safe=True, protected_patterns=["AUDIT_FLAG"], min_tokens_to_crush=10
),
with_compaction=False,
)
messages = [
{"role": "user", "content": "check the results"},
{"role": "assistant", "tool_calls": [{"id": "t1", "function": {"name": "query"}}]},
{"role": "tool", "tool_call_id": "t1", "content": content},
]
provider = OpenAIProvider()
tokenizer = Tokenizer(provider.get_token_counter("gpt-4o"), "gpt-4o")
result = crusher.apply(messages, tokenizer)
tool_message = result.messages[-1]
assert tool_message["role"] == "tool"
# Content is `<crushed>\n<digest_marker>` — strip the marker line.
crushed_body = tool_message["content"].rsplit("\n", 1)[0]
kept = json.loads(crushed_body)
for row in protected_rows:
assert row in kept, f"protected row missing from apply() output: {row!r}"
assert len(kept) < len(items)
def test_content_protection_falls_back_to_pattern_count_for_non_array_output() -> None:
"""Direct unit test of `_apply_audit_safe_protection_to_content`'s
non-list branch: when `crushed` isn't a JSON array (e.g. a
lossless CSV/table render, or an opaque marker string), there's no
row structure to splice into, so verification counts protected
pattern matches in the raw text instead. A drop in match count
still fails closed."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
protected_row = {"id": 1, "note": "AUDIT_FLAG"}
original_content = json.dumps([protected_row, {"id": 2, "note": "fine"}])
# Simulate a lossless render that dropped the marker text entirely.
crushed_without_marker = "id,note\n2,fine"
out_text, was_modified, info = crusher._apply_audit_safe_protection_to_content(
[protected_row], original_content, crushed_without_marker, True, "lossless:table"
)
assert out_text == original_content
assert was_modified is False
assert info == "audit_safe:fail_closed"
def test_content_protection_no_op_when_pattern_count_preserved() -> None:
"""Non-list `crushed` output that still contains every protected
pattern occurrence is left untouched — the fallback only fires on
an actual count decrease."""
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
crusher = SmartCrusher(
SmartCrusherConfig(audit_safe=True, protected_patterns=["AUDIT_FLAG"]),
with_compaction=False,
)
protected_row = {"id": 1, "note": "AUDIT_FLAG"}
original_content = json.dumps([protected_row])
crushed_with_marker = "id,note\n1,AUDIT_FLAG"
out_text, was_modified, info = crusher._apply_audit_safe_protection_to_content(
[protected_row], original_content, crushed_with_marker, True, "lossless:table"
)
assert out_text == crushed_with_marker
assert was_modified is True
assert info == "lossless:table"