headroom/tests/test_cache_breakpoint_diagnostics.py

"""Tests for cache_control breakpoint diagnostics and log-privacy switches.

Covers the three pieces added for the uncached-tail investigation:
- ``count_cache_breakpoints`` / ``log_cache_breakpoints`` (proxy helpers)
- the ``HEADROOM_LOG_PAYLOAD_PREVIEW`` kill switch (compression store)
- the injection guard that keeps proactive expansion out of breakpointed blocks
"""

from __future__ import annotations

import logging

from headroom.cache.compression_store import _payload_for_retrieval_log
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
from headroom.proxy.helpers import count_cache_breakpoints, log_cache_breakpoints

_CC = {"cache_control": {"type": "ephemeral"}}


def _claude_code_style_request() -> tuple[list[dict], list[dict], list[dict]]:
    """System/messages/tools shaped like a real Claude Code request."""
    system = [
        {"type": "text", "text": "You are Claude Code."},
        {"type": "text", "text": "project instructions", **_CC},
    ]
    tools = [
        {"name": "Bash", "input_schema": {}},
        {"name": "Read", "input_schema": {}, **_CC},
    ]
    messages = [
        {"role": "user", "content": [{"type": "text", "text": "hi", **_CC}]},
        {"role": "assistant", "content": [{"type": "text", "text": "ack"}]},
        {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": "t1",
                    "content": [{"type": "text", "text": "big output"}],
                    **_CC,
                }
            ],
        },
    ]
    return system, messages, tools


def test_count_cache_breakpoints_counts_all_sections() -> None:
    system, messages, tools = _claude_code_style_request()
    stats = count_cache_breakpoints(system, messages, tools)
    assert stats["system"] == 1
    assert stats["tools"] == 1
    assert stats["messages"] == 2
    assert stats["total"] == 4
    assert stats["message_count"] == 3
    assert stats["last_marker_tail"] == 0  # last message carries a marker


def test_count_cache_breakpoints_counts_nested_tool_result_markers() -> None:
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "tool_result",
                    "tool_use_id": "t1",
                    "content": [{"type": "text", "text": "out", **_CC}],
                }
            ],
        }
    ]
    stats = count_cache_breakpoints("plain system string", messages, None)
    assert stats["system"] == 0
    assert stats["tools"] == 0
    assert stats["messages"] == 1
    assert stats["last_marker_tail"] == 0


def test_count_cache_breakpoints_tail_tracks_last_marker() -> None:
    messages = [
        {"role": "user", "content": [{"type": "text", "text": "a", **_CC}]},
        {"role": "assistant", "content": [{"type": "text", "text": "b"}]},
        {"role": "user", "content": [{"type": "text", "text": "c"}]},
    ]
    stats = count_cache_breakpoints(None, messages, None)
    assert stats["last_marker_tail"] == 2
    assert count_cache_breakpoints(None, [], None)["last_marker_tail"] == -1


def test_log_cache_breakpoints_warns_on_dropped_marker(caplog) -> None:
    system, messages, tools = _claude_code_style_request()
    inbound = count_cache_breakpoints(system, messages, tools)
    # Transform "lost" the final breakpoint: strip it from the last message.
    stripped = [dict(m) for m in messages]
    stripped[2] = {
        "role": "user",
        "content": [{"type": "tool_result", "tool_use_id": "t1", "content": "compressed"}],
    }
    outbound = count_cache_breakpoints(system, stripped, tools)
    with caplog.at_level(logging.INFO, logger="headroom.proxy"):
        log_cache_breakpoints(request_id="r1", inbound=inbound, outbound=outbound)
    [record] = caplog.records
    assert record.levelno == logging.WARNING
    assert "dropped=true" in record.getMessage()
    assert "tail_grew=true" in record.getMessage()


def test_log_cache_breakpoints_info_when_preserved(caplog) -> None:
    system, messages, tools = _claude_code_style_request()
    stats = count_cache_breakpoints(system, messages, tools)
    with caplog.at_level(logging.INFO, logger="headroom.proxy"):
        log_cache_breakpoints(request_id="r1", inbound=stats, outbound=stats)
    [record] = caplog.records
    assert record.levelno == logging.INFO
    assert "dropped=false" in record.getMessage()


def test_payload_preview_disabled_omits_content(monkeypatch) -> None:
    monkeypatch.setenv("HEADROOM_LOG_PAYLOAD_PREVIEW", "0")
    payload = "secret file contents: api_key=sk-abcdefghijklmnop"
    event = _payload_for_retrieval_log(payload)
    assert event["payload_preview"] == ""
    assert event["payload_preview_chars"] == 0
    assert event["payload_chars"] == len(payload)
    assert event["payload_truncated"] is True


def test_payload_preview_enabled_by_default(monkeypatch) -> None:
    monkeypatch.delenv("HEADROOM_LOG_PAYLOAD_PREVIEW", raising=False)
    event = _payload_for_retrieval_log("hello world")
    assert event["payload_preview"] == "hello world"


def test_append_context_skips_breakpointed_text_block() -> None:
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "breakpointed", **_CC},
                {"type": "text", "text": "free"},
            ],
        }
    ]
    result = AnthropicHandlerMixin._append_context_to_latest_non_frozen_user_turn(
        messages, "CTX", frozen_message_count=0
    )
    blocks = result[0]["content"]
    assert blocks[0]["text"] == "breakpointed"  # untouched
    assert blocks[1]["text"].endswith("CTX")


def test_append_context_no_eligible_block_returns_unchanged() -> None:
    messages = [
        {
            "role": "user",
            "content": [{"type": "text", "text": "breakpointed", **_CC}],
        }
    ]
    result = AnthropicHandlerMixin._append_context_to_latest_non_frozen_user_turn(
        messages, "CTX", frozen_message_count=0
    )
    assert result == messages


def test_count_cache_breakpoints_tolerates_malformed_shapes() -> None:
    messages = [
        "not-a-dict",
        {"role": "user", "content": ["scalar-block", {"type": "text", "text": "x", **_CC}]},
        {"role": "user", "content": "plain string"},
    ]
    stats = count_cache_breakpoints("system-as-string", messages, "tools-as-string")
    assert stats["system"] == 0
    assert stats["tools"] == 0
    assert stats["messages"] == 1
    assert stats["message_count"] == 3
    assert stats["last_marker_tail"] == 1

    empty = count_cache_breakpoints(None, None, None)
    assert empty["total"] == 0
    assert empty["message_count"] == 0