feat(proxy): let extensions report cost savings and their own latency

An extension that changes the bill has to be able to say so, or the operator
sees a different total with nothing to explain it. Two gaps stopped that.

SAVINGS WERE DROPPED ON GEMINI TRAFFIC
    `bind_scope` shares one attribution ledger between ASGI middleware and the
    request handler. Anthropic and OpenAI call it; Gemini never did, so
    anything an extension recorded into the request scope was discarded for
    Gemini traffic only -- silently, because an empty ledger and an unbound one
    are indistinguishable at the outcome funnel. Bound at all four Gemini tag
    sites.

AN EXTENSION'S OWN LATENCY WAS INVISIBLE
    `overhead_ms` is measured inside the handler, and an ASGI extension wraps
    that handler, so every millisecond it spends reaches the client while every
    timing surface stays flat. An extension that halves the bill and adds 200ms
    per request is a trade the operator has to see both halves of, and only one
    half was reaching the dashboard.

    `record_scope_timing(scope, stage, ms)` is the symmetric counterpart to the
    existing `record_scope_savings`, carried on the same bound ledger and merged
    into `RequestOutcome.pipeline_timing` at the outcome funnel -- one place, so
    every provider picks it up at once. It lands in `/stats.pipeline_timing`,
    the dashboard's Performance panel, and `headroom_transform_timing_ms_*`.

    Stage names are extension-supplied, so they are capped (16) and namespaced
    `ext:` -- `deep_copy` reported by a plugin must never accumulate into the
    same series as `deep_copy` measured by the pipeline. A handler's own timing
    wins a collision, which is unreachable while the prefix stands and is the
    safe way round if it ever goes.

Both ledgers are now stripped by `public_tags`: they ride on `tags` because
that is the one dict reaching the outcome funnel from every handler, and a list
and a dict must not land in a string-keyed label store.

`extensions.py` documents both calls. `record_scope_savings` already accepted
`usd`, which is the one channel in the proxy that can express savings WITHOUT
tokens -- routing a request to a cheaper model sends the same tokens for less
money -- but nothing in the extension-facing contract said so, and the module
is where extension authors look.

Real behavior proof in the PR body: a demo ASGI extension reporting
`tokens=0, usd=0.173` shows up on /stats as `savings.by_source`, in
`pipeline_timing` as `ext:routemegood`, and in /metrics as
`headroom_savings_attributed_usd_total{source="routemegood"} 0.519`.

Test suite: 10,989 passed, 578 skipped. The 3 failures on this branch
(test_graceful_shutdown ordering, test_learn integration, release-workflow
cargo) reproduce identically on clean main.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Tejas Chopra 2026-08-15 17:57:34 -07:00
parent 2f4d001c9f
commit c814b950c2
8 changed files with 469 additions and 4 deletions

View file

@ -18,6 +18,41 @@ Each ``install`` callable is invoked with the FastAPI ``app`` and the
OSS makes no assumptions about what extensions do. The interface is
deliberately minimal; extensions own the complexity behind it.
Reporting what an extension saved, and what it cost
---------------------------------------------------
An extension that changes the bill should say so, or the operator sees a
different total with nothing to attribute it to. Two calls, both taking the
ASGI ``scope`` so they work from middleware which runs outside the request
handler and has no other way in::
from headroom.proxy.savings_attribution import (
record_scope_savings, record_scope_timing,
)
record_scope_savings(scope, "my_extension", tokens=1200, usd=0.004)
record_scope_timing(scope, "my_extension", elapsed_ms)
``record_scope_savings`` takes ``tokens``, ``usd``, or both, so an extension
that saves money WITHOUT saving tokens routing a request to a cheaper model,
say can report a real number instead of a token count nobody saved. Pass
``realized=False`` for a projection rather than a measured amount; the two are
kept apart everywhere they surface. Savings land on ``/stats`` under
``savings.by_source``, on the dashboard as their own card, and in Prometheus as
``headroom_savings_attributed_usd_total{source=...}``. **Attribution only**
these rows explain the headline total, they are never added to it.
``record_scope_timing`` is the other half of the trade: an extension's own
latency, which is otherwise invisible because ``overhead_ms`` is measured
inside the handler that the extension wraps. It lands in ``/stats`` under
``pipeline_timing``, in the dashboard's Performance panel, and in
``headroom_transform_timing_ms_*``, namespaced ``ext:<name>`` so it can never
collide with a built-in transform.
Both are bounded (32 sources, 16 stages), never raise, and never change a
response telemetry from a plugin must not be able to break the request it is
describing.
**Extensions are opt-in.** Discovery enumerates every registered extension,
but ``install_all`` only invokes those explicitly enabled by the operator.
This protects users from silent behavior changes when a package they didn't

