diff --git a/CHANGELOG.md b/CHANGELOG.md index 64f8a8757..dbf607f2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 `<>` 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. diff --git a/headroom/transforms/smart_crusher.py b/headroom/transforms/smart_crusher.py index 249e37d34..796b8ce61 100644 --- a/headroom/transforms/smart_crusher.py +++ b/headroom/transforms/smart_crusher.py @@ -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 `<>` 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 `<>` 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. diff --git a/tests/test_transforms/test_smart_crusher_audit_safe.py b/tests/test_transforms/test_smart_crusher_audit_safe.py new file mode 100644 index 000000000..bd53a13ff --- /dev/null +++ b/tests/test_transforms/test_smart_crusher_audit_safe.py @@ -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 `<>` 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": "<>", + "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": "<>", + "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": "<>", + "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 `\n` — 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"