diff --git a/headroom/proxy/extensions.py b/headroom/proxy/extensions.py index 9ead7a6cd..183028f71 100644 --- a/headroom/proxy/extensions.py +++ b/headroom/proxy/extensions.py @@ -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:`` 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 diff --git a/headroom/proxy/handlers/gemini.py b/headroom/proxy/handlers/gemini.py index 3311fe69c..dfdc3e90f 100644 --- a/headroom/proxy/handlers/gemini.py +++ b/headroom/proxy/handlers/gemini.py @@ -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, diff --git a/headroom/proxy/outcome.py b/headroom/proxy/outcome.py index f98c88753..d96190b4d 100644 --- a/headroom/proxy/outcome.py +++ b/headroom/proxy/outcome.py @@ -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, diff --git a/headroom/proxy/savings_attribution.py b/headroom/proxy/savings_attribution.py index db7f517e9..56e2fe209 100644 --- a/headroom/proxy/savings_attribution.py +++ b/headroom/proxy/savings_attribution.py @@ -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: diff --git a/tests/test_extension_attribution.py b/tests/test_extension_attribution.py new file mode 100644 index 000000000..58a227341 --- /dev/null +++ b/tests/test_extension_attribution.py @@ -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)] diff --git a/tests/test_gemini_nonjson_status.py b/tests/test_gemini_nonjson_status.py index 341fbaab2..2409b2390 100644 --- a/tests/test_gemini_nonjson_status.py +++ b/tests/test_gemini_nonjson_status.py @@ -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: diff --git a/tests/test_proxy_handler_helpers.py b/tests/test_proxy_handler_helpers.py index a11e437b3..e0e2a90ea 100644 --- a/tests/test_proxy_handler_helpers.py +++ b/tests/test_proxy_handler_helpers.py @@ -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="", diff --git a/tests/test_proxy_handlers_batch.py b/tests/test_proxy_handlers_batch.py index 53e823fcf..504e0c6f1 100644 --- a/tests/test_proxy_handlers_batch.py +++ b/tests/test_proxy_handlers_batch.py @@ -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