chore: import upstream snapshot with attribution

tests/test_proxy_streaming_request_logger.py

@@ -0,0 +1,314 @@
+"""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_logs_original_and_compressed_messages():
+    """With log_full_messages enabled, both sides of the compression are
+    recorded: `request_messages` is the pre-compression snapshot the caller
+    threads in via `original_messages`, `compressed_messages` is what was
+    actually sent upstream (i.e. `body["messages"]` after in-place mutation)."""
+    proxy = _build_proxy_with_real_logger(log_full_messages=True)
+    # `body["messages"]` models the post-compression list - the proxy mutates
+    # `body` in place before calling `_finalize_stream_response`, so this is
+    # already what was shipped over the wire.
+    body = {"messages": [{"role": "user", "content": "[compressed]"}]}
+    original = [{"role": "user", "content": "[original, pre-compression]"}]
+
+    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,
+        original_messages=original,
+    )
+
+    entries = proxy.logger.get_recent_with_messages(10)
+    assert len(entries) == 1
+    assert entries[0]["request_messages"] == original
+    assert entries[0]["compressed_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,
+        original_messages=[{"role": "user", "content": "dropped"}],
+    )
+
+    entries = proxy.logger.get_recent_with_messages(10)
+    assert len(entries) == 1
+    # Both sides share the same gate - neither leaks when log_full_messages
+    # is off.
+    assert entries[0]["request_messages"] is None
+    assert entries[0]["compressed_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,
+    )