headroom/tests/test_request_outcome.py
inix 4aac068814
fix(proxy/metrics): move the savings-ledger append off the event loop (#2439)
## Description

`PrometheusMetrics.record_request` appends one durable JSONL event per
compressed request. That append is synchronous: `open` + `fcntl.flock` +
`write`, plus a full-file rewrite once the ledger passes 1 MB. It runs
on the event loop, inside `self._lock`.

`export()` takes that same lock and holds it for the entire Prometheus
serialization, so a slow ledger write stops `/metrics` cold. In a repro
run of 200 compressed requests, `/metrics` completed zero scrapes and
the event loop never yielded once across 6.4 seconds.

The append now runs in a thread, outside the lock. `savings_ledger`
already takes its own `flock` across processes, so the metrics lock was
never what made the write safe.

Both halves are one change. Awaiting inside the lock would hold it for
the whole write rather than just the syscall, which is worse than what
is on main today.

The file already documents this hazard against itself.
`record_stage_timings` (`prometheus_metrics.py:867-874`) picks a plain
`threading.Lock` over `self._lock` specifically because "the async lock
is also held by `export()` during Prometheus scrapes." The ledger append
was the pattern that docstring warns about.

No filed issue for this one.

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- Move the `savings_ledger.record_savings_event` call in
`record_request` out of `async with self._lock` and run it through
`asyncio.to_thread`. The call site keeps its keyword arguments verbatim;
`to_thread` forwards `**kwargs`, so no `functools.partial` wrapper is
needed.
- Keep the `await`. Callers still see the event on disk when
`record_request` returns, which
`tests/test_savings_ledger_before_forwarded.py` asserts synchronously.
- Add `tests/test_savings_ledger_offload.py`: lock scope, event-loop
responsiveness, durability on return, and both arms of the `tokens_saved
> 0 and not stateless` gate.

`savings_ledger.py` is untouched. It stays synchronous so the MCP
`headroom_compress` caller in `ccr/mcp_server.py:789` does not have to
change.

Sizing the executor is left alone on purpose. `asyncio.to_thread` uses
the default pool, which is the documented tool for blocking I/O and
already the idiom here (`helpers.py:1297`, `server.py:1694`, `:3557`,
`:3613`, `:4244`). The compression pools are sized `max(1,
os.cpu_count())` for CPU-bound work, and `PrometheusMetrics` holds no
reference to `HeadroomProxy` anyway, so reaching them would mean a new
constructor parameter.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed

### Test Output

```text
$ pytest tests/test_savings_ledger_offload.py tests/test_savings_ledger.py tests/test_savings_ledger_before_forwarded.py -q
======================== 26 passed, 1 warning in 4.08s =========================

$ ruff check . && ruff format --check headroom/proxy/prometheus_metrics.py tests/test_savings_ledger_offload.py
All checks passed!
2 files already formatted

$ mypy headroom/proxy/prometheus_metrics.py
Success: no issues found in 1 source file
```

Broader sweep across the blast radius, 145 test files matching savings /
metrics / outcome / stats / proxy / handler / server / ledger / cost /
prometheus, each run under a per-file wall-clock watchdog:

```text
138 files pass, 1470 tests passed
7 non-green:
  HANG tests/test_agent_savings.py
  HANG tests/test_ccr_mcp_server.py
  HANG tests/test_netcost_gate.py
  HANG tests/test_proxy_compress_endpoint.py
  HANG tests/test_proxy_mode_benchmark.py
  HANG tests/test_read_maturation_handler_nobust.py
  FAIL tests/test_proxy_copilot_auth_hooks.py::test_openai_passthrough_applies_copilot_auth

Same 7 files re-run with headroom/proxy/prometheus_metrics.py reverted to c400f908:
  identical set, identical failure. diff of the two non-green lists is empty.
```

The before and after non-green sets match exactly, so nothing here is a
regression from this PR. See Additional Notes for the hang.

## Real Behavior Proof

- Environment: macOS 15.4 (Darwin 25.4.0) arm64, Python 3.13.13,
uv-managed venv, git worktree at `upstream/main` `c400f908`.
- Exact command / steps: a standalone asyncio script, not the unit
tests. It builds a real `PrometheusMetrics` (no injected tracker, so it
self-constructs with `save_flush_every=PROXY_SAVINGS_FLUSH_EVERY`
exactly as the proxy does) against a real on-disk ledger pre-seeded to
3.00 MB so `_maybe_compact`'s full-file rewrite actually fires. It then
drives 200 `record_request` calls at concurrency 16 while a `/metrics`
scraper calls `export()` every 20 ms and a canary coroutine ticks every
5 ms. Ran twice from the same script: once with
`headroom/proxy/prometheus_metrics.py` reverted to `c400f908`, once with
this change. Seeding the ledger past 1 MB is the part that matters. On a
fresh ledger the write is microseconds, compaction never fires, and the
run shows no delta at all.
- Observed result: before, `/metrics` completed 0 scrapes and the canary
ticked once in 6419 ms. After, 206 scrapes at p50 0.1 ms and max 0.2 ms,
and 418 canary ticks with a 92.3 ms worst gap. Total wall clock barely
moved, 6419 ms to 6473 ms, which is the expected result and not a null
one: the same disk work still serializes on the ledger's own `flock`,
now in a thread instead of on the loop. Unit-test view of the same
behavior, with a 500 ms stub standing in for the write: before, `event
loop stalled 0.506s during a 0.500s ledger write` and the competing lock
holder waited `+0.502s`; after, both pass.
- Not tested: Windows, where `savings_ledger` already skips locking
because `fcntl` is unavailable. Multi-process contention on one ledger
file, which this change does not alter. The residual 92.3 ms loop gap
after the fix, which traces to `SavingsTracker._save_locked`'s
`os.fsync` (`savings_tracker.py:1445`) firing every 25th request from
inside the same lock, a separate path this PR leaves alone.

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I did **not** edit `CHANGELOG.md` — it is generated by
release-please from my Conventional Commit PR title (a CI guard enforces
this)

## Screenshots (if applicable)

N/A, no user-visible surface.

## Additional Notes

Docs checklist item is N/A. Nothing user-facing moves; `headroom
savings` reads the same ledger, with the same contents, written from a
thread.

Two related in-lock costs on this same code path are deliberately out of
scope, one logical change per commit:
`_current_savings_tracker_totals()` at `:798` rebuilds
`cost_tracker.stats()` per request, and `_resolve_litellm_model` in
`savings_tracker.py` is uncached across roughly seven calls per request.
Happy to follow up on either.

`make ci-precheck` was not run end to end. The uv-managed worktree venv
has no `pip`, so the `ci-precheck-python` hook's `pip install -e .` step
fails on this machine for reasons unrelated to the change. Ran `pytest`,
`ruff`, and `mypy` directly instead, output above. No Rust touched.

One heads up worth passing on, since it is why the numbers above are a
sweep and not a single full-suite line. Seven test files do not complete
on this macOS box: six hang and one fails. The hangs park the main
thread in `_dispatch_semaphore_wait_slow` with CPU time frozen and never
recover, and `pytest-timeout --timeout-method=signal` cannot break them
out, so the block is native, below the interpreter.
`tests/test_adversarial_grid.py::test_grid_shape_and_schema` is the
first one a full run reaches.

All seven reproduce identically on unmodified `c400f908` with this
change reverted, so they predate the PR. I ran the reverted comparison
specifically to rule out a thread-before-fork interaction from the new
`to_thread` call, which was the plausible way this change could have
caused it. It did not. Happy to open a separate issue with the sample
output if that is useful.


---

## Follow-up: a cancellation bug this change introduced

Self-review turned up a second problem in this change, so the fix rides
along here.

Moving the append into a thread added the first suspension point in
`record_request` that can tear state. The metrics lock at
`prometheus_metrics.py:701` suspends too, but only under contention, and
it sits ahead of every mutation, so a cancellation there recorded
nothing at all. The new await is different. It sits after the Prometheus
counters commit and before OTel and the funnel's effects 2/3/4, and it
suspends on every compressed request.

Four of the funnel's call sites are `finally:` blocks inside streaming
async generators (`streaming.py:1611`, `:1859`, `:2069`,
`openai.py:8614`). A client disconnect cancels that task. The
cancellation lands on the new await, so Prometheus counts the request
while the cost tracker, the request log, and the PERF line `headroom
perf` reads never see it. `emit_request_outcome` has one try/except and
it sits before `record_request`, so nothing catches this.

`_record_request_outcome` now wraps the funnel in `asyncio.shield`. One
line at a single choke point, covering all 28 call sites. The shield
leaves the cancellation itself alone: the await still raises
`CancelledError`, so generator teardown propagates as before. Only the
bookkeeping survives.

Real stack, uvicorn 0.40.0 + starlette 1.3.1, raw-socket disconnect
mid-stream:

| effect | before | after |
|---|---|---|
| Prometheus counters | committed | committed |
| ledger write | ran | ran |
| OTel | **skipped** | ran |
| cost tracker | **skipped** | ran |
| request log | **skipped** | ran |
| PERF line | **skipped** | ran |
| caller sees `CancelledError` | yes | yes |

The new test fails on its parent commit with a `TimeoutError`.

## Test changes

Dropped `test_event_loop_keeps_running_during_the_ledger_write`. It
detected a strict subset of what the lock test already detects:

| scenario | lock test | loop test |
|---|---|---|
| correct: outside lock + `to_thread` | PASS | PASS |
| regress: INSIDE lock + `to_thread` | FAIL | **PASS** |
| regress: outside lock + sync write | FAIL | FAIL |
| pre-fix: INSIDE lock + sync write | FAIL | FAIL |

Added a concurrency test in its place, which covers what the offload
actually introduces: before the move every proxy ledger write ran on the
one event-loop thread and was serialised for free, and now N in-flight
requests append from N worker threads.

One thing left open. That same intra-process concurrency reaches
`_maybe_compact`, which rewrites the file in place. On POSIX the
ledger's own `flock` serialises it. On Windows `_HAS_FCNTL` is false and
all locking is skipped, so a single Windows proxy can now interleave
writers where the loop thread used to serialise them. The cross-process
form of that is pre-existing and called out at `savings_ledger.py:38`.
Happy to take the intra-process guard here or in a follow-up.

Two notes on the sweep above, now that the diff is three files. The
blast radius re-run at this head is 136 of 145 files green, and the nine
non-green are identical with and without the change. Two of them
(`test_gemini_function_response_waste.py`,
`test_openai_responses_context_compaction.py`) are not in the seven
listed earlier; I re-ran both against a reverted `server.py` and they
hang the same way on both sides.
2026-07-20 11:01:34 -07:00

653 lines
25 KiB
Python

"""Tests for :class:`headroom.proxy.outcome.RequestOutcome` and the
:meth:`HeadroomProxy._record_request_outcome` funnel.
The point of this file is the *contract* — every behavioural assertion
here is a thing that, prior to the funnel, lived inline at one or more
of the 18 metrics-emit sites identified in
``docs/superpowers/specs/P0-proxy-pipeline-audit.md``. Locking the
contract in tests means future migrations onto the funnel cannot
silently regress the wire shape.
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
from dataclasses import FrozenInstanceError
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
from headroom.proxy.outcome import RequestOutcome
# ── Value-type contract ────────────────────────────────────────────────
def _outcome(**overrides: Any) -> RequestOutcome:
"""Construct a RequestOutcome with sensible defaults; override fields per test."""
defaults: dict[str, Any] = {
"request_id": "req-1",
"provider": "anthropic",
"model": "claude-sonnet-4",
"original_tokens": 1000,
"optimized_tokens": 300,
"output_tokens": 50,
"tokens_saved": 700,
"attempted_input_tokens": 800,
}
defaults.update(overrides)
return RequestOutcome(**defaults)
def test_outcome_is_frozen() -> None:
"""Mutability would let a handler patch the outcome after handing it
to the funnel — bypassing the contract. Must error."""
o = _outcome()
with pytest.raises(FrozenInstanceError):
o.cache_read_tokens = 999 # type: ignore[misc]
def test_cache_hit_is_derived_not_stored() -> None:
"""Pre-refactor, 9 of 18 ``RequestLog`` sites hardcoded ``cache_hit=False``
even when ``cache_read_tokens > 0``. Deriving from the actual value
makes "forgot to compute it" structurally impossible."""
assert _outcome(cache_read_tokens=0).cache_hit is False
assert _outcome(cache_read_tokens=1).cache_hit is True
assert _outcome(cache_read_tokens=500).cache_hit is True
def test_cache_hit_pct_handles_zero_denominator() -> None:
"""No reads + no writes is a no-cache request, not a 0%-hit cache request.
Returning 0 here is correct as long as dashboards distinguish via the
absolute ``cache_read_tokens`` / ``cache_write_tokens`` values."""
assert _outcome(cache_read_tokens=0, cache_write_tokens=0).cache_hit_pct == 0
def test_cache_hit_pct_rounds_to_int() -> None:
"""PERF log line consumed by ``headroom perf`` parses an integer here;
keep the type contract tight."""
o = _outcome(cache_read_tokens=2, cache_write_tokens=1) # 66.66%
assert o.cache_hit_pct == 67
assert isinstance(o.cache_hit_pct, int)
def test_savings_pct_handles_zero_original() -> None:
"""A request with 0 original tokens — e.g. an empty body — should not
raise ZeroDivisionError. Sites pre-refactor handled this inconsistently."""
assert _outcome(original_tokens=0).savings_pct == 0.0
def test_savings_pct_basic() -> None:
assert _outcome(original_tokens=1000, tokens_saved=300).savings_pct == 30.0
def test_provider_specific_fields_default_to_zero() -> None:
"""Anthropic's 5m/1h cache TTL splits don't exist on OpenAI / Gemini.
The dataclass defaults them to 0 so non-Anthropic handlers don't
have to know about them."""
o = _outcome(provider="openai", cache_read_tokens=100, cache_write_tokens=200)
assert o.cache_write_5m_tokens == 0
assert o.cache_write_1h_tokens == 0
# And OpenAI's "inferred" flag defaults False — only the OpenAI
# handler sets it True after running _infer_openai_cache_write_tokens.
assert o.cache_inferred is False
def test_optional_fields_default_to_neutral_values() -> None:
"""Handlers that don't have a field (e.g. Bedrock with no waste_signals)
must not have to pass anything — defaults handle it."""
o = _outcome()
assert o.ttfb_ms == 0.0
assert o.pipeline_timing is None
assert o.waste_signals is None
assert o.transforms_applied == ()
assert o.turn_id is None
assert o.request_messages is None
assert o.tags == {}
assert o.client is None # unidentified harness
def test_client_field_round_trips() -> None:
"""The ``client`` field is the proof point that the refactor pays
out across harnesses — one field-add gives every dashboard a
per-harness dimension for free.
"""
o = _outcome(client="codex")
assert o.client == "codex"
def test_stream_outcome_derives_gemini_contents_metadata() -> None:
outcome = RequestOutcome.from_stream(
body={
"systemInstruction": {"parts": [{"text": "sys"}]},
"contents": [{"role": "user", "parts": [{"text": "hello"}]}],
},
provider="vertex:google",
model="gemini-2.0-flash",
request_id="req-gemini-stream",
original_tokens=12,
optimized_tokens=10,
output_tokens=3,
tokens_saved=2,
transforms_applied=["compress"],
total_latency_ms=25.0,
overhead_ms=4.0,
tags={"route": "vertex"},
client="codex",
log_full_messages=True,
)
assert outcome.provider == "vertex:google"
assert outcome.num_messages == 1
assert outcome.request_messages == [{"role": "user", "parts": [{"text": "hello"}]}]
assert outcome.turn_id is not None
# ── classify_client — the harness ID source ─────────────────────────
def test_classify_client_recognises_known_harness_user_agents() -> None:
from headroom.proxy.auth_mode import classify_client
cases = [
({"User-Agent": "codex-cli/0.30.0 (osx)"}, "codex"),
({"User-Agent": "claude-code/1.4.2"}, "claude-code"),
({"User-Agent": "claude-cli/2.0"}, "claude-code"), # aliased
({"User-Agent": "cursor/0.42.1 (electron)"}, "cursor"),
({"User-Agent": "aider/0.50.0"}, "aider"),
({"User-Agent": "zed/0.143.0"}, "zed"),
({"User-Agent": "opencode/1.0"}, "opencode"),
({"User-Agent": "github-copilot/x.y.z"}, "copilot"),
]
for headers, expected in cases:
assert classify_client(headers) == expected, headers
def test_classify_client_x_client_header_wins_over_user_agent() -> None:
from headroom.proxy.auth_mode import classify_client
# X-Client wins even when UA matches a different harness
h = {"User-Agent": "codex-cli/0.30.0", "X-Client": "my-custom-harness"}
assert classify_client(h) == "my-custom-harness"
def test_classify_client_returns_none_for_unknown_traffic() -> None:
"""``None`` is the loud "unidentified" signal — downstream consumers
can group these as "unknown" rather than silently bucketing into
a default that would mislead dashboards."""
from headroom.proxy.auth_mode import classify_client
assert classify_client({"User-Agent": "Mozilla/5.0"}) is None
assert classify_client({}) is None
assert classify_client({"User-Agent": ""}) is None
# ── Funnel contract (_record_request_outcome) ──────────────────────────
class _CollectingLogger:
"""Minimal stand-in for ``RequestLogger``."""
def __init__(self) -> None:
self.logs: list[Any] = []
def log(self, entry: Any) -> None:
self.logs.append(entry)
class _FunnelHarness:
"""Pulls just enough of HeadroomProxy onto an object to exercise
``_record_request_outcome`` without instantiating the full proxy.
The harness assigns the real method to ``self`` via descriptor
binding so the implementation is exactly the production one — no
forking, no mock-the-thing-you're-testing.
"""
def __init__(self, *, with_cost_tracker: bool = True, with_logger: bool = True) -> None:
from headroom.proxy.server import HeadroomProxy
self.metrics = MagicMock()
self.metrics.record_request = AsyncMock()
self.cost_tracker = MagicMock() if with_cost_tracker else None
self.logger = _CollectingLogger() if with_logger else None
# Bind the real method to this harness.
self._record_request_outcome = HeadroomProxy._record_request_outcome.__get__(
self, type(self)
)
@pytest.mark.asyncio
async def test_funnel_calls_metrics_with_full_kwargs() -> None:
"""The funnel must pass EVERY field that
``PrometheusMetrics.record_request`` knows about, not the
pre-refactor "pass-what-was-convenient" subset. Otherwise a handler
that forgets to populate a field silently degrades dashboard data."""
h = _FunnelHarness()
o = _outcome(
provider="openai",
model="gpt-4",
optimized_tokens=300,
output_tokens=50,
tokens_saved=700,
attempted_input_tokens=800,
cache_read_tokens=200,
cache_write_tokens=100,
cache_write_5m_tokens=50,
cache_write_1h_tokens=50,
uncached_input_tokens=0,
total_latency_ms=1234.5,
overhead_ms=12.3,
ttfb_ms=200.0,
pipeline_timing={"phase": 1.0},
waste_signals={"skipped": 3},
)
await h._record_request_outcome(o)
h.metrics.record_request.assert_awaited_once()
kwargs = h.metrics.record_request.await_args.kwargs
assert kwargs["provider"] == "openai"
assert kwargs["model"] == "gpt-4"
assert kwargs["input_tokens"] == 300 # optimized → input
assert kwargs["output_tokens"] == 50
assert kwargs["tokens_saved"] == 700
assert kwargs["latency_ms"] == 1234.5
assert kwargs["cached"] is True # derived from cache_read > 0
assert kwargs["overhead_ms"] == 12.3
assert kwargs["ttfb_ms"] == 200.0
assert kwargs["pipeline_timing"] == {"phase": 1.0}
assert kwargs["waste_signals"] == {"skipped": 3}
assert kwargs["cache_read_tokens"] == 200
assert kwargs["cache_write_tokens"] == 100
assert kwargs["cache_write_5m_tokens"] == 50
assert kwargs["cache_write_1h_tokens"] == 50
assert kwargs["uncached_input_tokens"] == 0
assert kwargs["attempted_input_tokens"] == 800
@pytest.mark.asyncio
async def test_funnel_passes_canonical_record_tokens_shape() -> None:
"""``cost_tracker.record_tokens`` takes ``(model, tokens_saved,
optimized_tokens)`` positionally and the cache args as kwargs. The
funnel preserves this — moving anything to positional would break
sites that pass kwargs explicitly."""
h = _FunnelHarness()
o = _outcome(
model="claude-sonnet-4",
optimized_tokens=300,
tokens_saved=700,
cache_read_tokens=200,
cache_write_tokens=100,
cache_write_5m_tokens=80,
cache_write_1h_tokens=20,
uncached_input_tokens=0,
)
await h._record_request_outcome(o)
h.cost_tracker.record_tokens.assert_called_once()
args, kwargs = h.cost_tracker.record_tokens.call_args
assert args == ("claude-sonnet-4", 700, 300)
assert kwargs == {
"cache_read_tokens": 200,
"cache_write_tokens": 100,
"cache_write_5m_tokens": 80,
"cache_write_1h_tokens": 20,
"uncached_tokens": 0,
"output_tokens": 50,
}
@pytest.mark.asyncio
async def test_funnel_skips_cost_tracker_when_absent() -> None:
"""When the proxy was started with ``--no-cost``, ``cost_tracker``
is None and the funnel must skip step 2 silently."""
h = _FunnelHarness(with_cost_tracker=False)
await h._record_request_outcome(_outcome())
# No crash, metrics still recorded.
h.metrics.record_request.assert_awaited_once()
@pytest.mark.asyncio
async def test_funnel_logs_request_with_derived_cache_hit() -> None:
"""The RequestLog row needs cache_hit derived from cache_read>0, not
the hardcoded False that 9 of 18 pre-refactor sites used."""
h = _FunnelHarness()
await h._record_request_outcome(_outcome(cache_read_tokens=200, cache_write_tokens=100))
assert len(h.logger.logs) == 1
log_entry = h.logger.logs[0]
assert log_entry.cache_hit is True
@pytest.mark.asyncio
async def test_funnel_skips_request_log_when_logger_absent() -> None:
"""Same pattern as cost_tracker — optional surface."""
h = _FunnelHarness(with_logger=False)
await h._record_request_outcome(_outcome())
h.metrics.record_request.assert_awaited_once() # still happens
@pytest.mark.asyncio
async def test_funnel_tail_survives_cancellation_inside_record_request() -> None:
"""A client disconnect must not tear per-request bookkeeping in half.
Four call sites are ``finally:`` blocks inside streaming async generators
(``streaming.py:1611``, ``:1859``, ``:2069``, ``openai.py:8614``), and
``record_request`` suspends partway through — it awaits the savings-ledger
append in a worker thread after the Prometheus counters have already been
committed. A cancellation landing on that await used to skip every effect
below it, leaving the request counted in Prometheus but absent from the cost
tracker, the request log, and the PERF line ``headroom perf`` reads.
Without the ``asyncio.shield`` in ``_record_request_outcome`` the release
below never resumes the funnel and this test times out on ``logged``.
"""
h = _FunnelHarness()
logged = asyncio.Event()
collect = h.logger.log
def log_and_signal(entry: Any) -> None:
collect(entry)
logged.set()
h.logger.log = log_and_signal # type: ignore[method-assign]
entered = asyncio.Event()
release = asyncio.Event()
async def suspending_record_request(**kwargs: Any) -> None:
# Stands in for the `await asyncio.to_thread(...)` ledger append: the
# counters are in, and the funnel is now parked on an await.
entered.set()
await release.wait()
h.metrics.record_request = suspending_record_request
task = asyncio.create_task(h._record_request_outcome(_outcome()))
await asyncio.wait_for(entered.wait(), timeout=5)
task.cancel()
# The shield deliberately does not swallow the cancellation — the caller
# still sees CancelledError, so generator teardown propagates unchanged.
with contextlib.suppress(asyncio.CancelledError):
await task
release.set()
await asyncio.wait_for(logged.wait(), timeout=5)
assert h.cost_tracker.record_tokens.called, "cost tracker was skipped by the cancellation"
assert len(h.logger.logs) == 1, "request log was skipped by the cancellation"
@pytest.mark.asyncio
async def test_funnel_emits_perf_log_with_canonical_shape(
caplog: pytest.LogCaptureFixture,
) -> None:
"""``headroom perf`` parses this exact ``key=value`` format. Changing
it breaks the analyzer. The contract: model, msgs, tok_before,
tok_after, tok_saved, cache_read, cache_write, cache_hit_pct,
opt_ms, transforms — in that order, space-separated."""
h = _FunnelHarness()
# Direct handler attach: caplog otherwise drops propagation-disabled
# records (the proxy disables ``headroom.*`` propagation once started).
target = logging.getLogger("headroom.proxy")
captured: list[logging.LogRecord] = []
class _H(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
captured.append(record)
handler = _H(level=logging.INFO)
target.addHandler(handler)
prior_level = target.level
target.setLevel(logging.INFO)
try:
await h._record_request_outcome(
_outcome(
request_id="req-perf",
model="gpt-4",
original_tokens=1000,
optimized_tokens=300,
tokens_saved=700,
cache_read_tokens=200,
cache_write_tokens=100,
num_messages=5,
overhead_ms=12.0,
transforms_applied=("smart_crusher", "content_router"),
)
)
finally:
target.removeHandler(handler)
target.setLevel(prior_level)
perf_lines = [r.getMessage() for r in captured if " PERF " in r.getMessage()]
assert len(perf_lines) == 1
line = perf_lines[0]
assert "[req-perf] PERF " in line
assert "model=gpt-4" in line
assert "msgs=5" in line
assert "tok_before=1000" in line
assert "tok_after=300" in line
assert "tok_saved=700" in line
assert "cache_read=200" in line
assert "cache_write=100" in line
assert "cache_hit_pct=67" in line # 200/(200+100) * 100 = 67
assert "opt_ms=12" in line
# ── Funnel: per-client analytics surface ─────────────────────────────
@pytest.mark.asyncio
async def test_funnel_appends_client_to_perf_log_when_set() -> None:
"""``headroom perf --client X`` filtering relies on the ``client=X``
token at the end of the PERF line. Absent client means no token —
the PERF line stays clean for unidentified traffic."""
h = _FunnelHarness()
target = logging.getLogger("headroom.proxy")
captured: list[logging.LogRecord] = []
class _H(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
captured.append(record)
handler = _H(level=logging.INFO)
target.addHandler(handler)
prior_level = target.level
target.setLevel(logging.INFO)
try:
await h._record_request_outcome(_outcome(client="codex"))
finally:
target.removeHandler(handler)
target.setLevel(prior_level)
perf_lines = [r.getMessage() for r in captured if " PERF " in r.getMessage()]
assert len(perf_lines) == 1
assert "client=codex" in perf_lines[0]
@pytest.mark.asyncio
async def test_funnel_omits_client_from_perf_log_when_unidentified() -> None:
"""When ``client`` is None the PERF line must NOT include a
bogus ``client=`` token — that would mislead the parser into
bucketing unidentified traffic as the empty string."""
h = _FunnelHarness()
target = logging.getLogger("headroom.proxy")
captured: list[logging.LogRecord] = []
class _H(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
captured.append(record)
handler = _H(level=logging.INFO)
target.addHandler(handler)
prior_level = target.level
target.setLevel(logging.INFO)
try:
await h._record_request_outcome(_outcome(client=None))
finally:
target.removeHandler(handler)
target.setLevel(prior_level)
perf_lines = [r.getMessage() for r in captured if " PERF " in r.getMessage()]
assert len(perf_lines) == 1
assert "client=" not in perf_lines[0]
@pytest.mark.asyncio
async def test_funnel_stamps_client_into_request_log_tags() -> None:
"""Dashboards already filter on RequestLog.tags. Copying ``client``
into tags gives per-harness slicing for free with no new column."""
h = _FunnelHarness()
await h._record_request_outcome(_outcome(client="aider"))
assert len(h.logger.logs) == 1
assert h.logger.logs[0].tags.get("client") == "aider"
# ── from_stream classmethod (streaming-finalizer construction shape) ──
#
# Three streaming finalizers (``_finalize_stream_response``,
# ``_stream_response_bedrock``, ``_stream_openai_via_backend``) each used
# to construct ``RequestOutcome(...)`` inline with the same body- and
# config-derived fields — ``attempted_input_tokens``, ``num_messages``,
# ``request_messages``, ``turn_id``, tuple-conversion of
# ``transforms_applied``, ``tags`` normalization. One site (Bedrock)
# computed ``turn_id``; the other two silently dropped it — a real bug
# the helper fixes by computing it uniformly. ``from_stream`` is the
# canonical construction point so the three finalizers cannot drift
# apart on derivation logic again.
def _stream_kwargs(**overrides: Any) -> dict[str, Any]:
"""Minimal kwargs for ``RequestOutcome.from_stream``; override per test."""
base: dict[str, Any] = {
"body": {"messages": [{"role": "user", "content": "hi"}]},
"provider": "anthropic",
"model": "claude-sonnet-4",
"request_id": "req-1",
"original_tokens": 1000,
"optimized_tokens": 300,
"output_tokens": 50,
"tokens_saved": 700,
"transforms_applied": ["smart_crusher"],
"total_latency_ms": 1234.5,
"overhead_ms": 12.3,
"tags": {"a": "b"},
"client": "codex",
"log_full_messages": False,
}
base.update(overrides)
return base
def test_from_stream_returns_request_outcome() -> None:
o = RequestOutcome.from_stream(**_stream_kwargs())
assert isinstance(o, RequestOutcome)
def test_from_stream_derives_attempted_input_tokens_from_optimized_plus_saved() -> None:
"""One of the six derivations the three finalizers each computed
inline. Centralising it makes the dashboard's active-savings
denominator structurally consistent across providers (#454/#455)."""
o = RequestOutcome.from_stream(**_stream_kwargs(optimized_tokens=300, tokens_saved=700))
assert o.attempted_input_tokens == 1000
def test_from_stream_counts_messages_from_body() -> None:
"""``num_messages`` powers PERF ``msgs=N``. Computing it from the
body in one place prevents the historical drift where some sites
used ``original_messages`` and others used ``body["messages"]``."""
body = {"messages": [{"role": "user", "content": "1"}, {"role": "user", "content": "2"}]}
assert RequestOutcome.from_stream(**_stream_kwargs(body=body)).num_messages == 2
def test_from_stream_handles_missing_messages_key() -> None:
"""Empty body — e.g. a probe request — must yield num_messages=0,
not raise KeyError. All three pre-refactor sites used
``len(body.get("messages", []))`` so the contract is already
"default to 0"."""
assert RequestOutcome.from_stream(**_stream_kwargs(body={})).num_messages == 0
def test_from_stream_always_computes_turn_id() -> None:
"""The bug the helper is fixing: pre-refactor, only the Bedrock
finalizer called ``compute_turn_id``. Sites 1 and 3 silently dropped
it, breaking the dashboard's multi-turn-session grouping for every
Anthropic-SSE and OpenAI-via-backend request. The helper computes
it uniformly."""
body = {
"messages": [{"role": "user", "content": "hi"}],
"system": "you are helpful",
}
o = RequestOutcome.from_stream(**_stream_kwargs(body=body, model="claude-sonnet-4"))
assert o.turn_id is not None
assert isinstance(o.turn_id, str)
# Stable: same body+model produces the same turn_id.
o2 = RequestOutcome.from_stream(**_stream_kwargs(body=body, model="claude-sonnet-4"))
assert o.turn_id == o2.turn_id
def test_from_stream_converts_transforms_to_tuple() -> None:
"""``transforms_applied`` is typed as ``tuple[str, ...]`` on the
dataclass (frozen → must be hashable/immutable). Callers pass lists.
The helper does the conversion so no caller has to remember."""
o = RequestOutcome.from_stream(**_stream_kwargs(transforms_applied=["a", "b"]))
assert o.transforms_applied == ("a", "b")
assert isinstance(o.transforms_applied, tuple)
def test_from_stream_normalises_none_tags_to_empty_dict() -> None:
"""``tags=None`` is the common case (no routing tags); the dataclass
contract is ``dict[str, str]``. Pre-refactor each site wrote
``tags or {}``; the helper does it once."""
assert RequestOutcome.from_stream(**_stream_kwargs(tags=None)).tags == {}
def test_from_stream_omits_request_messages_when_log_full_messages_disabled() -> None:
"""``log_full_messages=False`` is the default; message bodies are
sensitive (tool outputs, secrets) and must not land in the request
log unless explicitly enabled."""
body = {"messages": [{"role": "user", "content": "secret"}]}
o = RequestOutcome.from_stream(**_stream_kwargs(body=body, log_full_messages=False))
assert o.request_messages is None
def test_from_stream_includes_request_messages_when_log_full_messages_enabled() -> None:
"""Same path, opt-in for full-message logging — used by
/transformations/feed when the operator enables it."""
body = {"messages": [{"role": "user", "content": "hi"}]}
o = RequestOutcome.from_stream(**_stream_kwargs(body=body, log_full_messages=True))
assert o.request_messages == body["messages"]
def test_from_stream_threads_provider_specific_cache_fields() -> None:
"""Anthropic populates all five cache fields; OpenAI-via-backend
sets ``cache_inferred=True``; Gemini populates read only. The
helper must pass each through without forcing every caller to
pass every field."""
o = RequestOutcome.from_stream(
**_stream_kwargs(),
cache_read_tokens=100,
cache_write_tokens=200,
cache_write_5m_tokens=150,
cache_write_1h_tokens=50,
uncached_input_tokens=10,
)
assert o.cache_read_tokens == 100
assert o.cache_write_tokens == 200
assert o.cache_write_5m_tokens == 150
assert o.cache_write_1h_tokens == 50
assert o.uncached_input_tokens == 10
assert o.cache_inferred is False # default — only set True by OpenAI sites
def test_from_stream_threads_waste_signals_for_openai_via_backend_site() -> None:
"""Only the OpenAI-via-backend finalizer populates ``waste_signals``;
the helper threads it through as an optional kwarg."""
o = RequestOutcome.from_stream(
**_stream_kwargs(),
waste_signals={"skipped_units": 3, "applied_units": 7},
)
assert o.waste_signals == {"skipped_units": 3, "applied_units": 7}