diff --git a/headroom/cli/evals.py b/headroom/cli/evals.py index 97001e414..573e2dbfb 100644 --- a/headroom/cli/evals.py +++ b/headroom/cli/evals.py @@ -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}") diff --git a/headroom/evals/session_probes.py b/headroom/evals/session_probes.py new file mode 100644 index 000000000..5574841ab --- /dev/null +++ b/headroom/evals/session_probes.py @@ -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=|<]+>>") + +# 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) diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index 15b9774ad..e03ef5744 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -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: diff --git a/headroom/proxy/handlers/openai.py b/headroom/proxy/handlers/openai.py index 7e3afc455..06103ea52 100644 --- a/headroom/proxy/handlers/openai.py +++ b/headroom/proxy/handlers/openai.py @@ -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: diff --git a/headroom/proxy/probe_recorder.py b/headroom/proxy/probe_recorder.py new file mode 100644 index 000000000..eaad8b0f4 --- /dev/null +++ b/headroom/proxy/probe_recorder.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index badc6a530..24c43ab35 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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, ) diff --git a/tests/test_probe_recorder.py b/tests/test_probe_recorder.py new file mode 100644 index 000000000..818602569 --- /dev/null +++ b/tests/test_probe_recorder.py @@ -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 diff --git a/tests/test_proxy_pipeline_lifecycle.py b/tests/test_proxy_pipeline_lifecycle.py index e61559265..e68be2198 100644 --- a/tests/test_proxy_pipeline_lifecycle.py +++ b/tests/test_proxy_pipeline_lifecycle.py @@ -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) diff --git a/tests/test_session_probes.py b/tests/test_session_probes.py new file mode 100644 index 000000000..c18e3b882 --- /dev/null +++ b/tests/test_session_probes.py @@ -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}