headroom/tests/test_sse_thinking_blocks.py

"""PR-A8 / P1-9: SSE delta arms for thinking, signature, citations.

The proxy used to handle only ``text_delta`` and ``input_json_delta``
events on Anthropic's stream. The remaining delta types
(``thinking_delta``, ``signature_delta``, ``citations_delta``) and the
``redacted_thinking`` content_block_start were silently dropped, so any
non-streaming retry path that reconstructed the response from the SSE
stream produced an unsigned thinking block (rejected by Anthropic on
replay) or empty citations.

These tests pin the new contract:

- ``thinking_delta`` text appends to ``block.thinking_buffer`` and is
  promoted to ``block.thinking`` on ``content_block_stop``.
- ``signature_delta`` sets ``block.signature`` (last-write-wins).
- ``citations_delta`` appends each citation object to ``block.citations``.
- ``redacted_thinking`` content_block_start preserves the opaque
  ``data`` field as-is.
"""

from __future__ import annotations

import json
from typing import Any

import pytest

from headroom.proxy.handlers.streaming import StreamingMixin


class _Parser(StreamingMixin):
    """Subclass that exposes the parser without the rest of the proxy."""


def _build_sse(events: list[dict[str, Any]]) -> str:
    """Render a list of event dicts as an SSE payload string."""
    out: list[str] = []
    for ev in events:
        out.append(f"event: {ev['type']}")
        out.append(f"data: {json.dumps(ev)}")
        out.append("")  # event terminator
    return "\n".join(out) + "\n"


def test_thinking_delta_accumulated() -> None:
    parser = _Parser()
    events = [
        {"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
        {
            "type": "content_block_start",
            "index": 0,
            "content_block": {"type": "thinking", "thinking": ""},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "thinking_delta", "thinking": "Let me consider "},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "thinking_delta", "thinking": "the question carefully."},
        },
        {"type": "content_block_stop", "index": 0},
    ]
    sse = _build_sse(events)
    response = parser._parse_sse_to_response(sse, "anthropic")
    assert response is not None
    assert len(response["content"]) == 1
    block = response["content"][0]
    assert block["type"] == "thinking"
    assert block["thinking"] == "Let me consider the question carefully."


def test_signature_delta_preserved() -> None:
    parser = _Parser()
    events = [
        {"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
        {
            "type": "content_block_start",
            "index": 0,
            "content_block": {"type": "thinking", "thinking": ""},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "thinking_delta", "thinking": "hmm"},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "signature_delta", "signature": "sig_abc123_v1"},
        },
        {"type": "content_block_stop", "index": 0},
    ]
    sse = _build_sse(events)
    response = parser._parse_sse_to_response(sse, "anthropic")
    assert response is not None
    block = response["content"][0]
    assert block["signature"] == "sig_abc123_v1"
    # Last-write-wins semantics — second signature_delta overrides.
    events2 = events + [
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "signature_delta", "signature": "sig_xyz999_v2"},
        },
    ]
    # Re-emit with the corrected ordering: stop must come after all deltas.
    events2 = [e for e in events2 if e["type"] != "content_block_stop"]
    events2.append({"type": "content_block_stop", "index": 0})
    response2 = parser._parse_sse_to_response(_build_sse(events2), "anthropic")
    assert response2 is not None
    assert response2["content"][0]["signature"] == "sig_xyz999_v2"


def test_citations_delta_accumulated() -> None:
    parser = _Parser()
    events = [
        {"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
        {
            "type": "content_block_start",
            "index": 0,
            "content_block": {"type": "text", "text": ""},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {"type": "text_delta", "text": "Per source A"},
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {
                "type": "citations_delta",
                "citation": {
                    "type": "page_location",
                    "cited_text": "abc",
                    "document_index": 0,
                },
            },
        },
        {
            "type": "content_block_delta",
            "index": 0,
            "delta": {
                "type": "citations_delta",
                "citation": {
                    "type": "page_location",
                    "cited_text": "def",
                    "document_index": 1,
                },
            },
        },
        {"type": "content_block_stop", "index": 0},
    ]
    sse = _build_sse(events)
    response = parser._parse_sse_to_response(sse, "anthropic")
    assert response is not None
    block = response["content"][0]
    citations = block["citations"]
    assert len(citations) == 2
    assert citations[0]["cited_text"] == "abc"
    assert citations[1]["cited_text"] == "def"


def test_redacted_thinking_data_preserved() -> None:
    parser = _Parser()
    redacted_blob = "ENC:" + ("x" * 200)
    events = [
        {"type": "message_start", "message": {"id": "msg_1", "model": "claude-opus-4"}},
        {
            "type": "content_block_start",
            "index": 0,
            "content_block": {"type": "redacted_thinking", "data": redacted_blob},
        },
        {"type": "content_block_stop", "index": 0},
    ]
    sse = _build_sse(events)
    response = parser._parse_sse_to_response(sse, "anthropic")
    assert response is not None
    block = response["content"][0]
    assert block["type"] == "redacted_thinking"
    # `data` field MUST be preserved byte-for-byte for signature
    # validation on the next turn.
    assert block["data"] == redacted_blob


def test_response_to_sse_preserves_thinking_redacted_and_citations() -> None:
    parser = _Parser()
    redacted_blob = "ENC:" + ("y" * 200)
    response = {
        "id": "msg_2",
        "model": "claude-opus-4",
        "role": "assistant",
        "content": [
            {"type": "thinking", "thinking": "plan carefully", "signature": "sig_123"},
            {
                "type": "text",
                "text": "Per source A",
                "citations": [
                    {
                        "type": "page_location",
                        "cited_text": "abc",
                        "document_index": 0,
                    }
                ],
            },
            {"type": "redacted_thinking", "data": redacted_blob},
        ],
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": 3},
    }

    sse_text = b"".join(parser._response_to_sse(response, "anthropic")).decode("utf-8")

    assert "thinking_delta" in sse_text
    assert "signature_delta" in sse_text
    assert "citations_delta" in sse_text
    assert "redacted_thinking" in sse_text
    assert redacted_blob in sse_text

    round_tripped = parser._parse_sse_to_response(sse_text, "anthropic")
    assert round_tripped is not None
    assert round_tripped["content"][0]["thinking"] == "plan carefully"
    assert round_tripped["content"][0]["signature"] == "sig_123"
    assert round_tripped["content"][1]["citations"][0]["cited_text"] == "abc"
    assert round_tripped["content"][2]["data"] == redacted_blob


def test_response_to_sse_rejects_unknown_content_block() -> None:
    parser = _Parser()

    with pytest.raises(ValueError, match="Unsupported Anthropic content block type"):
        parser._response_to_sse(
            {"content": [{"type": "future_block", "payload": "preserve me"}]},
            "anthropic",
        )