mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
f9807fd69e
## What
Two changes that let a proxy extension report **what it saved** and
**what it cost**, so both show up under `/stats`, the dashboard, and
Prometheus.
`record_scope_savings` already existed and already accepted `usd` — the
one channel in the proxy that can express savings *without* tokens. Two
things stopped it working end to end.
### 1. Savings were silently dropped on Gemini traffic (bug)
`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. Now bound at
all four Gemini tag sites.
### 2. An extension's own latency was invisible (gap)
`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 200 ms 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
`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.
## API surface
```python
from headroom.proxy.savings_attribution import record_scope_savings, record_scope_timing
record_scope_savings(scope, "my_extension", tokens=0, usd=0.004) # money without tokens
record_scope_timing(scope, "my_extension", elapsed_ms)
```
Both take the ASGI `scope`, because middleware has no other way in.
Documented in `extensions.py` — the module extension authors actually
read, and the stability contract for this interface.
- Savings → `/stats` `savings.by_source`, dashboard card,
`headroom_savings_attributed_usd_total{source=...}`
- Timing → `/stats` `pipeline_timing`, dashboard Performance panel,
`headroom_transform_timing_ms_*`
**Attribution only.** These rows explain the headline total; they are
never added to it.
## Changes to existing behavior
- `public_tags` now strips `_headroom_stage_timing` as well as
`_headroom_savings_attribution`. Both 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.
- `pipeline_timing` passed to `metrics.record_request` is merged rather
than passed through **only when an extension contributed timings**; with
no extension the handler's own dict is passed through unchanged
(asserted by identity in the tests).
- Stage names are extension-supplied, so they are capped at 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 (unreachable while the prefix
stands; the safe way round if it ever goes).
## Failure modes
Both calls 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. Non-positive and non-numeric durations are
ignored: a zero is a clock artifact, not an observation, and averaging
it in would drag the mean down exactly where the stage is cheapest to
skip. `timings_from_tags` tolerates junk on the tag.
## Test-double fix
Three Gemini test fakes (`FakeRequest`, `_FakeRequest`,
`_VertexGeminiImageRequest`) had no `.scope`, which every real Starlette
`Request` has. They now do. This is a double that had drifted from the
type it stands in for; the alternative was weakening the handler to
tolerate a request shape that cannot occur in production.
---
## Real behavior proof
**Setup:** macOS 15.4 (darwin 25.4.0), Python 3.12.13, this branch at
`c814b950`, real `create_app` proxy with `respx`-mocked Anthropic
upstream, a demo ASGI extension added via `app.add_middleware`.
**The extension** — written as a third party would, reporting `tokens=0`
because it re-routed `claude-opus-5` → `claude-haiku-4-5`: same tokens,
cheaper model. That is precisely the case no existing Headroom savings
channel can express, since all of them compute `saved = before - after`.
```python
class DemoRouter:
def __init__(self, app): self.app = app
async def __call__(self, scope, receive, send):
if scope.get("type") != "http":
return await self.app(scope, receive, send)
started = time.perf_counter()
record_scope_savings(scope, "routemegood", tokens=0, usd=0.173)
record_scope_timing(scope, "routemegood", (time.perf_counter() - started) * 1000)
await self.app(scope, receive, send)
```
**Ran:** three POSTs to `/v1/messages`, then `GET /stats` and `GET
/metrics`.
**Observed:**
```
upstream call -> 200
upstream call -> 200
upstream call -> 200
=== /stats savings.by_source (what the dashboard renders) ===
[
{
"source": "routemegood",
"realized": true,
"events": 3,
"tokens": 0,
"usd": 0.519
}
]
=== /stats pipeline_timing (dashboard Performance panel) ===
{
"ext:routemegood": {
"average_ms": 0.01,
"max_ms": 0.02,
"count": 3
}
}
=== /metrics ===
# HELP headroom_savings_attributed_tokens_total Tokens attributed to a savings source
# TYPE headroom_savings_attributed_tokens_total counter
headroom_savings_attributed_tokens_total{realized="true",source="routemegood"} 0
# HELP headroom_savings_attributed_usd_total Cost savings attributed to a source; may be negative
# TYPE headroom_savings_attributed_usd_total gauge
headroom_savings_attributed_usd_total{realized="true",source="routemegood"} 0.519
headroom_transform_timing_ms_sum{transform="ext:routemegood"} 0.03
```
`$0.519 = 3 × $0.173` — three requests, correctly accumulated, with
`tokens: 0` throughout.
**Also have (not a substitute for the above):** 22 new unit tests in
`tests/test_extension_attribution.py`, including four that drive the
real `_record_request_outcome` funnel via the same descriptor-binding
harness `test_request_outcome.py` uses.
Full suite on this branch: **10,989 passed, 578 skipped**. Three
failures —
`test_graceful_shutdown.py::test_run_server_installs_cancelled_error_filter`
(full-suite ordering; passes in isolation),
`test_learn/test_integration.py::TestCodexIntegration::test_full_pipeline`,
and `test_release_workflows.py::test_no_native_tls_in_wheel_build_tree`
(needs `cargo`) — **reproduce identically on clean `main`** (`2f4d001c`,
10,967 passed, same 3 failed). Verified by stashing this branch and
re-running the full suite on main in the same tree.
**What I did not test:** a live provider (upstream is `respx`-mocked);
the Gemini `bind_scope` fix against real Google traffic (covered by the
existing 114 Gemini tests, which all pass); the dashboard rendered in a
browser — I verified the JSON shape its templates bind to
(`stats.savings?.by_source`, `stats.pipeline_timing`) rather than the
pixels.
---
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Tejas Chopra <tejas@Tejass-MacBook-Pro.local>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
93 lines
3.2 KiB
Python
93 lines
3.2 KiB
Python
"""A non-JSON Gemini upstream body must not be masked as a synthetic 502."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from headroom.proxy.handlers.gemini import GeminiHandlerMixin
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self) -> None:
|
|
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:
|
|
status_code = 503
|
|
content = b"<html>temporarily unavailable</html>"
|
|
headers = {"content-type": "text/html", "content-length": str(len(content))}
|
|
|
|
def json(self) -> object:
|
|
raise json.JSONDecodeError("not json", self.content.decode("utf-8"), 0)
|
|
|
|
|
|
class _FakeMetrics:
|
|
def __init__(self) -> None:
|
|
self.failed: list[str] = []
|
|
|
|
async def record_failed(self, *, provider: str, model: str = "") -> None:
|
|
self.failed.append(f"{provider}:{model}")
|
|
|
|
|
|
class _Handler(GeminiHandlerMixin):
|
|
GEMINI_API_URL = "https://gemini.example"
|
|
|
|
def __init__(self) -> None:
|
|
self.memory_handler = None
|
|
self.rate_limiter = None
|
|
self.usage_reporter = None
|
|
self.config = SimpleNamespace(
|
|
optimize=False,
|
|
anthropic_pre_upstream_memory_context_timeout_seconds=0.1,
|
|
)
|
|
self.metrics = _FakeMetrics()
|
|
self.outcomes = []
|
|
|
|
async def _next_request_id(self) -> str:
|
|
return "req-1"
|
|
|
|
async def _retry_request(self, method, url, headers, body): # noqa: ANN001, ANN201
|
|
return _NonJsonResponse()
|
|
|
|
async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001
|
|
self.outcomes.append(outcome)
|
|
|
|
async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201
|
|
# Test stub for HeadroomProxy._count_tokens_offloaded: resolve the
|
|
# tokenizer and count inline (the real method offloads to the executor).
|
|
from headroom.tokenizers import get_tokenizer
|
|
|
|
tokenizer = get_tokenizer(model)
|
|
return tokenizer, tokenizer.count_messages(messages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_forwards_non_json_upstream_status(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def payload(request): # noqa: ANN001, ANN201
|
|
return {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}
|
|
|
|
class _Tokenizer:
|
|
def count_messages(self, messages): # noqa: ANN001, ANN201
|
|
return 7
|
|
|
|
monkeypatch.setattr("headroom.proxy.helpers._read_request_json", payload)
|
|
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _Tokenizer())
|
|
|
|
handler = _Handler()
|
|
response = await handler.handle_gemini_generate_content(_FakeRequest(), "gemini-pro")
|
|
|
|
assert response.status_code == 503
|
|
assert response.body == _NonJsonResponse.content
|
|
assert response.headers["content-type"] == "text/html"
|
|
assert response.headers["x-headroom-tokens-before"] == "7"
|
|
assert response.headers["x-headroom-tokens-after"] == "7"
|
|
assert handler.metrics.failed == []
|
|
assert handler.outcomes[0].status_code == 503
|