|
| 1 | +"""A buffered-CCR turn is a billed call and its usage must reach accounting. |
| 2 | +
|
| 3 | +When a ``stream:true`` request carries the ``headroom_retrieve`` tool, the |
| 4 | +Anthropic handler rewrites it to ``stream:false`` upstream so it can resolve |
| 5 | +retrievals server-side, then re-synthesizes SSE for the client. That buffered |
| 6 | +response carries the same ``usage`` block any non-stream reply does, so the |
| 7 | +provider's cache-read / cache-write / output counts must land on the outcome. |
| 8 | +
|
| 9 | +If they don't, every cached token on the dominant Claude Code path is invisible: |
| 10 | +``metrics.cache_by_provider`` only records a provider row when cache read or |
| 11 | +write is non-zero, so the whole prefix-cache card (hit rate, savings, TTL mix) |
| 12 | +is computed from whatever small fraction of traffic took another path, while |
| 13 | +compression savings on the buffered turns still count in full. |
| 14 | +""" |
| 15 | + |
| 16 | +from __future__ import annotations |
| 17 | + |
| 18 | +from typing import Any |
| 19 | + |
| 20 | +import httpx |
| 21 | +import pytest |
| 22 | +import respx |
| 23 | + |
| 24 | +pytest.importorskip("fastapi") |
| 25 | + |
| 26 | +from fastapi.testclient import TestClient # noqa: E402 |
| 27 | + |
| 28 | +from headroom.proxy.loopback_guard import require_loopback # noqa: E402 |
| 29 | +from headroom.proxy.server import ProxyConfig, create_app # noqa: E402 |
| 30 | + |
| 31 | +# Provider-reported usage for one warm turn: most of the prompt served from |
| 32 | +# cache, a slice newly written, a little uncached, and a real completion. |
| 33 | +_USAGE = { |
| 34 | + "input_tokens": 300, |
| 35 | + "output_tokens": 120, |
| 36 | + "cache_read_input_tokens": 46893, |
| 37 | + "cache_creation_input_tokens": 8197, |
| 38 | +} |
| 39 | + |
| 40 | +_RESPONSE = { |
| 41 | + "id": "msg_buffered", |
| 42 | + "type": "message", |
| 43 | + "role": "assistant", |
| 44 | + "model": "claude-opus-5", |
| 45 | + "content": [{"type": "text", "text": "done"}], |
| 46 | + "stop_reason": "end_turn", |
| 47 | + "usage": _USAGE, |
| 48 | +} |
| 49 | + |
| 50 | +_RETRIEVE_TOOL = { |
| 51 | + "name": "headroom_retrieve", |
| 52 | + "description": "Retrieve compressed content", |
| 53 | + "input_schema": {"type": "object", "properties": {"ref": {"type": "string"}}}, |
| 54 | +} |
| 55 | + |
| 56 | + |
| 57 | +def _app_and_outcomes(monkeypatch, **overrides): |
| 58 | + """App with a spy on the outcome record — where billed counts land.""" |
| 59 | + kwargs: dict[str, Any] = { |
| 60 | + "optimize": False, |
| 61 | + "cache_enabled": False, |
| 62 | + "rate_limit_enabled": False, |
| 63 | + "cost_tracking_enabled": False, |
| 64 | + "log_requests": False, |
| 65 | + } |
| 66 | + kwargs.update(overrides) |
| 67 | + app = create_app(ProxyConfig(**kwargs)) |
| 68 | + app.dependency_overrides[require_loopback] = lambda: None |
| 69 | + outcomes: list[Any] = [] |
| 70 | + |
| 71 | + async def _spy(_self, outcome, *a, **kw): |
| 72 | + outcomes.append(outcome) |
| 73 | + |
| 74 | + monkeypatch.setattr(type(app.state.proxy), "_record_request_outcome", _spy, raising=True) |
| 75 | + return app, outcomes |
| 76 | + |
| 77 | + |
| 78 | +def _post(app, *, stream: bool, tools: list[dict] | None): |
| 79 | + body: dict[str, Any] = { |
| 80 | + "model": "claude-opus-5", |
| 81 | + "max_tokens": 256, |
| 82 | + "messages": [{"role": "user", "content": "hi"}], |
| 83 | + } |
| 84 | + if stream: |
| 85 | + body["stream"] = True |
| 86 | + if tools is not None: |
| 87 | + body["tools"] = tools |
| 88 | + with TestClient(app) as client: |
| 89 | + return client.post("/v1/messages", json=body, headers={"x-api-key": "sk-ant-test"}) |
| 90 | + |
| 91 | + |
| 92 | +@respx.mock |
| 93 | +def test_non_stream_turn_records_provider_usage(monkeypatch) -> None: |
| 94 | + """Control: the plain non-stream path already books the usage block.""" |
| 95 | + app, outcomes = _app_and_outcomes(monkeypatch) |
| 96 | + respx.post("https://api.anthropic.com/v1/messages").mock( |
| 97 | + return_value=httpx.Response(200, json=_RESPONSE) |
| 98 | + ) |
| 99 | + |
| 100 | + r = _post(app, stream=False, tools=None) |
| 101 | + |
| 102 | + assert r.status_code == 200 |
| 103 | + assert outcomes, "an outcome must be recorded" |
| 104 | + o = outcomes[-1] |
| 105 | + assert o.cache_read_tokens == 46893 |
| 106 | + assert o.cache_write_tokens == 8197 |
| 107 | + assert o.output_tokens == 120 |
| 108 | + |
| 109 | + |
| 110 | +@respx.mock |
| 111 | +def test_buffered_ccr_turn_records_provider_usage(monkeypatch) -> None: |
| 112 | + """A stream:true + headroom_retrieve turn must book the same usage block.""" |
| 113 | + app, outcomes = _app_and_outcomes(monkeypatch) |
| 114 | + route = respx.post("https://api.anthropic.com/v1/messages").mock( |
| 115 | + return_value=httpx.Response(200, json=_RESPONSE) |
| 116 | + ) |
| 117 | + |
| 118 | + r = _post(app, stream=True, tools=[_RETRIEVE_TOOL]) |
| 119 | + |
| 120 | + assert r.status_code == 200 |
| 121 | + # The handler must have buffered it: upstream saw stream:false. |
| 122 | + sent = route.calls.last.request |
| 123 | + assert b'"stream": false' in sent.content or b'"stream":false' in sent.content |
| 124 | + |
| 125 | + assert outcomes, "an outcome must be recorded" |
| 126 | + o = outcomes[-1] |
| 127 | + assert o.cache_read_tokens == 46893, ( |
| 128 | + f"buffered-CCR turn dropped the provider cache-read count: {o.cache_read_tokens}" |
| 129 | + ) |
| 130 | + assert o.cache_write_tokens == 8197, ( |
| 131 | + f"buffered-CCR turn dropped the provider cache-write count: {o.cache_write_tokens}" |
| 132 | + ) |
| 133 | + assert o.output_tokens == 120, ( |
| 134 | + f"buffered-CCR turn dropped the provider output count: {o.output_tokens}" |
| 135 | + ) |
| 136 | + |
| 137 | + |
| 138 | +@respx.mock |
| 139 | +def test_buffered_ccr_records_usage_with_compression_on(monkeypatch) -> None: |
| 140 | + """Same turn with the live token-mode pipeline running, as in production.""" |
| 141 | + app, outcomes = _app_and_outcomes(monkeypatch, optimize=True, mode="token") |
| 142 | + respx.post("https://api.anthropic.com/v1/messages").mock( |
| 143 | + return_value=httpx.Response(200, json=_RESPONSE) |
| 144 | + ) |
| 145 | + |
| 146 | + r = _post(app, stream=True, tools=[_RETRIEVE_TOOL]) |
| 147 | + |
| 148 | + assert r.status_code == 200 |
| 149 | + assert outcomes, "an outcome must be recorded" |
| 150 | + o = outcomes[-1] |
| 151 | + assert (o.cache_read_tokens, o.cache_write_tokens, o.output_tokens) == (46893, 8197, 120), ( |
| 152 | + f"buffered-CCR turn dropped usage under compression: " |
| 153 | + f"cr={o.cache_read_tokens} cw={o.cache_write_tokens} out={o.output_tokens}" |
| 154 | + ) |
0 commit comments