chopratejas--headroom/tests/test_proxy_streaming_request_logger.py

"""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,
    )