headroomlabs-ai
diff --git a/‎headroom/proxy/extensions.py‎
Lines changed: 35 additions & 0 deletions b/‎headroom/proxy/extensions.py‎
Lines changed: 35 additions & 0 deletions
diff --git a/‎headroom/proxy/handlers/gemini.py‎
Lines changed: 16 additions & 0 deletions b/‎headroom/proxy/handlers/gemini.py‎
Lines changed: 16 additions & 0 deletions
diff --git a/‎headroom/proxy/outcome.py‎
Lines changed: 21 additions & 2 deletions b/‎headroom/proxy/outcome.py‎
Lines changed: 21 additions & 2 deletions
diff --git a/‎headroom/proxy/savings_attribution.py‎
Lines changed: 132 additions & 4 deletions b/‎headroom/proxy/savings_attribution.py‎
Lines changed: 132 additions & 4 deletions
@@ -18,6 +18,41 @@
 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
 
@@ -316,6 +316,13 @@ async def handle_gemini_generate_content(
         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 @@ async def handle_google_cloudcode_stream(
         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 @@ async def handle_gemini_stream_generate_content(
         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 @@ async def handle_gemini_count_tokens(
         # 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,
 
@@ -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,
 
@@ -4,6 +4,7 @@
 
 import base64
 import json
+import math
 import re
 from collections.abc import MutableMapping
 from typing import Any
@@ -13,6 +14,43 @@
 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:"
+
+# NON-FINITE VALUES POISON EVERY CONSUMER DOWNSTREAM, and they do it long after
+# the call that introduced them. Starlette's JSONResponse encodes with
+# ``allow_nan=False``, so a single ``inf`` reaching ``/stats`` raises
+# ``ValueError: Out of range float values are not JSON compliant`` -- and the
+# value sits in the process-wide metrics totals, so the endpoint stays broken
+# until restart. Prometheus is no better: Python renders ``inf``, the exposition
+# format wants ``+Inf``, and the scrape fails to parse.
+#
+# The request itself still returns 200 throughout, which is the worst shape a
+# bug can have: the extension looks healthy while the operator's dashboard and
+# scrape are dead.
+#
+# One hour bounds a single stage inside one request -- unreachable in practice,
+# and it makes overflow-to-infinity on accumulation structurally impossible
+# (16 stages x 1h is nowhere near the float ceiling).
+MAX_STAGE_MS = 3_600_000.0
+
 
 def _source_name(value: object) -> str:
     name = _NAME_RE.sub("_", str(value or "other").strip().lower()).strip("_.-")
@@ -29,14 +67,25 @@ 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):
         ledger = []
         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 +96,63 @@ 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
+    # 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. Non-finite and
+    # absurdly large are not measurements either, and they break consumers
+    # rather than merely skewing them -- see MAX_STAGE_MS.
+    if not math.isfinite(elapsed) or not 0.0 < elapsed <= MAX_STAGE_MS:
+        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
+        # Re-checked rather than trusted: the ledger is a plain dict reachable
+        # through ``tags``, so a handler can be handed one this module never
+        # wrote. The guarantee has to hold at the read, not only at the write.
+        #
+        # Finiteness ONLY. ``MAX_STAGE_MS`` bounds a single sample at the write,
+        # where it prevents overflow; applying it here would test it against an
+        # ACCUMULATED total and silently discard a stage that legitimately ran
+        # for longer across many samples -- throwing away real data to guard
+        # against a value this path cannot produce.
+        if math.isfinite(elapsed) and elapsed > 0.0:
+            out[str(name)] = elapsed
+    return out
+
+
 def record_savings(
     tags: MutableMapping[str, Any],
     source: object,
@@ -61,12 +167,25 @@ def record_savings(
     ledger = _ledger(tags)
     if len(ledger) >= MAX_SOURCES:
         return
+    # Same hazard as MAX_STAGE_MS, on the amounts rather than the durations:
+    # ``usd=inf`` reaches ``/stats`` and raises out of the JSON encoder, and
+    # ``int(inf)`` raises OverflowError right here, inside the handler, on a
+    # request that would otherwise have succeeded. Neither is a saving, so
+    # neither is recorded -- the alternative is a plugin's arithmetic bug
+    # taking down an endpoint it has nothing to do with.
+    try:
+        amount = float(usd or 0.0)
+        count = int(tokens or 0)
+    except (TypeError, ValueError, OverflowError):
+        return
+    if not math.isfinite(amount):
+        return
     item: dict[str, Any] = {
         "source": _source_name(source),
         "realized": bool(realized),
         "estimated": bool(estimated),
-        "tokens": max(0, int(tokens or 0)),
-        "usd": round(float(usd or 0.0), 12),
+        "tokens": max(0, count),
+        "usd": round(amount, 12),
     }
     if details:
         item["details"] = {
@@ -84,8 +203,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: