mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
2f4d001c9f
commit
c814b950c2
8 changed files with 469 additions and 4 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
304
tests/test_extension_attribution.py
Normal file
304
tests/test_extension_attribution.py
Normal 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)]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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="",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue