feat: probe-based retention scoring of recorded compression events (#862)

Closes #861

## What

First piece of compression quality measurement on **real proxied
sessions** (vs the existing public-benchmark evals): an opt-in recorder
captures (original, compressed) message pairs at each compression event,
and a deterministic offline prober scores what survived.

**Recorder** (`headroom/proxy/probe_recorder.py`)
- `CompressionEventRecorder` implements the existing `PipelineExtension`
protocol, subscribed to `INPUT_COMPRESSED`. Registered ONLY when
`HEADROOM_PROBE_RECORD_DIR` is set — off means not even constructed,
zero request-path overhead.
- One JSONL line per compression event that changed tokens: `{ts,
request_id, provider, model, tokens_before, tokens_after,
transforms_applied, original_messages, compressed_messages}`. One file
per PID (no interleaving), directory mode 0700.
- Fail-open everywhere: construction failure logs a warning and disables
recording; runtime exceptions are already swallowed by
`PipelineExtensionManager.emit`.
- Enabling handler change: the two `INPUT_COMPRESSED` emit sites
(anthropic + openai) add a read-only `original_messages` reference to
event metadata. No copies, no behavior change for other consumers.

**Probes** (`headroom/evals/session_probes.py` + `headroom evals probes`
CLI)
- Probe targets extracted from ORIGINAL tool-result content across three
dimensions: **exact numerics** (number + key context, incl. JSON-quoted
keys), **artifact trail** (paths, URLs, hex hashes, UUIDs), **error
evidence** (lines matching the existing `is_error_content` heuristic).
- Each target classified as **retained** (verbatim, or surviving a
legitimate format conversion — punctuation-normalized match; numerics
require key AND value to survive; error lines tolerate dropped JSON key
prefixes), **recoverable** (absent but a CCR retrieval marker is
present), or **lost** (gone with no retrieval path).
- Report: aggregate retention per dimension, bucketed by compression
ratio (the quality-per-ratio curve), and grouped per transform.
`--json-output` for machine-readable results.
- Fully offline: no LLM, no API key. The recording format is designed to
feed an LLM-judge pass later (out of scope per #861).

## Tests

33 new tests (red before, green after): `tests/test_probe_recorder.py`
(11 — event filtering, JSONL shape, env activation, fail-open on
unusable path, 0700 dir mode) and `tests/test_session_probes.py` (22 —
extraction per dimension incl. JSON-quoted numerics,
retained/recoverable/lost classification, format-change survival for
numerics and error lines, ratio bucketing, transform dedup,
malformed-line skipping, report rendering/serialization). Both proxy
lifecycle tests additionally assert the INPUT_COMPRESSED
`original_messages` metadata contract end-to-end through the real
anthropic/openai handlers, so a refactor cannot silently disable the
recorder.

Local runs: new tests + `tests/test_proxy_pipeline_lifecycle.py` +
`tests/test_canonical_pipeline.py` + `tests/test_pipeline.py` — 45
passed. `ruff check` + `uvx ruff format --check` clean.

### Self-review hardening (second commit)

- Hex artifact regex now requires at least one `a-f`, so bare decimal
runs (timestamps, counters) no longer inflate the artifact dimension.
- Inflation events (ratio > 1, the #847 territory) get an explicit
`1.00+ (inflated)` ratio bucket instead of silently dropping out of the
bucketed view.
- `run_probes` streams recording files line by line instead of slurping
them.
- Documented honestly: marker recoverability is event-scoped
(comparative metric, not absolute); recorder writes synchronously on the
request path (diagnostic sessions, not always-on).
- `headroom/evals/README.md` gained a Session Probes usage section.

## Real behavior proof

**Setup:** macOS (Darwin 25.5), Python 3.11.9, this branch, real proxy
(`python -m headroom.proxy.server --port 18994 --anthropic-api-url
http://127.0.0.1:18995`) with a local mock Anthropic upstream (no key),
`HEADROOM_PROBE_RECORD_DIR=/tmp/headroom-probe-proof/recordings`.

**Steps:** POSTed three Anthropic-format conversations whose tool
results carry large JSON arrays (220-row uniform logs, 3000-row logs,
800 heterogeneous events), each containing known numerics, paths, trace
hashes, and one error line.

**Observed** — recorder wrote `compression-events-<pid>.jsonl` (one line
per event); `headroom evals probes --recordings ...`:

```
Probed 3 compression events

Aggregate retention:
  numerics    97.7% retained,   2.3% recoverable,   0.0% lost (527 targets)
  artifacts  100.0% retained,   0.0% recoverable,   0.0% lost (6555 targets)
  errors     100.0% retained,   0.0% recoverable,   0.0% lost (3 targets)

By compression ratio (tokens_after / tokens_before):
  ratio 0.50-0.75:  numerics 100.0% retained  (CSV compaction — lossless, correctly recognized)
  ratio 0.75-1.00:  numerics  95.7% retained, 4.3% recoverable  (SmartCrusher sampling — dropped values carried a CCR marker)
```

All three classifications exercised: verbatim/format-change retention on
the CSV-compacted events, **recoverable** on the heterogeneous event
where SmartCrusher sampled rows out behind a `Retrieve more: hash=`
marker, and the injected error lines retained in every event (the
error-protection gate held). The `lost` path is covered by unit tests.
The first iteration of this proof exposed two real bugs — naive verbatim
matching misreported lossless JSON→CSV compaction as 100% lost, and
duplicated transform markers double-counted tallies — both fixed with
regression tests.

**Not tested live:** Gemini path (no `INPUT_COMPRESSED` emit parity —
pre-existing, same gap as #819); LLM-judge scoring (out of scope per
#861).

## Security

Recordings contain full conversation content in plaintext: opt-in env
var only, local disk only, dir mode 0700, documented in CLI help.

## Out of scope (per #861)

LLM-judge dimensions (decisions/intent, next steps), ACON-style
counterfactual replay, automatic rule revision.

---------

Co-authored-by: Ash Rhodes <ashley.rhodes@king.com>
This commit is contained in:
Focused Instability 2026-06-11 20:02:36 +02:00 committed by GitHub
parent c71592d421
commit c2106cbdab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 940 additions and 1 deletions

View file

@ -3,6 +3,7 @@
from __future__ import annotations
from collections.abc import Callable
from pathlib import Path
import click
@ -661,3 +662,36 @@ Running evaluation...
if output:
result.save(output)
click.echo(f"\nResults saved to: {output}")
@evals.command("probes")
@click.option(
"--recordings",
"recordings_dir",
required=True,
type=click.Path(exists=True, file_okay=False, path_type=Path),
help="Directory of JSONL recordings written via HEADROOM_PROBE_RECORD_DIR.",
)
@click.option(
"--json-output",
type=click.Path(dir_okay=False, path_type=Path),
help="Optional machine-readable JSON report output.",
)
def probes(recordings_dir: Path, json_output: Path | None) -> None:
"""Score retention of recorded compression events (offline, no LLM).
\b
Record sessions first by running the proxy with
HEADROOM_PROBE_RECORD_DIR set. Recordings contain full conversation
content in plaintext and stay on this machine.
"""
import json as json_module
from headroom.evals.session_probes import render_report, run_probes
report = run_probes(recordings_dir)
click.echo(render_report(report))
if json_output:
json_output.parent.mkdir(parents=True, exist_ok=True)
json_output.write_text(json_module.dumps(report.to_dict(), indent=2), encoding="utf-8")
click.echo(f"\nWrote JSON report: {json_output}")

View file

@ -0,0 +1,360 @@
"""Deterministic retention probes over recorded compression events.
Offline scoring of what compression removed from real proxied sessions no
LLM, no API key. For each event recorded by
``headroom.proxy.probe_recorder``, probe targets are extracted from the
ORIGINAL tool-result content and each is classified against the compressed
messages as:
- ``retained`` appears verbatim in the compressed content, or survives in
punctuation-normalized form (compressors legitimately
reshape JSON into tables/KV; for numerics the key and the
value must both survive the format change)
- ``recoverable`` absent, but the compressed content carries a CCR
retrieval marker, so the agent can fetch the original back
- ``lost`` absent with no retrieval path
Dimensions follow production session-replay findings: exact numerics are the
leakiest under compression, artifact trails (paths, hashes, URLs) the weakest,
and error evidence the most critical to keep.
Known limitation: marker recoverability is event-scoped, not block-scoped a
retrieval marker anywhere in the compressed messages marks every missing
target as recoverable, which can overcount when the marker belongs to a
different block than the loss. The metric is comparative (across ratios,
transforms, and versions), not absolute.
"""
from __future__ import annotations
import json
import re
from collections.abc import Iterable, Iterator
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from headroom.learn.scanner import is_error_content
DIMENSIONS = ("numerics", "artifacts", "errors")
# Mirrors the marker shapes matched by
# headroom.transforms.compression_units._CCR_MARKER_RE (kept local so the
# evals layer does not depend on a private transforms symbol).
_CCR_MARKER_RE = re.compile(r"Retrieve more: hash=|Retrieve original: hash=|<<ccr:[^>]+>>")
# A number with its immediate key context ("retry_limit: 3", "port=8787",
# JSON's '"latency_ms": 12'). Bare numbers are skipped: without context they
# are unverifiable noise.
_NUMERIC_RE = re.compile(r"[A-Za-z_][\w.-]{0,24}\"?[ =:]{1,3}\d+(?:\.\d+)?")
_URL_RE = re.compile(r"https?://[^\s\"'<>)\]]+")
_PATH_RE = re.compile(r"(?:~/|\.{1,2}/|/)?(?:[\w.-]+/){2,}[\w.@-]+")
# Requires at least one a-f so bare decimal runs (timestamps, row counts) are
# not mistaken for content hashes.
_HEX_RE = re.compile(r"\b(?=[0-9a-f]*[a-f])[0-9a-f]{7,64}\b")
_UUID_RE = re.compile(
r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b"
)
_MIN_TARGET_LEN = 4
_ERROR_LINE_PREFIX_LEN = 160
# Final bucket catches inflation events (tokens_after > tokens_before), which
# the recorder captures because compression changed the token count.
_RATIO_BUCKETS = ((0.0, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.01), (1.01, float("inf")))
# Collapse punctuation that format conversions (JSON -> table/KV/CSV) rewrite,
# keeping path/url/hash-significant characters.
_NORMALIZE_RE = re.compile(r"[^\w./-]+")
_NUMERIC_SPLIT_RE = re.compile(r"(.+?)[\"' =:]+(\d+(?:\.\d+)?)$")
@dataclass
class DimensionTally:
"""Counts for one probe dimension."""
total: int = 0
retained: int = 0
recoverable: int = 0
@property
def lost(self) -> int:
return self.total - self.retained - self.recoverable
def add(self, other: DimensionTally) -> None:
self.total += other.total
self.retained += other.retained
self.recoverable += other.recoverable
def to_dict(self) -> dict[str, int]:
return {
"total": self.total,
"retained": self.retained,
"recoverable": self.recoverable,
"lost": self.lost,
}
@dataclass
class EventProbeResult:
"""Probe outcome for a single recorded compression event."""
request_id: str
ratio: float
transforms: list[str]
dims: dict[str, DimensionTally]
def to_dict(self) -> dict[str, Any]:
return {
"request_id": self.request_id,
"ratio": round(self.ratio, 4),
"transforms": self.transforms,
"dimensions": {name: tally.to_dict() for name, tally in self.dims.items()},
}
@dataclass
class ProbeReport:
"""Aggregate probe outcomes across all recorded events."""
events: list[EventProbeResult] = field(default_factory=list)
skipped_lines: int = 0
def aggregate(self) -> dict[str, DimensionTally]:
totals = {name: DimensionTally() for name in DIMENSIONS}
for event in self.events:
for name, tally in event.dims.items():
totals[name].add(tally)
return totals
def by_ratio_bucket(self) -> dict[str, dict[str, DimensionTally]]:
buckets: dict[str, dict[str, DimensionTally]] = {}
for low, high in _RATIO_BUCKETS:
label = (
"1.00+ (inflated)" if high == float("inf") else f"{low:.2f}-{min(high, 1.0):.2f}"
)
buckets[label] = {name: DimensionTally() for name in DIMENSIONS}
for event in self.events:
if low <= event.ratio < high:
for name, tally in event.dims.items():
buckets[label][name].add(tally)
return buckets
def by_transform(self) -> dict[str, dict[str, DimensionTally]]:
transforms: dict[str, dict[str, DimensionTally]] = {}
for event in self.events:
for transform in set(event.transforms):
per_dim = transforms.setdefault(
transform, {name: DimensionTally() for name in DIMENSIONS}
)
for name, tally in event.dims.items():
per_dim[name].add(tally)
return transforms
def to_dict(self) -> dict[str, Any]:
return {
"events": [event.to_dict() for event in self.events],
"skipped_lines": self.skipped_lines,
"aggregate": {name: tally.to_dict() for name, tally in self.aggregate().items()},
"by_ratio_bucket": {
label: {name: tally.to_dict() for name, tally in dims.items()}
for label, dims in self.by_ratio_bucket().items()
},
"by_transform": {
transform: {name: tally.to_dict() for name, tally in dims.items()}
for transform, dims in self.by_transform().items()
},
}
def _to_text(content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for item in content:
if isinstance(item, dict) and item.get("type") == "text":
parts.append(str(item.get("text", "")))
elif isinstance(item, str):
parts.append(item)
return "\n".join(parts)
return "" if content is None else str(content)
def _tool_texts(messages: Iterable[Any] | None) -> list[str]:
"""Extract tool-result text from OpenAI (role=tool) and Anthropic blocks."""
out: list[str] = []
for msg in messages or []:
if not isinstance(msg, dict):
continue
content = msg.get("content")
if msg.get("role") == "tool":
out.append(_to_text(content))
elif isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("type") == "tool_result":
out.append(_to_text(block.get("content")))
return [text for text in out if text]
def _all_text(node: Any) -> Iterator[str]:
"""Yield every string leaf in a message structure (survival haystack)."""
if isinstance(node, str):
yield node
elif isinstance(node, dict):
for value in node.values():
yield from _all_text(value)
elif isinstance(node, list):
for item in node:
yield from _all_text(item)
def extract_probe_targets(text: str) -> dict[str, set[str]]:
"""Extract probe targets per dimension from original tool-result text."""
targets: dict[str, set[str]] = {name: set() for name in DIMENSIONS}
targets["numerics"].update(
match for match in _NUMERIC_RE.findall(text) if len(match) >= _MIN_TARGET_LEN
)
for pattern in (_URL_RE, _PATH_RE, _HEX_RE, _UUID_RE):
targets["artifacts"].update(
match for match in pattern.findall(text) if len(match) >= _MIN_TARGET_LEN
)
for line in text.splitlines():
stripped = line.strip()
if len(stripped) >= _MIN_TARGET_LEN and is_error_content(stripped):
targets["errors"].add(stripped[:_ERROR_LINE_PREFIX_LEN])
return targets
def _normalize(text: str) -> str:
return _NORMALIZE_RE.sub(" ", text).strip()
def _target_survives(dimension: str, value: str, haystack: str, normalized_haystack: str) -> bool:
if value in haystack:
return True
normalized_value = _normalize(value)
if normalized_value and normalized_value in normalized_haystack:
return True
if dimension == "errors":
# Format conversions drop JSON key prefixes ('"msg": "Error..."'
# becomes a bare CSV/KV cell); the error substance is what matters.
_, _, remainder = normalized_value.partition(" ")
if len(remainder) >= _MIN_TARGET_LEN and remainder in normalized_haystack:
return True
if dimension == "numerics":
# Format conversions (JSON -> table) separate key from value; count the
# probe as retained only when both still appear.
match = _NUMERIC_SPLIT_RE.match(value)
if match:
key, number = match.groups()
normalized_key = _normalize(key)
if (
normalized_key
and normalized_key in normalized_haystack
and re.search(rf"\b{re.escape(number)}\b", normalized_haystack)
):
return True
return False
def probe_event(record: dict[str, Any]) -> EventProbeResult | None:
"""Score one recorded compression event; None if it cannot be scored."""
tokens_before = record.get("tokens_before")
tokens_after = record.get("tokens_after")
if not isinstance(tokens_before, (int, float)) or not isinstance(tokens_after, (int, float)):
return None
if tokens_before <= 0:
return None
original_text = "\n".join(_tool_texts(record.get("original_messages")))
compressed_text = "\n".join(_all_text(record.get("compressed_messages")))
normalized_compressed = _normalize(compressed_text)
has_marker = bool(_CCR_MARKER_RE.search(compressed_text))
dims: dict[str, DimensionTally] = {}
for name, values in extract_probe_targets(original_text).items():
tally = DimensionTally(total=len(values))
for value in values:
if _target_survives(name, value, compressed_text, normalized_compressed):
tally.retained += 1
elif has_marker:
tally.recoverable += 1
dims[name] = tally
transforms = [str(item) for item in record.get("transforms_applied") or []]
return EventProbeResult(
request_id=str(record.get("request_id", "")),
ratio=float(tokens_after) / float(tokens_before),
transforms=transforms,
dims=dims,
)
def run_probes(recordings_dir: Path) -> ProbeReport:
"""Probe every event in every ``*.jsonl`` recording under a directory."""
report = ProbeReport()
for path in sorted(recordings_dir.glob("*.jsonl")):
with path.open(encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if not line:
continue
try:
record = json.loads(line)
except json.JSONDecodeError:
report.skipped_lines += 1
continue
result = probe_event(record) if isinstance(record, dict) else None
if result is None:
report.skipped_lines += 1
continue
report.events.append(result)
return report
def _format_tally(tally: DimensionTally) -> str:
if tally.total == 0:
return "n/a (0 targets)"
retained_pct = 100.0 * tally.retained / tally.total
recoverable_pct = 100.0 * tally.recoverable / tally.total
lost_pct = 100.0 * tally.lost / tally.total
return (
f"{retained_pct:5.1f}% retained, {recoverable_pct:5.1f}% recoverable, "
f"{lost_pct:5.1f}% lost ({tally.total} targets)"
)
def render_report(report: ProbeReport) -> str:
"""Render a human-readable retention report."""
lines = [
f"Probed {len(report.events)} compression events"
+ (f" ({report.skipped_lines} lines skipped)" if report.skipped_lines else ""),
"",
"Aggregate retention:",
]
for name, tally in report.aggregate().items():
lines.append(f" {name:<10} {_format_tally(tally)}")
lines += ["", "By compression ratio (tokens_after / tokens_before):"]
for label, dims in report.by_ratio_bucket().items():
if all(tally.total == 0 for tally in dims.values()):
continue
lines.append(f" ratio {label}:")
for name, tally in dims.items():
lines.append(f" {name:<10} {_format_tally(tally)}")
by_transform = report.by_transform()
if by_transform:
lines += ["", "By transform:"]
for transform in sorted(by_transform):
lines.append(f" {transform}:")
for name, tally in by_transform[transform].items():
lines.append(f" {name:<10} {_format_tally(tally)}")
return "\n".join(lines)

View file

@ -1210,6 +1210,9 @@ class AnthropicHandlerMixin:
"tokens_before": original_tokens,
"tokens_after": optimized_tokens,
"transforms_applied": transforms_applied,
# Read-only reference for recording extensions (probe
# recorder); extensions must not mutate it.
"original_messages": original_messages,
},
)
if compressed_event.messages is not None:

View file

@ -1792,6 +1792,9 @@ class OpenAIHandlerMixin:
"tokens_before": original_tokens,
"tokens_after": optimized_tokens,
"transforms_applied": transforms_applied,
# Read-only reference for recording extensions (probe
# recorder); extensions must not mutate it.
"original_messages": original_messages,
},
)
if compressed_event.messages is not None:

View file

@ -0,0 +1,94 @@
"""Opt-in JSONL recorder for compression events (probe-based replay evals).
Records (original, compressed) message pairs at ``INPUT_COMPRESSED`` so that
``headroom evals probes`` can measure what compression removed from real
proxied sessions. Activated only when ``HEADROOM_PROBE_RECORD_DIR`` is set;
recordings contain full conversation content in plaintext, are written with
directory mode 0700, and never leave the machine.
Writes happen synchronously on the request path, so this is a diagnostic
tool for bounded recording sessions, not an always-on production setting.
"""
from __future__ import annotations
import json
import logging
import os
import threading
import time
from pathlib import Path
from headroom.pipeline import PipelineEvent, PipelineStage
logger = logging.getLogger(__name__)
RECORD_DIR_ENV = "HEADROOM_PROBE_RECORD_DIR"
class CompressionEventRecorder:
"""Pipeline extension appending one JSONL line per compression event.
Only ``INPUT_COMPRESSED`` events that carry ``original_messages`` in their
metadata and actually changed the token count are recorded. The extension
never mutates the event; ``PipelineExtensionManager.emit`` already swallows
extension exceptions, so a broken recorder cannot break a request.
"""
def __init__(self, record_dir: Path) -> None:
self._dir = record_dir
self._dir.mkdir(parents=True, exist_ok=True)
os.chmod(self._dir, 0o700)
# One file per process so concurrent proxy workers never interleave
# partial lines.
self._path = self._dir / f"compression-events-{os.getpid()}.jsonl"
self._lock = threading.Lock()
@property
def path(self) -> Path:
return self._path
def on_pipeline_event(self, event: PipelineEvent) -> None:
if event.stage is not PipelineStage.INPUT_COMPRESSED:
return None
metadata = event.metadata or {}
original = metadata.get("original_messages")
tokens_before = metadata.get("tokens_before")
tokens_after = metadata.get("tokens_after")
if original is None or event.messages is None:
return None
if tokens_before is None or tokens_after is None or tokens_before == tokens_after:
return None
record = {
"ts": time.time(),
"request_id": event.request_id,
"provider": event.provider,
"model": event.model,
"tokens_before": tokens_before,
"tokens_after": tokens_after,
"transforms_applied": metadata.get("transforms_applied") or [],
"original_messages": original,
"compressed_messages": event.messages,
}
line = json.dumps(record, ensure_ascii=False, default=str)
with self._lock:
with self._path.open("a", encoding="utf-8") as fh:
fh.write(line + "\n")
return None
def probe_recorder_from_env() -> CompressionEventRecorder | None:
"""Build a recorder when ``HEADROOM_PROBE_RECORD_DIR`` is set, else None.
Fail-open: any error constructing the recorder (unwritable path, etc.)
disables recording with a warning instead of breaking proxy startup.
"""
record_dir = os.environ.get(RECORD_DIR_ENV, "").strip()
if not record_dir:
return None
try:
return CompressionEventRecorder(Path(record_dir).expanduser())
except Exception as exc: # noqa: BLE001 - recorder must never break proxy startup
logger.warning("probe recorder disabled (%s): %s", RECORD_DIR_ENV, exc)
return None

View file

@ -139,6 +139,7 @@ from headroom.proxy.modes import (
is_token_mode,
normalize_proxy_mode,
)
from headroom.proxy.probe_recorder import probe_recorder_from_env
from headroom.proxy.project_context import (
classify_project,
set_current_project,
@ -314,9 +315,13 @@ class HeadroomProxy(
def __init__(self, config: ProxyConfig):
self.config = config
self.config.mode = normalize_proxy_mode(self.config.mode)
pipeline_extensions = list(config.pipeline_extensions or [])
probe_recorder = probe_recorder_from_env()
if probe_recorder is not None:
pipeline_extensions.append(probe_recorder)
self.pipeline_extensions = PipelineExtensionManager(
hooks=config.hooks,
extensions=config.pipeline_extensions,
extensions=pipeline_extensions,
discover=config.discover_pipeline_extensions,
)

View file

@ -0,0 +1,139 @@
"""Tests for the opt-in compression event probe recorder."""
import json
import stat
from headroom.pipeline import PipelineEvent, PipelineStage
from headroom.proxy.probe_recorder import (
RECORD_DIR_ENV,
CompressionEventRecorder,
probe_recorder_from_env,
)
ORIGINAL = [
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": "retry_limit: 3"}],
}
]
COMPRESSED = [
{
"role": "user",
"content": [{"type": "tool_result", "tool_use_id": "t1", "content": "[compressed]"}],
}
]
def _metadata(**overrides):
metadata = {
"tokens_before": 100,
"tokens_after": 40,
"transforms_applied": ["smart_crusher"],
"original_messages": ORIGINAL,
}
metadata.update(overrides)
return metadata
def _event(stage=PipelineStage.INPUT_COMPRESSED, messages=COMPRESSED, metadata=None):
return PipelineEvent(
stage=stage,
operation="proxy.request",
request_id="req-1",
provider="anthropic",
model="claude-test",
messages=messages,
metadata=_metadata() if metadata is None else metadata,
)
class TestCompressionEventRecorder:
def test_records_compression_event(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event())
lines = recorder.path.read_text(encoding="utf-8").splitlines()
assert len(lines) == 1
record = json.loads(lines[0])
assert record["request_id"] == "req-1"
assert record["provider"] == "anthropic"
assert record["model"] == "claude-test"
assert record["tokens_before"] == 100
assert record["tokens_after"] == 40
assert record["transforms_applied"] == ["smart_crusher"]
assert record["original_messages"] == ORIGINAL
assert record["compressed_messages"] == COMPRESSED
assert record["ts"] > 0
def test_appends_one_line_per_event(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event())
recorder.on_pipeline_event(_event())
assert len(recorder.path.read_text(encoding="utf-8").splitlines()) == 2
def test_ignores_other_stages(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event(stage=PipelineStage.INPUT_ROUTED))
assert not recorder.path.exists()
def test_skips_without_original_messages(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event(metadata=_metadata(original_messages=None)))
assert not recorder.path.exists()
def test_skips_zero_token_delta(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event(metadata=_metadata(tokens_after=100)))
assert not recorder.path.exists()
def test_skips_without_compressed_messages(self, tmp_path):
recorder = CompressionEventRecorder(tmp_path)
recorder.on_pipeline_event(_event(messages=None))
assert not recorder.path.exists()
def test_record_dir_is_private(self, tmp_path):
record_dir = tmp_path / "recordings"
CompressionEventRecorder(record_dir)
mode = stat.S_IMODE(record_dir.stat().st_mode)
assert mode == 0o700
class TestProbeRecorderFromEnv:
def test_disabled_without_env(self, monkeypatch):
monkeypatch.delenv(RECORD_DIR_ENV, raising=False)
assert probe_recorder_from_env() is None
def test_disabled_with_blank_env(self, monkeypatch):
monkeypatch.setenv(RECORD_DIR_ENV, " ")
assert probe_recorder_from_env() is None
def test_enabled_with_env(self, tmp_path, monkeypatch):
record_dir = tmp_path / "recordings"
monkeypatch.setenv(RECORD_DIR_ENV, str(record_dir))
recorder = probe_recorder_from_env()
assert isinstance(recorder, CompressionEventRecorder)
assert record_dir.is_dir()
def test_fail_open_on_unusable_path(self, tmp_path, monkeypatch):
blocker = tmp_path / "not-a-dir"
blocker.write_text("file", encoding="utf-8")
monkeypatch.setenv(RECORD_DIR_ENV, str(blocker / "recordings"))
assert probe_recorder_from_env() is None

View file

@ -14,9 +14,11 @@ from headroom.proxy.server import ProxyConfig, create_app
class _RecordingExtension:
def __init__(self) -> None:
self.stages: list[PipelineStage] = []
self.events: list = []
def on_pipeline_event(self, event):
self.stages.append(event.stage)
self.events.append(event)
return None
@ -25,6 +27,23 @@ class _DummyTokenizer:
return len(messages)
def _assert_compressed_event_carries_originals(events: list) -> None:
"""INPUT_COMPRESSED must expose the pre-compression messages to extensions.
The probe recorder (headroom.proxy.probe_recorder) depends on this
metadata contract; dropping it silently disables session recording.
"""
compressed = [event for event in events if event.stage is PipelineStage.INPUT_COMPRESSED]
assert compressed
original = compressed[0].metadata.get("original_messages")
assert isinstance(original, list)
assert any(
message.get("role") == "user" and "hello" in str(message.get("content"))
for message in original
if isinstance(message, dict)
)
def _assert_stage_order(stages: list[PipelineStage]) -> None:
expected = [
PipelineStage.SETUP,
@ -136,6 +155,7 @@ def test_openai_chat_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None:
assert response.status_code == 200
_assert_stage_order(recorder.stages)
_assert_compressed_event_carries_originals(recorder.events)
def test_anthropic_messages_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None:
@ -214,3 +234,4 @@ def test_anthropic_messages_pipeline_events_cover_proxy_lifecycle(monkeypatch) -
assert response.status_code == 200
_assert_stage_order(recorder.stages)
_assert_compressed_event_carries_originals(recorder.events)

View file

@ -0,0 +1,280 @@
"""Tests for deterministic retention probes over recorded compression events."""
import json
from headroom.evals.session_probes import (
DIMENSIONS,
DimensionTally,
extract_probe_targets,
probe_event,
render_report,
run_probes,
)
ORIGINAL_TOOL_TEXT = (
"Deploy summary\n"
"retry_limit: 3\n"
"port=8787\n"
"see headroom/proxy/server.py and https://example.com/build/42\n"
"commit d293b77ab12\n"
"ModuleNotFoundError: No module named 'left_pad'\n"
)
def _record(compressed_content, tokens_before=100, tokens_after=40, transforms=None):
return {
"request_id": "req-1",
"tokens_before": tokens_before,
"tokens_after": tokens_after,
"transforms_applied": transforms or ["smart_crusher"],
"original_messages": [
{
"role": "user",
"content": [{"type": "tool_result", "content": ORIGINAL_TOOL_TEXT}],
}
],
"compressed_messages": [
{
"role": "user",
"content": [{"type": "tool_result", "content": compressed_content}],
}
],
}
class TestExtractProbeTargets:
def test_extracts_contextual_numerics(self):
targets = extract_probe_targets("retry_limit: 3 and port=8787")
assert "retry_limit: 3" in targets["numerics"]
assert "port=8787" in targets["numerics"]
def test_extracts_json_quoted_numerics(self):
targets = extract_probe_targets('{"latency_ms": 12, "status": 200}')
assert 'latency_ms": 12' in targets["numerics"]
assert 'status": 200' in targets["numerics"]
def test_extracts_artifacts(self):
text = (
"path headroom/proxy/server.py url https://example.com/build/42 "
"hash d293b77ab12 id 123e4567-e89b-42d3-a456-426614174000"
)
targets = extract_probe_targets(text)
assert "headroom/proxy/server.py" in targets["artifacts"]
assert "https://example.com/build/42" in targets["artifacts"]
assert "d293b77ab12" in targets["artifacts"]
assert "123e4567-e89b-42d3-a456-426614174000" in targets["artifacts"]
def test_bare_decimal_runs_are_not_artifacts(self):
targets = extract_probe_targets("run id 27344471690 at ts 1765449600")
assert "27344471690" not in targets["artifacts"]
assert "1765449600" not in targets["artifacts"]
def test_extracts_error_lines(self):
targets = extract_probe_targets("all good\nModuleNotFoundError: No module named 'x'\n")
assert any("ModuleNotFoundError" in value for value in targets["errors"])
assert all("all good" not in value for value in targets["errors"])
def test_openai_role_tool_messages_supported(self):
record = _record("anything")
record["original_messages"] = [{"role": "tool", "content": ORIGINAL_TOOL_TEXT}]
result = probe_event(record)
assert result is not None
assert result.dims["numerics"].total > 0
class TestProbeEvent:
def test_everything_retained_when_content_survives(self):
result = probe_event(_record(ORIGINAL_TOOL_TEXT))
assert result is not None
for name in DIMENSIONS:
tally = result.dims[name]
assert tally.total > 0
assert tally.retained == tally.total
assert tally.lost == 0
def test_recoverable_when_ccr_marker_present(self):
result = probe_event(_record("[60 items compressed to 5. Retrieve more: hash=abc123def]"))
assert result is not None
numerics = result.dims["numerics"]
assert numerics.total > 0
assert numerics.retained == 0
assert numerics.recoverable == numerics.total
assert numerics.lost == 0
def test_lost_when_dropped_without_marker(self):
result = probe_event(_record("everything went fine"))
assert result is not None
for name in DIMENSIONS:
tally = result.dims[name]
assert tally.retained == 0
assert tally.recoverable == 0
assert tally.lost == tally.total
def test_ratio_and_transforms(self):
result = probe_event(_record("x", tokens_before=200, tokens_after=50))
assert result is not None
assert result.ratio == 0.25
assert result.transforms == ["smart_crusher"]
def test_numerics_retained_across_format_change(self):
record = _record("| latency_ms | status |\n| 12 | 200 |")
record["original_messages"] = [
{"role": "tool", "content": '{"latency_ms": 12, "status": 200}'}
]
result = probe_event(record)
assert result is not None
numerics = result.dims["numerics"]
assert numerics.total > 0
assert numerics.retained == numerics.total
def test_numerics_lost_when_value_dropped_after_format_change(self):
record = _record("| latency_ms |\n| 99 |")
record["original_messages"] = [{"role": "tool", "content": '{"latency_ms": 12}'}]
result = probe_event(record)
assert result is not None
assert result.dims["numerics"].lost == result.dims["numerics"].total
def test_error_line_survives_punctuation_rewrite(self):
record = _record("msg=ModuleNotFoundError: No module named 'left_pad'")
record["original_messages"] = [
{
"role": "tool",
"content": '{"msg": "ModuleNotFoundError: No module named \'left_pad\'"}',
}
]
result = probe_event(record)
assert result is not None
errors = result.dims["errors"]
assert errors.total > 0
assert errors.retained == errors.total
def test_error_line_survives_json_to_csv_compaction(self):
record = _record("error,ModuleNotFoundError: No module named 'left_pad',src/imports.py")
record["original_messages"] = [
{
"role": "tool",
"content": '{"msg": "ModuleNotFoundError: No module named \'left_pad\'"}',
}
]
result = probe_event(record)
assert result is not None
errors = result.dims["errors"]
assert errors.total > 0
assert errors.retained == errors.total
def test_rejects_unscorable_records(self):
assert probe_event({"tokens_before": 0, "tokens_after": 0}) is None
assert probe_event({"tokens_before": "x", "tokens_after": 5}) is None
assert probe_event({}) is None
class TestRunProbesAndReport:
def test_run_probes_reads_jsonl_and_skips_garbage(self, tmp_path):
records = [
_record(ORIGINAL_TOOL_TEXT),
_record("gone", tokens_before=100, tokens_after=80),
]
lines = [json.dumps(record) for record in records]
lines.insert(1, "{not valid json")
lines.append(json.dumps(["not", "a", "dict"]))
(tmp_path / "compression-events-1.jsonl").write_text(
"\n".join(lines) + "\n", encoding="utf-8"
)
report = run_probes(tmp_path)
assert len(report.events) == 2
assert report.skipped_lines == 2
def test_aggregate_sums_dimensions(self, tmp_path):
path = tmp_path / "compression-events-1.jsonl"
path.write_text(
json.dumps(_record(ORIGINAL_TOOL_TEXT)) + "\n" + json.dumps(_record("gone")) + "\n",
encoding="utf-8",
)
report = run_probes(tmp_path)
aggregate = report.aggregate()
for name in DIMENSIONS:
single = report.events[0].dims[name]
assert aggregate[name].total == single.total * 2
assert aggregate[name].retained == single.total
assert aggregate[name].lost == single.total
def test_bucketing_and_transform_grouping(self, tmp_path):
path = tmp_path / "compression-events-1.jsonl"
path.write_text(
json.dumps(_record("gone", tokens_before=100, tokens_after=10)) + "\n",
encoding="utf-8",
)
report = run_probes(tmp_path)
buckets = report.by_ratio_bucket()
assert buckets["0.00-0.25"]["numerics"].total > 0
assert buckets["0.75-1.00"]["numerics"].total == 0
assert "smart_crusher" in report.by_transform()
def test_inflated_events_land_in_inflation_bucket(self, tmp_path):
record = _record("gone", tokens_before=100, tokens_after=130)
path = tmp_path / "compression-events-1.jsonl"
path.write_text(json.dumps(record) + "\n", encoding="utf-8")
report = run_probes(tmp_path)
buckets = report.by_ratio_bucket()
assert buckets["1.00+ (inflated)"]["numerics"].total > 0
assert all(
dims["numerics"].total == 0
for label, dims in buckets.items()
if label != "1.00+ (inflated)"
)
def test_transform_grouping_dedupes_repeated_markers(self, tmp_path):
record = _record("gone", transforms=["smart_crusher", "smart_crusher"])
path = tmp_path / "compression-events-1.jsonl"
path.write_text(json.dumps(record) + "\n", encoding="utf-8")
report = run_probes(tmp_path)
per_dim = report.by_transform()["smart_crusher"]
assert per_dim["numerics"].total == report.events[0].dims["numerics"].total
def test_to_dict_and_render(self, tmp_path):
path = tmp_path / "compression-events-1.jsonl"
path.write_text(json.dumps(_record("gone")) + "\n", encoding="utf-8")
report = run_probes(tmp_path)
payload = report.to_dict()
rendered = render_report(report)
assert payload["aggregate"]["numerics"]["lost"] > 0
assert payload["events"][0]["ratio"] == 0.4
assert "Aggregate retention" in rendered
assert "smart_crusher" in rendered
def test_dimension_tally_lost_property(self):
tally = DimensionTally(total=5, retained=2, recoverable=1)
assert tally.lost == 2
assert tally.to_dict() == {"total": 5, "retained": 2, "recoverable": 1, "lost": 2}