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

chore(proxy): per-strategy compression observability

Browse source 
Adds a `CompressionObserver` Protocol (`headroom.transforms.observability`)
and wires `ContentRouter` and `SmartCrusher` to call it once per real
compression event. `PrometheusMetrics` implements the protocol and
accumulates per-strategy counters (`compressions_by_strategy`,
`tokens_saved_by_strategy`).

Why: the TOIN→SmartCrusher silent disconnect was invisible for three
weeks because no signal distinguished by strategy. With per-strategy
counters in place, the next regression of that shape fails the test
suite the day it lands instead of waiting on a manual audit.

Counters live as in-process state only — deliberately NOT exported via
the Prometheus scrape or OTel surface. The metric→Supabase pipeline
treats each metric name as a column, and we cannot add new columns.
CI-level observability via `tests/test_compression_observability.py` is
sufficient to catch silent regressions; production export waits on a
non-column-adding pipeline.
This commit is contained in:
chopratejas 2026-04-28 13:11:48 -07:00
parent 9d5f15bdeb
commit cf97995834
6 changed files with 537 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

50
headroom/proxy/prometheus_metrics.py
View file

@ -81,6 +81,17 @@ class PrometheusMetrics:
self.tokens_output_total = 0
self.tokens_saved_total = 0
# Per-strategy compression counters. Populated lazily as we see
# each strategy tag — no hardcoded list of strategies; the keys
# come from ContentRouter's `CompressionStrategy.value` and
# SmartCrusher's literal `"smart_crusher"`. The forcing
# function for catching strategy-level silent regressions:
# if SmartCrusher events drop to zero in production, the
# `headroom_compressions_total{strategy="smart_crusher"}`
# counter shows it on day 1, not week 3.
self.compressions_by_strategy: dict[str, int] = defaultdict(int)
self.tokens_saved_by_strategy: dict[str, int] = defaultdict(int)
self.latency_sum_ms = 0.0
self.latency_min_ms = float("inf")
self.latency_max_ms = 0.0
@ -242,6 +253,36 @@ class PrometheusMetrics:
return
self.requests_by_stack[slug] += 1
def record_compression(
self,
strategy: str,
original_tokens: int,
compressed_tokens: int,
) -> None:
"""Implements `headroom.transforms.observability.CompressionObserver`.
Called once per real compression event by the configured
transforms (ContentRouter at routing-decision granularity;
SmartCrusher at message granularity in the legacy direct-
pipeline path). Increments the per-strategy counters that
get exported as labelled Prometheus metrics, so silent
regressions in any single strategy become visible in the
scrape.
Synchronous + lock-free: `defaultdict(int)` writes are
atomic under the GIL for these key types; the proxy serves
many requests concurrently and the contention here would be
a single dict write per routing decision.
Tokens saved is `max(0, original - compressed)` — the
observer never records "negative savings" even if a
compressor goofs and emits more tokens than it received.
"""
self.compressions_by_strategy[strategy] += 1
saved = original_tokens - compressed_tokens
if saved > 0:
self.tokens_saved_by_strategy[strategy] += saved
async def record_request(
self,
provider: str,
@ -523,6 +564,15 @@ class PrometheusMetrics:
help_text="Tokens saved by optimization",
value=self.tokens_saved_total,
)
# NOTE: per-strategy compression breakdown is tracked
# internally on `self.compressions_by_strategy` and
# `self.tokens_saved_by_strategy` (populated by
# `record_compression`) but **deliberately not exported
# here**. The proxy's metric→Supabase pipeline treats
# each metric name as a column, and we cannot add new
# columns. The state is still observable for tests +
# programmatic introspection; if/when a non-column-
# adding export path exists, surface it there.
_append_metric(
lines,
name="headroom_latency_ms_sum",

32
headroom/proxy/server.py
View file

@ -236,6 +236,22 @@ class HeadroomProxy(
self.anthropic_provider = self.provider_runtime.pipeline_provider("anthropic")
self.openai_provider = self.provider_runtime.pipeline_provider("openai")
# `metrics` is hoisted ahead of transform construction so the
# transforms can receive `self.metrics` as their compression
# observer at __init__ time. The forcing function for catching
# silent strategy regressions: per-strategy counters increment
# only when wired up here, so the wiring is mandatory, not
# something we patch in later. (See `RUST_DEV.md` audit notes.)
self.cost_tracker = (
CostTracker(
budget_limit_usd=config.budget_limit_usd,
budget_period=config.budget_period,
)
if config.cost_tracking_enabled
else None
)
self.metrics = PrometheusMetrics(cost_tracker=self.cost_tracker)
# Initialize transforms based on routing mode
# Choose context manager: IntelligentContextManager (smart) or RollingWindow (legacy)
context_manager: Transform # Can be either IntelligentContextManager or RollingWindow
@ -277,7 +293,7 @@ class HeadroomProxy(
router_config.protect_recent_reads_fraction = 0.3
transforms = [
CacheAligner(CacheAlignerConfig(enabled=False)),
ContentRouter(router_config),
ContentRouter(router_config, observer=self.metrics),
context_manager,
]
self._code_aware_status = "lazy" if config.code_aware_enabled else "disabled"
@ -295,6 +311,7 @@ class HeadroomProxy(
enabled=config.ccr_inject_tool,
inject_retrieval_marker=config.ccr_inject_tool, # Add CCR markers
),
observer=self.metrics,
),
context_manager,
]
@ -329,16 +346,9 @@ class HeadroomProxy(
else None
)
self.cost_tracker = (
CostTracker(
budget_limit_usd=config.budget_limit_usd,
budget_period=config.budget_period,
)
if config.cost_tracking_enabled
else None
)
self.metrics = PrometheusMetrics(cost_tracker=self.cost_tracker)
# `cost_tracker` and `metrics` were hoisted to before transforms so
# ContentRouter / SmartCrusher could take `self.metrics` as their
# compression observer at __init__ time.
# Prefix cache tracking: freeze already-cached messages to avoid
# invalidating the provider's prefix cache with our transforms

57
headroom/transforms/content_router.py
View file

@ -643,13 +643,26 @@ class ContentRouter(Transform):
name: str = "content_router"
def __init__(self, config: ContentRouterConfig | None = None):
def __init__(
self,
config: ContentRouterConfig | None = None,
observer: Any = None,
):
"""Initialize content router.
Args:
config: Router configuration. Uses defaults if None.
observer: Optional `CompressionObserver` (see
`headroom.transforms.observability`) called once per
routing decision after `compress()` finishes. The
proxy's `PrometheusMetrics` is the production
implementation — it increments per-strategy counters
so silent regressions become visible. `None` disables
observation; pick one explicitly per the no-fallback
rule in the audit doc.
"""
self.config = config or ContentRouterConfig()
self._observer = observer
# Lazy-loaded compressors
self._code_compressor: Any = None
@ -766,20 +779,46 @@ class ContentRouter(Transform):
RouterCompressionResult with compressed content and routing metadata.
"""
if not content or not content.strip():
return RouterCompressionResult(
result = RouterCompressionResult(
compressed=content,
original=content,
strategy_used=CompressionStrategy.PASSTHROUGH,
routing_log=[],
)
# Determine strategy from content analysis
strategy = self._determine_strategy(content)
if strategy == CompressionStrategy.MIXED:
return self._compress_mixed(content, context, question, bias=bias)
else:
return self._compress_pure(content, strategy, context, question, bias=bias)
# Determine strategy from content analysis
strategy = self._determine_strategy(content)
if strategy == CompressionStrategy.MIXED:
result = self._compress_mixed(content, context, question, bias=bias)
else:
result = self._compress_pure(content, strategy, context, question, bias=bias)
# One observer call per routing decision; the observer is the
# forcing function for catching strategy-level regressions.
# Empty routing_log (passthrough fast path) → no calls.
self._observe(result)
return result
def _observe(self, result: RouterCompressionResult) -> None:
"""Forward each `RoutingDecision` in `result.routing_log` to the
configured `CompressionObserver`. No-op when no observer is set.
Observers MUST NOT raise per the protocol contract; if one does
anyway, swallow at debug level. Compression already succeeded;
a buggy observer must not turn a 200 into a 500.
"""
if self._observer is None:
return
for d in result.routing_log:
try:
self._observer.record_compression(
strategy=d.strategy.value,
original_tokens=d.original_tokens,
compressed_tokens=d.compressed_tokens,
)
except Exception as e: # pragma: no cover - defensive
logger.debug("CompressionObserver raised (non-fatal): %s", e)
def _determine_strategy(self, content: str) -> CompressionStrategy:
"""Determine the compression strategy from content analysis.

77
headroom/transforms/observability.py Normal file
View file

@ -0,0 +1,77 @@
"""Observability protocol for compression events.
A single `CompressionObserver` interface that any transform can call
after a real compression event. Concrete observers — Prometheus, OTel,
structured logs — implement this; transforms only see the protocol.
The motivating regression: `ContentRouter._record_to_toin` skipped
SmartCrusher on the assumption SmartCrusher recorded its own TOIN
events (it did when SmartCrusher was Python; it stopped when the Rust
port took over). The disconnect was invisible for three weeks because
no metric distinguished compression events by strategy. This module
exists so the next regression of that shape alerts on day 1: if
SmartCrusher events drop to zero in production, the Prometheus
counter shows it immediately.
Design choices, called out for posterity:
- **No fallback observer.** Callers pass `None` or pass a real
observer. There is no "default no-op" instance — that would let a
caller silently disable observability by forgetting to pass one,
and we just spent a PR fixing exactly that class of bug. Be
explicit.
- **No observer registry.** A single observer per transform instance.
If you need multi-fanout, compose at the call site (one wrapper
observer that forwards to N children) — but the trivial pattern
doesn't need a registry baked in.
- **No batching.** Each compression event is one call. Volume is
bounded by the number of routing decisions per request — small.
Batching would only matter if observers had to round-trip to a
remote system; production observers (Prometheus) are in-process
counter increments, which are cheaper than the protocol dispatch.
- **Strategy as a string.** The router and crusher both already
serialize their strategy as the enum's `.value` tag. Passing the
string keeps observers from importing `CompressionStrategy` and
lets non-router callers (e.g. SmartCrusher in legacy mode) emit
the same shape without round-tripping through the enum.
"""
from __future__ import annotations
from typing import Protocol, runtime_checkable
@runtime_checkable
class CompressionObserver(Protocol):
"""Receive one notification per real compression event.
Implementations should be cheap — this lives on the proxy hot path,
one call per routing decision per request. A Prometheus-counter
increment is the right order of magnitude.
Args:
strategy: Lowercase tag identifying the compression strategy
that ran. Matches `CompressionStrategy.<NAME>.value` for
ContentRouter; SmartCrusher's legacy direct-call path
passes the literal `"smart_crusher"`.
original_tokens: Token count of the input the strategy
received.
compressed_tokens: Token count of the output the strategy
produced. Equal to `original_tokens` for passthrough;
less when compression saved tokens.
Implementations MUST NOT raise. If the observer needs to fail-
over (Prometheus client misconfigured, OTel exporter offline)
handle that internally — bubbling exceptions out of an observer
would break the compression that just succeeded, which is the
opposite of what observability should do. (See the audit
in `RUST_DEV.md`: any silent regression is bad, but a noisy
observer that breaks compression is worse.)
"""
def record_compression(
self,
strategy: str,
original_tokens: int,
compressed_tokens: int,
) -> None: ...

27
headroom/transforms/smart_crusher.py
View file

@ -142,6 +142,7 @@ class SmartCrusher(Transform):
scorer: Any = None,
ccr_config: CCRConfig | None = None,
with_compaction: bool = True,
observer: Any = None,
):
# Hard import — no Python fallback. If the wheel is missing the
# caller must build it (scripts/build_rust_extension.sh) or
@ -157,6 +158,12 @@ class SmartCrusher(Transform):
cfg = config or SmartCrusherConfig()
self.config = cfg
self._with_compaction = with_compaction
# `observer`: see `headroom.transforms.observability`. The
# legacy proxy pipeline uses SmartCrusher.apply() directly
# (no ContentRouter); without an observer here, those
# compressions would be invisible to per-strategy metrics —
# exactly the silent-regression class we're guarding against.
self._observer = observer
# CCR config is preserved on `self` for callers that read it
# back (`headroom.proxy.server` does), but the Rust port doesn't
@ -334,6 +341,24 @@ class SmartCrusher(Transform):
return " ".join(context_parts)
def _notify_observer(self, original_tokens: int, compressed_tokens: int) -> None:
"""Forward a compression event to the configured
`CompressionObserver` (see `headroom.transforms.observability`).
No-op when no observer is set; swallows observer exceptions at
debug level so a buggy metrics impl doesn't break the
compression that just succeeded.
"""
if self._observer is None:
return
try:
self._observer.record_compression(
strategy="smart_crusher",
original_tokens=original_tokens,
compressed_tokens=compressed_tokens,
)
except Exception as e: # pragma: no cover - defensive
logger.debug("CompressionObserver raised (non-fatal): %s", e)
def apply(
self,
messages: list[dict[str, Any]],
@ -377,6 +402,7 @@ class SmartCrusher(Transform):
markers_inserted.append(marker)
if info:
transforms_applied.append(f"smart:{info}")
self._notify_observer(tokens, tokenizer.count_text(crushed))
# Anthropic-style: content is a list of blocks; each tool_result
# block has a string content field of its own.
@ -402,6 +428,7 @@ class SmartCrusher(Transform):
markers_inserted.append(marker)
if info:
transforms_applied.append(f"smart:{info}")
self._notify_observer(tokens, tokenizer.count_text(crushed))
if crushed_count > 0:
transforms_applied.insert(0, f"smart_crush:{crushed_count}")

314
tests/test_compression_observability.py Normal file
View file

@ -0,0 +1,314 @@
"""Per-strategy compression observability tests.
These guard the forcing function: when any compressor runs in
production, a `CompressionObserver` notification fires once per real
compression event, and `PrometheusMetrics` accumulates per-strategy
counters that the test suite asserts on directly.
The TOIN→SmartCrusher silent disconnect (caught three weeks late by
manual audit) was invisible because no signal distinguished by
strategy. These tests exist so the next regression of that shape
fails the suite the day it lands instead of waiting on an audit.
The counters live ONLY as in-process state on the metrics instance;
they are deliberately NOT exported through the Prometheus scrape or
OTel surface, because the metric→Supabase pipeline treats each
metric name as a column and we cannot add new columns. CI-level
observability via these tests is enough to catch silent regressions;
production export waits on a non-column-adding pipeline.
Coverage:
1. `ContentRouter.compress(...)` calls observer once per RoutingDecision.
2. `SmartCrusher.apply(...)` calls observer once per crushed message.
3. Both transforms tolerate an observer that raises (compression must
still succeed).
4. `PrometheusMetrics` correctly satisfies the `CompressionObserver`
protocol — `record_compression` increments per-strategy counters
and `tokens_saved_by_strategy` accumulates only positive savings.
5. The Prometheus scrape output (`export()`) does NOT emit any new
metric names — the per-strategy state stays internal.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
from headroom.transforms.content_detector import ContentType
from headroom.transforms.content_router import (
CompressionStrategy,
ContentRouter,
ContentRouterConfig,
RouterCompressionResult,
RoutingDecision,
)
from headroom.transforms.observability import CompressionObserver
from headroom.transforms.smart_crusher import SmartCrusher, SmartCrusherConfig
# ─── Test doubles ──────────────────────────────────────────────────────
@dataclass
class SpyObserver:
"""Captures every `record_compression` call for assertion."""
calls: list[tuple[str, int, int]] = field(default_factory=list)
def record_compression(
self,
strategy: str,
original_tokens: int,
compressed_tokens: int,
) -> None:
self.calls.append((strategy, original_tokens, compressed_tokens))
@dataclass
class ExplodingObserver:
"""Raises on every call. Used to assert observer failures don't
propagate out and break compression."""
raised: int = 0
def record_compression(self, *_a: Any, **_kw: Any) -> None:
self.raised += 1
raise RuntimeError("simulated observer outage")
# ─── Protocol conformance ──────────────────────────────────────────────
def test_spy_satisfies_observer_protocol():
spy = SpyObserver()
# `runtime_checkable` Protocol — isinstance check works.
assert isinstance(spy, CompressionObserver)
def test_prometheus_metrics_satisfies_observer_protocol():
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
assert isinstance(m, CompressionObserver)
# ─── ContentRouter wiring ──────────────────────────────────────────────
def test_content_router_records_observer_call_per_routing_decision():
spy = SpyObserver()
router = ContentRouter(ContentRouterConfig(), observer=spy)
# Forge a routing log directly via the result object — the observer
# call site walks `result.routing_log`, so we assert the contract
# without depending on which compressor would actually fire.
result = RouterCompressionResult(
compressed="x",
original="x",
strategy_used=CompressionStrategy.SMART_CRUSHER,
routing_log=[
RoutingDecision(
content_type=ContentType.JSON_ARRAY,
strategy=CompressionStrategy.SMART_CRUSHER,
original_tokens=200,
compressed_tokens=50,
),
RoutingDecision(
content_type=ContentType.SOURCE_CODE,
strategy=CompressionStrategy.CODE_AWARE,
original_tokens=300,
compressed_tokens=300, # passthrough — still recorded
),
],
)
router._observe(result)
assert spy.calls == [
("smart_crusher", 200, 50),
("code_aware", 300, 300),
]
def test_content_router_with_no_observer_is_silent():
router = ContentRouter(ContentRouterConfig()) # observer defaults None
result = RouterCompressionResult(
compressed="x",
original="x",
strategy_used=CompressionStrategy.PASSTHROUGH,
routing_log=[
RoutingDecision(
content_type=ContentType.PLAIN_TEXT,
strategy=CompressionStrategy.TEXT,
original_tokens=10,
compressed_tokens=5,
)
],
)
# Should not raise.
router._observe(result)
def test_content_router_swallows_observer_failures():
boom = ExplodingObserver()
router = ContentRouter(ContentRouterConfig(), observer=boom)
result = RouterCompressionResult(
compressed="x",
original="x",
strategy_used=CompressionStrategy.TEXT,
routing_log=[
RoutingDecision(
content_type=ContentType.PLAIN_TEXT,
strategy=CompressionStrategy.TEXT,
original_tokens=10,
compressed_tokens=5,
)
],
)
# Must not raise — observability failures are not compression failures.
router._observe(result)
assert boom.raised == 1
# ─── SmartCrusher wiring (legacy direct-pipeline path) ─────────────────
def _bigger_array(n: int = 60) -> str:
import json as _json
items = [{"status": "ok", "tag": "x", "n": i} for i in range(n)]
return _json.dumps(items)
def test_smart_crusher_apply_records_observer_per_crushed_message():
"""End-to-end: SmartCrusher.apply() walks messages, crushes the
big tool_result, fires the observer with strategy='smart_crusher'."""
from headroom.providers.openai import OpenAITokenCounter
from headroom.tokenizer import Tokenizer
spy = SpyObserver()
crusher = SmartCrusher(SmartCrusherConfig(), observer=spy)
tok = Tokenizer(OpenAITokenCounter("gpt-4o-mini"), model="gpt-4o-mini")
messages = [
{"role": "user", "content": "what's in the data?"},
{"role": "tool", "content": _bigger_array(60)},
]
result = crusher.apply(messages, tok)
# If the analyzer chose passthrough this run, the observer wasn't
# fired; that's fine for the wiring test — we only assert it WAS
# fired in the case it crushed.
if "smart_crush:" in ",".join(result.transforms_applied):
assert spy.calls, "smart_crusher crushed but observer wasn't notified"
for strategy, original, compressed in spy.calls:
assert strategy == "smart_crusher"
assert original > 0
assert compressed >= 0
def test_smart_crusher_apply_swallows_observer_failures():
"""Observer raises → compression still completes, returns valid
TransformResult, count of raises matches the crushed_count."""
from headroom.providers.openai import OpenAITokenCounter
from headroom.tokenizer import Tokenizer
boom = ExplodingObserver()
crusher = SmartCrusher(SmartCrusherConfig(), observer=boom)
tok = Tokenizer(OpenAITokenCounter("gpt-4o-mini"), model="gpt-4o-mini")
messages = [{"role": "tool", "content": _bigger_array(60)}]
result = crusher.apply(messages, tok)
# Either the analyzer didn't crush (boom.raised == 0) or it did
# (boom.raised >= 1) — but in both cases compression returned a
# valid TransformResult. No exception escaped.
assert result.messages is not None
# ─── PrometheusMetrics implementation ──────────────────────────────────
def test_prometheus_metrics_accumulates_per_strategy_counters():
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
m.record_compression("smart_crusher", original_tokens=200, compressed_tokens=50)
m.record_compression("smart_crusher", original_tokens=100, compressed_tokens=40)
m.record_compression("diff", original_tokens=80, compressed_tokens=80) # no savings
m.record_compression("code_aware", original_tokens=50, compressed_tokens=70) # negative savings
assert m.compressions_by_strategy == {
"smart_crusher": 2,
"diff": 1,
"code_aware": 1,
}
# Tokens saved is `max(0, original - compressed)` per strategy.
# smart_crusher: 150 + 60 = 210; diff: 0 (no savings, dict entry omitted);
# code_aware: 0 (negative).
assert m.tokens_saved_by_strategy == {"smart_crusher": 210}
def test_prometheus_export_does_not_leak_per_strategy_metrics():
"""Per-strategy state is tracked in-process only. The Prometheus
scrape output deliberately must NOT emit new metric names — the
metric→Supabase pipeline treats each metric name as a column, and
we cannot add new columns. This test guards that constraint: if a
future change adds the metric to the scrape, this fails and forces
a conscious decision."""
import asyncio
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
m.record_compression("smart_crusher", original_tokens=200, compressed_tokens=50)
m.record_compression("diff", original_tokens=120, compressed_tokens=70)
output = asyncio.run(m.export())
assert "headroom_compressions_total" not in output
assert "headroom_tokens_saved_by_strategy_total" not in output
# ─── End-to-end smoke (router + metrics together) ──────────────────────
def test_router_with_prometheus_observer_increments_counters():
"""Plumbing test: a router wired to a real PrometheusMetrics
instance lights up the per-strategy counters as routing decisions
accumulate. This is the production wiring shape from
`headroom/proxy/server.py`."""
from headroom.proxy.prometheus_metrics import PrometheusMetrics
m = PrometheusMetrics()
router = ContentRouter(ContentRouterConfig(), observer=m)
fake_result = RouterCompressionResult(
compressed="x",
original="x",
strategy_used=CompressionStrategy.MIXED,
routing_log=[
RoutingDecision(
content_type=ContentType.JSON_ARRAY,
strategy=CompressionStrategy.SMART_CRUSHER,
original_tokens=300,
compressed_tokens=80,
),
RoutingDecision(
content_type=ContentType.SOURCE_CODE,
strategy=CompressionStrategy.CODE_AWARE,
original_tokens=200,
compressed_tokens=120,
),
RoutingDecision(
content_type=ContentType.JSON_ARRAY,
strategy=CompressionStrategy.SMART_CRUSHER,
original_tokens=100,
compressed_tokens=40,
),
],
)
router._observe(fake_result)
assert m.compressions_by_strategy == {"smart_crusher": 2, "code_aware": 1}
assert m.tokens_saved_by_strategy == {
"smart_crusher": (300 - 80) + (100 - 40), # 280
"code_aware": (200 - 120), # 80
}