View file

@ -316,6 +316,13 @@ class GeminiHandlerMixin:
headers.pop("host", None)
headers.pop("content-length", None)
tags = extract_tags(headers)
# Anthropic and OpenAI bind here; Gemini did not, so anything an ASGI
# extension recorded into the request scope was dropped on the floor
# for Gemini traffic only — silently, because an empty ledger and an
# unbound one look identical at the outcome funnel.
from headroom.proxy.savings_attribution import bind_scope
bind_scope(tags, request.scope)
client = classify_client(headers)
# PR-A5 (P5-49): strip internal x-headroom-* from upstream-bound
# headers AFTER `_extract_tags` reads them. Memory user-id reads
@ -1027,6 +1034,9 @@ class GeminiHandlerMixin:
headers.pop("content-length", None)
headers.pop("accept-encoding", None)
tags = extract_tags(headers)
from headroom.proxy.savings_attribution import bind_scope
bind_scope(tags, request.scope)
# Note: streaming handlers delegate to _stream_response, which
# does its own classify_client. No need to compute here.
is_antigravity = self._is_cloudcode_antigravity_request(body, headers)
@ -1180,6 +1190,9 @@ class GeminiHandlerMixin:
headers.pop("host", None)
headers.pop("content-length", None)
tags = extract_tags(headers)
from headroom.proxy.savings_attribution import bind_scope
bind_scope(tags, request.scope)
# Streaming variant — delegates to _stream_response which
# classifies the client itself from headers.
# PR-A5 (P5-49): strip internal x-headroom-* before forwarding upstream.
@ -1328,6 +1341,9 @@ class GeminiHandlerMixin:
# outcome. Extract here so apply_to_tags below has a dict to
# mutate and the outcome at end-of-call inherits the tag.
tags = extract_tags(request.headers)
from headroom.proxy.savings_attribution import bind_scope
bind_scope(tags, request.scope)
_decision = CompressionDecision.decide(
headers=request.headers,
config=self.config,

View file

@ -399,7 +399,12 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
from headroom.proxy.cost import _summarize_transforms
from headroom.proxy.models import RequestLog
from headroom.proxy.project_context import get_current_project
from headroom.proxy.savings_attribution import encode, from_tags, public_tags
from headroom.proxy.savings_attribution import (
encode,
from_tags,
public_tags,
timings_from_tags,
)
from headroom.telemetry.session import record_outcome
# GitHub Copilot: requests routed to the Copilot API travel on the OpenAI or
@ -467,6 +472,20 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
tool_search_saved = tool_schema_saved_from_tags(outcome.tags or {})
savings_breakdown = from_tags(outcome.tags)
# Stage timings contributed from OUTSIDE the handler, folded in here rather
# than in each handler so every provider picks them up from one place.
#
# The handler's own timings win a name collision, which cannot happen while
# extension stages carry the ``ext:`` prefix but is the safe way round if
# that ever changes: a plugin must not be able to overwrite a measurement
# the pipeline made of itself.
extension_timing = timings_from_tags(outcome.tags)
pipeline_timing = (
{**extension_timing, **(outcome.pipeline_timing or {})}
if extension_timing
else outcome.pipeline_timing
)
# Billed input volume. Prefer the provider's own count where it reported one
# — that is what the invoice charges for, and it is the number cache math is
# already expressed in. Falls back to our local ``optimized_tokens`` when the
@ -488,7 +507,7 @@ async def emit_request_outcome(handler: Any, outcome: RequestOutcome) -> None:
cached=outcome.cache_hit,
overhead_ms=outcome.overhead_ms,
ttfb_ms=outcome.ttfb_ms,
pipeline_timing=outcome.pipeline_timing,
pipeline_timing=pipeline_timing,
waste_signals=outcome.waste_signals,
cache_read_tokens=outcome.cache_read_tokens,
cache_write_tokens=outcome.cache_write_tokens,

View file

@ -13,6 +13,26 @@ _NAME_RE = re.compile(r"[^a-z0-9_.-]+")
MAX_SOURCES = 32
_SCOPE_KEY = "headroom_savings_attribution"
# Per-request stage timings contributed from outside the handler, merged into
# ``RequestOutcome.pipeline_timing`` at the outcome funnel.
#
# An ASGI middleware wraps the handler, so every millisecond it spends lands in
# the client's latency while ``overhead_ms`` -- measured inside the handler --
# stays flat. An extension that halves the bill and adds 200ms per request is a
# trade the operator has to be able to see both halves of, and until now only
# one half reached the dashboard.
STAGE_TIMING_TAG = "_headroom_stage_timing"
_TIMING_SCOPE_KEY = "headroom_stage_timing"
# Stage names are extension-supplied, so they are capped like every other
# client-influenced label in this proxy (see MAX_DISTINCT_MODELS).
MAX_STAGES = 16
# Namespace, so an extension can never shadow a built-in transform's timing --
# ``deep_copy`` reported by a plugin and ``deep_copy`` reported by the pipeline
# must not accumulate into the same series.
STAGE_PREFIX = "ext:"
def _source_name(value: object) -> str:
name = _NAME_RE.sub("_", str(value or "other").strip().lower()).strip("_.-")
@ -29,7 +49,12 @@ def _ledger(tags: MutableMapping[str, Any]) -> list[dict[str, Any]]:
def bind_scope(tags: MutableMapping[str, Any], scope: MutableMapping[str, Any]) -> None:
"""Share one ledger between ASGI middleware and the request handler."""
"""Share the savings and timing ledgers between ASGI middleware and the handler.
Both are bound together because an extension that reports one usually
reports the other, and a handler that binds only savings would drop the
timings silently -- which is the failure this call is here to prevent.
"""
state = scope.setdefault("state", {})
ledger = state.get(_SCOPE_KEY)
if not isinstance(ledger, list):
@ -37,6 +62,12 @@ def bind_scope(tags: MutableMapping[str, Any], scope: MutableMapping[str, Any])
state[_SCOPE_KEY] = ledger
tags[SAVINGS_ATTRIBUTION_TAG] = ledger
timings = state.get(_TIMING_SCOPE_KEY)
if not isinstance(timings, dict):
timings = {}
state[_TIMING_SCOPE_KEY] = timings
tags[STAGE_TIMING_TAG] = timings
def record_scope_savings(scope: MutableMapping[str, Any], source: object, **values: Any) -> None:
state = scope.setdefault("state", {})
@ -47,6 +78,52 @@ def record_scope_savings(scope: MutableMapping[str, Any], source: object, **valu
record_savings({SAVINGS_ATTRIBUTION_TAG: ledger}, source, **values)
def record_scope_timing(scope: MutableMapping[str, Any], stage: object, ms: float) -> None:
"""Attribute milliseconds spent outside the handler to a named stage.
Additive within one request, so a middleware that works in two passes
(before and after ``call_next``) reports each and gets their sum. Never
raises and never changes a response: a plugin's telemetry must not be able
to break the request it is describing.
"""
try:
elapsed = float(ms)
except (TypeError, ValueError):
return
if not elapsed > 0.0:
# Non-positive is either a clock artifact or nothing happening. Either
# way it is not a measurement, and averaging it in would drag the mean
# toward zero exactly where the stage is cheapest to ignore.
return
state = scope.setdefault("state", {})
timings = state.get(_TIMING_SCOPE_KEY)
if not isinstance(timings, dict):
timings = {}
state[_TIMING_SCOPE_KEY] = timings
name = STAGE_PREFIX + _source_name(stage)
if name not in timings and len(timings) >= MAX_STAGES:
return
timings[name] = round(float(timings.get(name, 0.0)) + elapsed, 4)
def timings_from_tags(tags: MutableMapping[str, Any] | None) -> dict[str, float]:
"""Extension stage timings carried on the request's tags, if any."""
raw = (tags or {}).get(STAGE_TIMING_TAG)
if not isinstance(raw, dict):
return {}
out: dict[str, float] = {}
for name, value in list(raw.items())[:MAX_STAGES]:
try:
elapsed = float(value)
except (TypeError, ValueError):
continue
if elapsed > 0.0:
out[str(name)] = elapsed
return out
def record_savings(
tags: MutableMapping[str, Any],
source: object,
@ -84,8 +161,17 @@ def from_tags(tags: MutableMapping[str, Any] | None) -> list[dict[str, Any]]:
return [dict(item) for item in raw[:MAX_SOURCES] if isinstance(item, dict)]
_INTERNAL_TAGS = frozenset({SAVINGS_ATTRIBUTION_TAG, STAGE_TIMING_TAG})
def public_tags(tags: MutableMapping[str, Any] | None) -> dict[str, Any]:
return {key: value for key, value in (tags or {}).items() if key != SAVINGS_ATTRIBUTION_TAG}
"""Tags minus the internal ledgers, which are structures rather than labels.
They are carried on ``tags`` because that is the one dict that reaches the
outcome funnel from every handler; letting them through to ``RequestLog``
would put a list and a dict into a string-keyed label store.
"""
return {key: value for key, value in (tags or {}).items() if key not in _INTERNAL_TAGS}
def encode(items: list[dict[str, Any]]) -> str:

View file

@ -0,0 +1,304 @@
"""Attribution and timing contributed by proxy extensions.
An extension that changes the bill has to be able to say so, or the operator
sees a different total with nothing to explain it. The savings half of this
already existed but only reached two of the three handler families; the timing
half did not exist at all, so an extension's own latency was invisible —
``overhead_ms`` is measured inside the handler that the extension wraps.
"""
from __future__ import annotations
import pytest
from headroom.proxy.savings_attribution import (
MAX_STAGES,
SAVINGS_ATTRIBUTION_TAG,
STAGE_PREFIX,
STAGE_TIMING_TAG,
bind_scope,
from_tags,
public_tags,
record_scope_savings,
record_scope_timing,
timings_from_tags,
)
def _scope() -> dict:
return {"type": "http", "method": "POST"}
# --- savings, from middleware ------------------------------------------------
def test_middleware_savings_reach_the_handlers_tags() -> None:
"""The contract: middleware records into the scope before the handler runs,
the handler binds, and the outcome funnel reads one ledger."""
scope = _scope()
record_scope_savings(scope, "routemegood", usd=0.42)
tags: dict = {}
bind_scope(tags, scope)
(row,) = from_tags(tags)
assert row["source"] == "routemegood"
assert row["usd"] == 0.42
def test_savings_can_be_money_without_being_tokens() -> None:
"""The gap this closes. Every other savings channel computes
``saved = before - after`` and three of them refuse a non-positive value,
so an extension that routes a request to a cheaper model same tokens,
smaller bill could only report by inventing a token count nobody saved."""
scope = _scope()
record_scope_savings(scope, "model_router", tokens=0, usd=1.75)
tags: dict = {}
bind_scope(tags, scope)
(row,) = from_tags(tags)
assert row["tokens"] == 0
assert row["usd"] == 1.75
def test_a_projection_is_not_a_measurement() -> None:
scope = _scope()
record_scope_savings(scope, "guess", usd=1.0, realized=False)
record_scope_savings(scope, "guess", usd=1.0, realized=True)
tags: dict = {}
bind_scope(tags, scope)
assert sorted(row["realized"] for row in from_tags(tags)) == [False, True]
# --- timing ------------------------------------------------------------------
def test_middleware_timing_reaches_the_handlers_tags() -> None:
scope = _scope()
record_scope_timing(scope, "routemegood", 12.5)
tags: dict = {}
bind_scope(tags, scope)
assert timings_from_tags(tags) == {f"{STAGE_PREFIX}routemegood": 12.5}
def test_timing_is_additive_within_one_request() -> None:
"""A middleware works in two passes — before ``call_next`` and after — and
should be able to report each without tracking the total itself."""
scope = _scope()
record_scope_timing(scope, "ext", 4.0)
record_scope_timing(scope, "ext", 2.5)
tags: dict = {}
bind_scope(tags, scope)
assert timings_from_tags(tags) == {f"{STAGE_PREFIX}ext": 6.5}
def test_extension_stages_are_namespaced() -> None:
"""``deep_copy`` reported by a plugin and ``deep_copy`` measured by the
pipeline must not accumulate into the same series."""
scope = _scope()
record_scope_timing(scope, "deep_copy", 1.0)
tags: dict = {}
bind_scope(tags, scope)
assert list(timings_from_tags(tags)) == [f"{STAGE_PREFIX}deep_copy"]
@pytest.mark.parametrize("bad", [0, -1.0, None, "slow", float("nan")])
def test_a_non_measurement_is_not_recorded(bad) -> None:
"""Zero and negative are clock artifacts, not observations; averaging them
in would drag the mean down exactly where the stage is cheapest to skip."""
scope = _scope()
record_scope_timing(scope, "ext", bad)
tags: dict = {}
bind_scope(tags, scope)
assert timings_from_tags(tags) == {}
def test_stage_cardinality_is_capped() -> None:
"""Stage names are extension-supplied, so they are bounded like every other
client-influenced label in this proxy."""
scope = _scope()
for i in range(MAX_STAGES * 4):
record_scope_timing(scope, f"stage-{i}", 1.0)
tags: dict = {}
bind_scope(tags, scope)
assert len(timings_from_tags(tags)) == MAX_STAGES
def test_an_existing_stage_still_accumulates_at_the_cap() -> None:
"""The cap bounds distinct names, not measurements. A stage already being
tracked must keep accumulating or its total silently stops growing."""
scope = _scope()
for i in range(MAX_STAGES):
record_scope_timing(scope, f"stage-{i}", 1.0)
record_scope_timing(scope, "stage-0", 5.0)
tags: dict = {}
bind_scope(tags, scope)
assert timings_from_tags(tags)[f"{STAGE_PREFIX}stage-0"] == 6.0
def test_recording_before_any_bind_still_works() -> None:
"""Ordering is not guaranteed: middleware runs first, and on a path where
the handler never binds, nothing should raise."""
scope = _scope()
record_scope_timing(scope, "ext", 1.0)
record_scope_savings(scope, "ext", usd=1.0)
assert scope["state"]
def test_recording_after_bind_is_seen_by_the_already_bound_tags() -> None:
"""A middleware measures its own post-response work AFTER the handler has
bound. Sharing one object rather than copying is what makes that land."""
tags: dict = {}
scope = _scope()
bind_scope(tags, scope)
record_scope_timing(scope, "ext", 3.0)
record_scope_savings(scope, "ext", usd=0.5)
assert timings_from_tags(tags) == {f"{STAGE_PREFIX}ext": 3.0}
assert from_tags(tags)[0]["usd"] == 0.5
def test_bind_is_idempotent() -> None:
tags: dict = {}
scope = _scope()
bind_scope(tags, scope)
record_scope_timing(scope, "ext", 1.0)
bind_scope(tags, scope)
record_scope_timing(scope, "ext", 1.0)
assert timings_from_tags(tags) == {f"{STAGE_PREFIX}ext": 2.0}
def test_timings_from_tags_tolerates_junk() -> None:
for junk in (
None,
{},
{STAGE_TIMING_TAG: "nope"},
{STAGE_TIMING_TAG: []},
{STAGE_TIMING_TAG: {"a": "b"}},
):
assert timings_from_tags(junk) == {}
# --- the ledgers are structures, not labels ---------------------------------
def test_neither_ledger_leaks_into_request_log_tags() -> None:
"""They ride on ``tags`` because that is the one dict reaching the outcome
funnel from every handler. A list and a dict must not land in a
string-keyed label store."""
tags: dict = {"client": "claude-code"}
scope = _scope()
bind_scope(tags, scope)
record_scope_savings(scope, "ext", usd=1.0)
record_scope_timing(scope, "ext", 1.0)
assert public_tags(tags) == {"client": "claude-code"}
assert SAVINGS_ATTRIBUTION_TAG not in public_tags(tags)
assert STAGE_TIMING_TAG not in public_tags(tags)
# --- through the outcome funnel ---------------------------------------------
pytest.importorskip("fastapi")
class _Harness:
"""Just enough of HeadroomProxy to drive the real funnel method.
Mirrors ``tests/test_request_outcome.py::_FunnelHarness`` the real
implementation is bound to the harness, so nothing under test is mocked.
"""
def __init__(self) -> None:
from unittest.mock import AsyncMock, MagicMock
from headroom.proxy.server import HeadroomProxy
self.metrics = MagicMock()
self.metrics.record_request = AsyncMock()
self.cost_tracker = MagicMock()
self.logger = None
self._record_request_outcome = HeadroomProxy._record_request_outcome.__get__(
self, type(self)
)
def _outcome(**overrides):
from headroom.proxy.outcome import RequestOutcome
defaults = {
"request_id": "req-1",
"provider": "anthropic",
"model": "claude-sonnet-4",
"original_tokens": 1000,
"optimized_tokens": 1000,
"output_tokens": 50,
"tokens_saved": 0,
"attempted_input_tokens": 1000,
}
defaults.update(overrides)
return RequestOutcome(**defaults)
@pytest.mark.asyncio
async def test_extension_timing_reaches_pipeline_timing() -> None:
"""The whole point of the timing half: ``pipeline_timing`` is what
``/stats``, the dashboard's Performance panel and
``headroom_transform_timing_ms_*`` are all built on."""
scope = _scope()
record_scope_timing(scope, "routemegood", 8.0)
tags: dict = {}
bind_scope(tags, scope)
h = _Harness()
await h._record_request_outcome(_outcome(tags=tags, pipeline_timing={"deep_copy": 1.0}))
timing = h.metrics.record_request.await_args.kwargs["pipeline_timing"]
assert timing == {"deep_copy": 1.0, f"{STAGE_PREFIX}routemegood": 8.0}
@pytest.mark.asyncio
async def test_a_handler_timing_wins_a_name_collision() -> None:
"""Namespacing makes this unreachable today; it is asserted so that if the
prefix ever goes, a plugin still cannot overwrite a measurement the
pipeline made of itself."""
tags = {STAGE_TIMING_TAG: {"deep_copy": 99.0}}
h = _Harness()
await h._record_request_outcome(_outcome(tags=tags, pipeline_timing={"deep_copy": 1.0}))
timing = h.metrics.record_request.await_args.kwargs["pipeline_timing"]
assert timing["deep_copy"] == 1.0
@pytest.mark.asyncio
async def test_no_extension_timing_leaves_pipeline_timing_untouched() -> None:
"""Including identity: a request with no extension must pass the handler's
own dict through, not a rebuilt copy of it."""
original = {"deep_copy": 1.0}
h = _Harness()
await h._record_request_outcome(_outcome(pipeline_timing=original))
assert h.metrics.record_request.await_args.kwargs["pipeline_timing"] is original
@pytest.mark.asyncio
async def test_extension_savings_reach_the_metrics_call() -> None:
scope = _scope()
record_scope_savings(scope, "routemegood", usd=0.42, tokens=0)
tags: dict = {}
bind_scope(tags, scope)
h = _Harness()
await h._record_request_outcome(_outcome(tags=tags))
attribution = h.metrics.record_request.await_args.kwargs["savings_attribution"]
assert [(row["source"], row["usd"]) for row in attribution] == [("routemegood", 0.42)]

View file

@ -15,6 +15,7 @@ class _FakeRequest:
self.headers: dict[str, str] = {}
self.query_params: dict[str, str] = {}
self.url = SimpleNamespace(path="/v1beta/models/gemini-pro:generateContent", query="")
self.scope: dict = {"type": "http", "method": "POST"}
class _NonJsonResponse:

View file

@ -117,6 +117,7 @@ class _VertexGeminiImageRequest:
method = "POST"
headers = {}
query_params = {}
scope: dict = {"type": "http", "method": "POST"}
url = SimpleNamespace(
path="/v1/projects/p/locations/us-central1/publishers/google/models/gemini-2.0-flash:generateContent",
query="",

View file

@ -156,6 +156,9 @@ class FakeRequest:
self.method = method
self.url = SimpleNamespace(path=path, query=query)
self.query_params = {}
# Every real Starlette Request has one, and handlers now share a
# per-request attribution ledger through it (savings_attribution).
self.scope: dict = {"type": "http", "method": method}
async def body(self) -> bytes:
return self._body