mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
c7af8307d1
Resolve conflict in tests/test_proxy_streaming_request_logger.py by keeping both new tests: - test_finalize_openai_responses_stream_uses_provider_usage_for_dashboard (this branch — OpenAI Responses stream dashboard usage) - test_finalize_stream_response_recovers_usage_from_truncated_buffer (main — SSE buffer recovery on truncated streams) streaming.py auto-merged cleanly (additive only).
300 lines
10 KiB
Python
300 lines
10 KiB
Python
"""Tests that the Anthropic streaming finalizer logs requests for the feed.
|
|
|
|
Without this, the streaming Anthropic path (which is what Claude Code uses)
|
|
silently bypassed the request logger, leaving `/stats.recent_requests` and
|
|
`/transformations/feed` permanently empty even when `--log-messages` was set.
|
|
The non-streaming Anthropic path and the Bedrock streaming path were the
|
|
only ones that called `self.logger.log(...)`.
|
|
"""
|
|
|
|
import json
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from headroom.proxy.request_logger import RequestLogger
|
|
from headroom.proxy.server import HeadroomProxy
|
|
|
|
|
|
def _build_proxy_with_real_logger(*, log_full_messages: bool) -> HeadroomProxy:
|
|
"""Build a HeadroomProxy with mocks for everything except the request logger,
|
|
so we can assert what actually gets recorded."""
|
|
proxy = object.__new__(HeadroomProxy)
|
|
proxy.http_client = MagicMock(spec=httpx.AsyncClient)
|
|
proxy.metrics = MagicMock()
|
|
proxy.metrics.record_request = AsyncMock(return_value=None)
|
|
proxy.cost_tracker = MagicMock()
|
|
proxy.cost_tracker.record_tokens.return_value = None
|
|
proxy.memory_manager = None
|
|
proxy.memory_handler = None
|
|
proxy._config = MagicMock()
|
|
proxy._config.log_full_messages = log_full_messages
|
|
proxy._config.ccr_inject_tool = False
|
|
proxy.config = proxy._config
|
|
proxy.logger = RequestLogger(log_file=None, log_full_messages=log_full_messages)
|
|
return proxy
|
|
|
|
|
|
def _stream_state(output_tokens: int = 42) -> dict:
|
|
return {
|
|
"output_tokens": output_tokens,
|
|
"total_bytes": 200,
|
|
"ttfb_ms": 35.0,
|
|
"input_tokens": 1000,
|
|
"cache_read_input_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_creation_ephemeral_5m_input_tokens": 0,
|
|
"cache_creation_ephemeral_1h_input_tokens": 0,
|
|
"sse_buffer": "",
|
|
}
|
|
|
|
|
|
def test_parse_openai_responses_completed_usage_from_sse_buffer():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
completed = {
|
|
"type": "response.completed",
|
|
"response": {
|
|
"id": "resp_1",
|
|
"usage": {
|
|
"input_tokens": 844_000,
|
|
"input_tokens_details": {"cached_tokens": 657_400},
|
|
"output_tokens": 6_635,
|
|
},
|
|
},
|
|
}
|
|
state = {
|
|
"sse_buffer": bytearray(
|
|
f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode()
|
|
)
|
|
}
|
|
|
|
usage = proxy._parse_sse_usage_from_buffer(state, "openai")
|
|
|
|
assert usage == {
|
|
"input_tokens": 844_000,
|
|
"output_tokens": 6_635,
|
|
"cache_read_input_tokens": 657_400,
|
|
}
|
|
assert state["sse_buffer"] == bytearray()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_logs_request_for_feed():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
|
|
await proxy._finalize_stream_response(
|
|
body={"messages": [{"role": "user", "content": "hi"}]},
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-1",
|
|
original_tokens=1000,
|
|
optimized_tokens=600,
|
|
tokens_saved=400,
|
|
transforms_applied=["smart_crusher"],
|
|
optimization_latency=12.0,
|
|
stream_state=_stream_state(),
|
|
start_time=0.0,
|
|
tags={"stack": "wrap_claude"},
|
|
)
|
|
|
|
entries = proxy.logger.get_recent(10)
|
|
assert len(entries) == 1, "streaming finalizer must log exactly one entry per request"
|
|
entry = entries[0]
|
|
assert entry["request_id"] == "req-stream-1"
|
|
assert entry["provider"] == "anthropic"
|
|
assert entry["model"] == "claude-sonnet-4-6"
|
|
assert entry["input_tokens_original"] == 1000
|
|
assert entry["input_tokens_optimized"] == 600
|
|
assert entry["tokens_saved"] == 400
|
|
assert entry["savings_percent"] == pytest.approx(40.0)
|
|
assert entry["transforms_applied"] == ["smart_crusher"]
|
|
assert entry["tags"] == {"stack": "wrap_claude"}
|
|
assert entry["cache_hit"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_includes_messages_when_log_full_messages_enabled():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=True)
|
|
body = {"messages": [{"role": "user", "content": "hello"}]}
|
|
|
|
await proxy._finalize_stream_response(
|
|
body=body,
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-2",
|
|
original_tokens=10,
|
|
optimized_tokens=8,
|
|
tokens_saved=2,
|
|
transforms_applied=[],
|
|
optimization_latency=1.0,
|
|
stream_state=_stream_state(output_tokens=5),
|
|
start_time=0.0,
|
|
)
|
|
|
|
entries = proxy.logger.get_recent_with_messages(10)
|
|
assert len(entries) == 1
|
|
assert entries[0]["request_messages"] == body["messages"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_omits_messages_when_log_full_messages_disabled():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
|
|
await proxy._finalize_stream_response(
|
|
body={"messages": [{"role": "user", "content": "hello"}]},
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-3",
|
|
original_tokens=10,
|
|
optimized_tokens=8,
|
|
tokens_saved=2,
|
|
transforms_applied=[],
|
|
optimization_latency=1.0,
|
|
stream_state=_stream_state(output_tokens=5),
|
|
start_time=0.0,
|
|
)
|
|
|
|
entries = proxy.logger.get_recent_with_messages(10)
|
|
assert len(entries) == 1
|
|
assert entries[0]["request_messages"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_handles_zero_original_tokens():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
|
|
await proxy._finalize_stream_response(
|
|
body={"messages": []},
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-4",
|
|
original_tokens=0,
|
|
optimized_tokens=0,
|
|
tokens_saved=0,
|
|
transforms_applied=[],
|
|
optimization_latency=0.0,
|
|
stream_state=_stream_state(output_tokens=0),
|
|
start_time=0.0,
|
|
)
|
|
|
|
entries = proxy.logger.get_recent(10)
|
|
assert len(entries) == 1
|
|
assert entries[0]["savings_percent"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_openai_responses_stream_uses_provider_usage_for_dashboard():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
state = _stream_state(output_tokens=6_635)
|
|
state["input_tokens"] = 844_000
|
|
state["cache_read_input_tokens"] = 657_400
|
|
|
|
await proxy._finalize_stream_response(
|
|
body={"model": "gpt-5.5", "input": [{"type": "message", "role": "user"}]},
|
|
provider="openai",
|
|
model="gpt-5.5",
|
|
request_id="req-openai-responses-stream",
|
|
original_tokens=0,
|
|
optimized_tokens=0,
|
|
tokens_saved=663_000,
|
|
transforms_applied=["openai_responses_live_zone"],
|
|
optimization_latency=26.0,
|
|
stream_state=state,
|
|
start_time=0.0,
|
|
)
|
|
|
|
entries = proxy.logger.get_recent(10)
|
|
assert len(entries) == 1
|
|
entry = entries[0]
|
|
assert entry["input_tokens_optimized"] == 844_000
|
|
assert entry["input_tokens_original"] == 1_507_000
|
|
assert entry["tokens_saved"] == 663_000
|
|
assert entry["savings_percent"] == pytest.approx(663_000 / 1_507_000 * 100)
|
|
assert entry["output_tokens"] == 6_635
|
|
|
|
proxy.metrics.record_request.assert_awaited_once()
|
|
metrics_kwargs = proxy.metrics.record_request.await_args.kwargs
|
|
assert metrics_kwargs["input_tokens"] == 844_000
|
|
assert metrics_kwargs["output_tokens"] == 6_635
|
|
assert metrics_kwargs["tokens_saved"] == 663_000
|
|
assert metrics_kwargs["cache_read_tokens"] == 657_400
|
|
assert metrics_kwargs["uncached_input_tokens"] == 186_600
|
|
|
|
proxy.cost_tracker.record_tokens.assert_called_once()
|
|
cost_args, cost_kwargs = proxy.cost_tracker.record_tokens.call_args
|
|
assert cost_args[:3] == ("gpt-5.5", 663_000, 844_000)
|
|
assert cost_kwargs["cache_read_tokens"] == 657_400
|
|
assert cost_kwargs["uncached_tokens"] == 186_600
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_recovers_usage_from_truncated_buffer() -> None:
|
|
"""When upstream truncates mid-event (no trailing \\n\\n), the per-chunk
|
|
parser leaves the message_start usage event sitting in sse_buffer and
|
|
PERF logs cache_read=cache_write=0 — which then poisons the freeze
|
|
heuristic on the next request. The finalizer must flush the residual
|
|
buffer so the real cache_read / cache_creation tokens still land in
|
|
the log even on aborted streams.
|
|
"""
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
|
|
partial_message_start = (
|
|
b"event: message_start\n"
|
|
b'data: {"type":"message_start","message":{"id":"msg_x",'
|
|
b'"type":"message","role":"assistant","model":"claude-sonnet-4-6",'
|
|
b'"content":[],"stop_reason":null,"usage":{'
|
|
b'"input_tokens":1234,"cache_read_input_tokens":50000,'
|
|
b'"cache_creation_input_tokens":2500,"output_tokens":1}}}'
|
|
)
|
|
|
|
state = {
|
|
"output_tokens": None,
|
|
"total_bytes": len(partial_message_start),
|
|
"ttfb_ms": 35.0,
|
|
"input_tokens": None,
|
|
"cache_read_input_tokens": 0,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_creation_ephemeral_5m_input_tokens": 0,
|
|
"cache_creation_ephemeral_1h_input_tokens": 0,
|
|
"sse_buffer": bytearray(partial_message_start),
|
|
}
|
|
|
|
await proxy._finalize_stream_response(
|
|
body={"messages": [{"role": "user", "content": "hi"}]},
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-truncated",
|
|
original_tokens=2000,
|
|
optimized_tokens=1800,
|
|
tokens_saved=200,
|
|
transforms_applied=[],
|
|
optimization_latency=5.0,
|
|
stream_state=state,
|
|
start_time=0.0,
|
|
)
|
|
|
|
assert state["input_tokens"] == 1234
|
|
assert state["cache_read_input_tokens"] == 50000
|
|
assert state["cache_creation_input_tokens"] == 2500
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_stream_response_no_op_when_logger_disabled():
|
|
proxy = _build_proxy_with_real_logger(log_full_messages=False)
|
|
proxy.logger = None # `--no-log-requests` would put us here
|
|
|
|
# Should not raise.
|
|
await proxy._finalize_stream_response(
|
|
body={"messages": []},
|
|
provider="anthropic",
|
|
model="claude-sonnet-4-6",
|
|
request_id="req-stream-5",
|
|
original_tokens=10,
|
|
optimized_tokens=8,
|
|
tokens_saved=2,
|
|
transforms_applied=[],
|
|
optimization_latency=1.0,
|
|
stream_state=_stream_state(),
|
|
start_time=0.0,
|
|
)
|