free-claude-code/smoke/test_stream_contracts.py

from __future__ import annotations

from collections.abc import Iterable

import pytest

from messaging.event_parser import parse_cli_event
from messaging.transcript import RenderCtx, TranscriptBuffer
from providers.common import (
    ContentType,
    HeuristicToolParser,
    SSEBuilder,
    ThinkTagParser,
)
from smoke.lib.sse import (
    assert_anthropic_stream_contract,
    event_names,
    has_tool_use,
    parse_sse_text,
    text_content,
    thinking_content,
)

pytestmark = [
    pytest.mark.live,
    pytest.mark.smoke_target("contract"),
    pytest.mark.smoke_target("thinking"),
]


def test_interleaved_thinking_text_blocks_are_valid() -> None:
    events = _parse_builder_events(
        _interleaved_thinking_text_events(
            ("first thought", "first answer", "second thought", "final answer")
        )
    )
    assert_anthropic_stream_contract(events)
    assert event_names(events).count("content_block_start") == 4
    assert thinking_content(events) == "first thoughtsecond thought"
    assert text_content(events) == "first answerfinal answer"


def test_split_think_tags_preserve_text_and_thinking() -> None:
    events = _parse_builder_events(
        _events_from_text_chunks(["before <thi", "nk>hidden", "</think> after"])
    )
    assert_anthropic_stream_contract(events)
    assert thinking_content(events) == "hidden"
    assert text_content(events) == "before  after"


def test_mixed_reasoning_content_and_think_tags_keep_order() -> None:
    builder = SSEBuilder("msg_smoke", "smoke-model")
    chunks = [builder.message_start()]
    chunks.extend(builder.ensure_thinking_block())
    chunks.append(builder.emit_thinking_delta("reasoning field"))
    chunks.extend(
        _events_from_text_chunks([" visible <think>tagged</think> done"], builder)
    )
    chunks.extend(builder.close_all_blocks())
    chunks.append(builder.message_delta("end_turn", 10))
    chunks.append(builder.message_stop())

    events = parse_sse_text("".join(chunks))
    assert_anthropic_stream_contract(events)
    assert thinking_content(events) == "reasoning fieldtagged"
    assert text_content(events) == " visible  done"


def test_thinking_tool_text_and_transcript_order_contract() -> None:
    builder = SSEBuilder("msg_smoke", "smoke-model")
    chunks = [builder.message_start()]
    chunks.extend(builder.ensure_thinking_block())
    chunks.append(builder.emit_thinking_delta("inspect first"))
    chunks.extend(builder.close_content_blocks())
    tool_block_index = builder.blocks.allocate_index()
    chunks.append(
        builder.content_block_start(
            tool_block_index, "tool_use", id="toolu_1", name="Read"
        )
    )
    chunks.append(
        builder.content_block_delta(
            tool_block_index, "input_json_delta", '{"file":"README.md"}'
        )
    )
    chunks.append(builder.content_block_stop(tool_block_index))
    chunks.extend(builder.ensure_text_block())
    chunks.append(builder.emit_text_delta("done"))
    chunks.extend(builder.close_all_blocks())
    chunks.append(builder.message_delta("end_turn", 20))
    chunks.append(builder.message_stop())

    events = parse_sse_text("".join(chunks))
    assert_anthropic_stream_contract(events)
    assert has_tool_use(events)

    transcript = TranscriptBuffer()
    for event in events:
        for parsed in parse_cli_event(event.data):
            transcript.apply(parsed)
    rendered = transcript.render(_render_ctx(), limit_chars=3900, status=None)
    assert (
        rendered.find("inspect first")
        < rendered.find("Tool call:")
        < rendered.find("done")
    )


def test_enable_thinking_false_suppresses_reasoning_only() -> None:
    events = _parse_builder_events(
        _events_from_text_chunks(
            ["hello <think>secret</think> world"], enable_thinking=False
        )
    )
    assert_anthropic_stream_contract(events)
    assert "secret" not in thinking_content(events)
    assert text_content(events) == "hello  world"


def test_task_tool_arguments_force_foreground_execution() -> None:
    parser = HeuristicToolParser()
    filtered, detected = parser.feed(
        "● <function=Task><parameter=description>Inspect</parameter>"
        "<parameter=run_in_background>true</parameter> trailing"
    )
    detected.extend(parser.flush())
    assert "trailing" in filtered
    task = detected[0]
    assert task["name"] == "Task"
    if isinstance(task.get("input"), dict):
        task["input"]["run_in_background"] = False
    assert task["input"]["run_in_background"] is False


def _interleaved_thinking_text_events(
    parts: tuple[str, str, str, str],
) -> Iterable[str]:
    builder = SSEBuilder("msg_smoke", "smoke-model")
    yield builder.message_start()
    yield from builder.ensure_thinking_block()
    yield builder.emit_thinking_delta(parts[0])
    yield from builder.ensure_text_block()
    yield builder.emit_text_delta(parts[1])
    yield from builder.ensure_thinking_block()
    yield builder.emit_thinking_delta(parts[2])
    yield from builder.ensure_text_block()
    yield builder.emit_text_delta(parts[3])
    yield from builder.close_all_blocks()
    yield builder.message_delta("end_turn", 20)
    yield builder.message_stop()


def _events_from_text_chunks(
    chunks: list[str],
    builder: SSEBuilder | None = None,
    *,
    enable_thinking: bool = True,
) -> list[str]:
    sse = builder or SSEBuilder("msg_smoke", "smoke-model")
    out: list[str] = [] if builder else [sse.message_start()]
    parser = ThinkTagParser()

    for chunk in chunks:
        out.extend(_emit_parser_parts(sse, parser.feed(chunk), enable_thinking))

    remaining = parser.flush()
    if remaining is not None:
        out.extend(_emit_parser_parts(sse, [remaining], enable_thinking))

    if builder is None:
        out.extend(sse.close_all_blocks())
        out.append(sse.message_delta("end_turn", 20))
        out.append(sse.message_stop())
    return out


def _emit_parser_parts(
    builder: SSEBuilder,
    parts: Iterable,
    enable_thinking: bool,
) -> list[str]:
    out: list[str] = []
    for part in parts:
        if part.type == ContentType.THINKING:
            if enable_thinking:
                out.extend(builder.ensure_thinking_block())
                out.append(builder.emit_thinking_delta(part.content))
            continue
        out.extend(builder.ensure_text_block())
        out.append(builder.emit_text_delta(part.content))
    return out


def _parse_builder_events(chunks: Iterable[str]):
    return parse_sse_text("".join(chunks))


def _render_ctx() -> RenderCtx:
    return RenderCtx(
        bold=lambda text: f"*{text}*",
        code_inline=lambda text: f"`{text}`",
        escape_code=lambda text: text,
        escape_text=lambda text: text,
        render_markdown=lambda text: text,
    )