# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Cohere v2 SSE stream conversion in
``vllm/entrypoints/cohere/serving.py``.

The stream-translation entry point is
:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an
async iterable of OpenAI SSE chunks into Cohere's
``message-start → (content|tool-call|citation)* → message-end → [DONE]``
event stream.

We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in
isolation, plus a handful of end-to-end scenarios that exercise the
state machine.
"""

import json
from collections.abc import AsyncGenerator
from typing import Any

import pytest

from vllm.entrypoints.cohere.cohere_chat_message import (
    Citation as VLLMCitation,
)
from vllm.entrypoints.cohere.cohere_chat_message import (
    CitationSource,
)
from vllm.entrypoints.cohere.protocol import (
    CohereChatV2Request,
    MessageStartEvent,
)
from vllm.entrypoints.cohere.serving import (
    _DONE_FRAME,
    CohereServingChatV2,
    ContentBlockType,
    _emit,
    _sse,
    _StreamState,
)
from vllm.entrypoints.openai.chat_completion.protocol import (
    ChatCompletionStreamResponse,
)

# ----------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------


class _FakeServing(CohereServingChatV2):
    """Lightweight stand-in for :class:`CohereServingChatV2` that skips
    the heavy ``OpenAIServingChat.__init__`` chain (which would need a
    real engine client, model registry, render service, etc.).

    Only ``_is_reasoning_model`` is read by the methods under test
    (``_chat_completion_stream_to_v2`` and the per-delta handlers); the
    rest is dead weight for unit testing.
    """

    def __init__(self, is_reasoning_model: bool = True) -> None:
        # Intentionally skipping super().__init__ — see class docstring.
        self._is_reasoning_model = is_reasoning_model


def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2:
    return _FakeServing(is_reasoning_model=is_reasoning_model)


def _parse_event(frame: str) -> dict[str, Any]:
    """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload."""
    assert frame.startswith("data: ")
    assert frame.endswith("\n\n")
    return json.loads(frame[len("data: ") : -2])


def _make_chunk(
    *,
    chunk_id: str = "chunk_0",
    role: str | None = None,
    content: str | None = None,
    reasoning: str | None = None,
    tool_calls: list[dict[str, Any]] | None = None,
    finish_reason: str | None = None,
    stop_reason: int | str | None = None,
    usage: dict[str, Any] | None = None,
    omit_choices: bool = False,
    citations: list[Any] | None = None,
) -> str:
    """Build the ``data: {...}\\n\\n`` SSE frame the production code
    consumes."""
    delta: dict[str, Any] = {}
    if role is not None:
        delta["role"] = role
    if content is not None:
        delta["content"] = content
    if reasoning is not None:
        delta["reasoning"] = reasoning
    if tool_calls is not None:
        delta["tool_calls"] = tool_calls
    if citations is not None:
        delta["citations"] = citations

    payload: dict[str, Any] = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": 0,
        "model": "m",
    }
    if not omit_choices:
        payload["choices"] = [
            {
                "index": 0,
                "delta": delta,
                "finish_reason": finish_reason,
                "stop_reason": stop_reason,
            }
        ]
    else:
        payload["choices"] = []
    if usage is not None:
        payload["usage"] = usage
    # Use ChatCompletionStreamResponse to normalize the payload shape.
    chunk = ChatCompletionStreamResponse.model_validate(payload)
    return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n"


def _make_done() -> str:
    return "data: [DONE]\n\n"


async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]:
    for item in items:
        yield item


async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]:
    """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect
    the emitted SSE frames."""
    request = CohereChatV2Request(
        model="m", messages=[{"role": "user", "content": "hi"}], stream=True
    )
    gen = serving._chat_completion_stream_to_v2(_stream_from(items), request)
    return [frame async for frame in gen]


# ======================================================================
# Low-level helpers: _sse, _emit, _DONE_FRAME
# ======================================================================


class TestSSEHelpers:
    def test_sse_wraps_payload(self):
        assert _sse("hello") == "data: hello\n\n"

    def test_done_frame_constant(self):
        # cohere-python and Fern-generated clients key off this exact
        # sentinel; keep it byte-for-byte stable.
        assert _DONE_FRAME == "data: [DONE]\n\n"

    def test_emit_serializes_event_with_type_discriminator(self):
        frame = _emit(
            MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}})
        )
        payload = _parse_event(frame)
        assert payload["type"] == "message-start"
        assert payload["id"] == "abc"
        assert payload["delta"] == {"message": {"role": "assistant"}}


# ======================================================================
# _StreamState
# ======================================================================


class TestStreamState:
    def test_content_index_monotonic(self):
        st = _StreamState()
        assert st.next_content_index() == 0
        assert st.next_content_index() == 1
        assert st.next_content_index() == 2

    def test_citation_index_separate_from_content_index(self):
        st = _StreamState()
        st.next_content_index()  # 0
        st.next_content_index()  # 1
        # Citation indexing is independent of content indexing.
        assert st.next_citation_index() == 0
        assert st.next_citation_index() == 1


# ======================================================================
# _close_open_blocks
# ======================================================================


class TestCloseOpenBlocks:
    def test_no_block_open_emits_nothing(self):
        serving = _serving()
        state = _StreamState()
        assert serving._close_open_blocks(state) == []

    def test_text_block_emits_content_end(self):
        serving = _serving()
        state = _StreamState()
        state.active_block = ContentBlockType.TEXT
        state.active_block_index = 1
        out = serving._close_open_blocks(state)
        assert len(out) == 1
        payload = _parse_event(out[0])
        assert payload == {"type": "content-end", "index": 1}
        assert state.active_block is None
        assert state.active_block_index is None

    def test_thinking_block_emits_content_end(self):
        serving = _serving()
        state = _StreamState()
        state.active_block = ContentBlockType.THINKING
        state.active_block_index = 3
        out = serving._close_open_blocks(state)
        payload = _parse_event(out[0])
        assert payload == {"type": "content-end", "index": 3}

    def test_tool_call_block_emits_tool_call_end(self):
        serving = _serving()
        state = _StreamState()
        state.active_block = ContentBlockType.TOOL_CALL
        state.active_tool_index = 7
        out = serving._close_open_blocks(state)
        payload = _parse_event(out[0])
        assert payload == {"type": "tool-call-end", "index": 7}
        assert state.active_tool_index is None


# ======================================================================
# _handle_text_delta
# ======================================================================


class TestHandleTextDelta:
    def test_opens_block_first_time(self):
        serving = _serving()
        state = _StreamState()
        events = serving._handle_text_delta(state, "Hi")
        assert len(events) == 2
        start = _parse_event(events[0])
        delta = _parse_event(events[1])
        assert start["type"] == "content-start"
        assert start["index"] == 0
        assert start["delta"]["message"]["content"]["type"] == "text"
        assert delta["type"] == "content-delta"
        assert delta["index"] == 0
        assert delta["delta"]["message"]["content"]["text"] == "Hi"
        assert state.active_block == ContentBlockType.TEXT
        assert state.active_block_index == 0

    def test_continues_block_with_just_delta(self):
        serving = _serving()
        state = _StreamState()
        serving._handle_text_delta(state, "Hi")
        events = serving._handle_text_delta(state, " there")
        # Only a delta event, no new content-start.
        assert len(events) == 1
        delta = _parse_event(events[0])
        assert delta["type"] == "content-delta"
        assert delta["delta"]["message"]["content"]["text"] == " there"

    def test_switches_from_thinking_block(self):
        serving = _serving(is_reasoning_model=True)
        state = _StreamState()
        # Open a thinking block first, then switch to text.
        serving._handle_thinking_delta(state, "ponder")
        events = serving._handle_text_delta(state, "answer")
        types = [_parse_event(ev)["type"] for ev in events]
        assert types == ["content-end", "content-start", "content-delta"]
        # The text block gets a new index (1), distinct from thinking's 0.
        assert _parse_event(events[1])["index"] == 1
        assert state.active_block == ContentBlockType.TEXT


# ======================================================================
# _handle_thinking_delta
# ======================================================================


class TestHandleThinkingDelta:
    def test_reasoning_model_opens_thinking_block(self):
        serving = _serving(is_reasoning_model=True)
        state = _StreamState()
        events = serving._handle_thinking_delta(state, "thought")
        assert len(events) == 2
        start = _parse_event(events[0])
        delta = _parse_event(events[1])
        assert start["type"] == "content-start"
        assert start["delta"]["message"]["content"]["type"] == "thinking"
        assert delta["type"] == "content-delta"
        assert delta["delta"]["message"]["content"]["thinking"] == "thought"
        assert state.active_block == ContentBlockType.THINKING

    def test_reasoning_model_continues_thinking_block(self):
        serving = _serving(is_reasoning_model=True)
        state = _StreamState()
        serving._handle_thinking_delta(state, "first")
        events = serving._handle_thinking_delta(state, " more")
        assert len(events) == 1
        assert _parse_event(events[0])["type"] == "content-delta"

    def test_non_reasoning_model_emits_tool_plan_delta(self):
        # Older Command models stream reasoning as ``tool_plan`` deltas;
        # no content-start/end pair is emitted.
        serving = _serving(is_reasoning_model=False)
        state = _StreamState()
        events = serving._handle_thinking_delta(state, "planning")
        assert len(events) == 1
        payload = _parse_event(events[0])
        assert payload["type"] == "tool-plan-delta"
        assert payload["delta"]["message"]["tool_plan"] == "planning"
        # ``tool_plan`` deltas don't claim an active content block.
        assert state.active_block is None

    def test_non_reasoning_model_closes_open_text_block(self):
        serving = _serving(is_reasoning_model=False)
        state = _StreamState()
        # Open a text block first.
        serving._handle_text_delta(state, "answer")
        events = serving._handle_thinking_delta(state, "rethink")
        types = [_parse_event(ev)["type"] for ev in events]
        assert types == ["content-end", "tool-plan-delta"]


# ======================================================================
# _handle_tool_call_deltas
# ======================================================================


class TestHandleToolCallDeltas:
    def test_new_tool_call_opens_tool_call_start(self):
        serving = _serving()
        state = _StreamState()
        deltas = [
            type(
                "Delta",
                (),
                {
                    "index": 0,
                    "id": "c1",
                    "function": type(
                        "Fn", (), {"name": "calc", "arguments": '{"x":'}
                    )(),
                },
            )()
        ]
        events = serving._handle_tool_call_deltas(state, deltas)
        assert len(events) == 1
        payload = _parse_event(events[0])
        assert payload["type"] == "tool-call-start"
        assert payload["index"] == 0
        tc = payload["delta"]["message"]["tool_calls"]
        assert tc["id"] == "c1"
        assert tc["function"]["name"] == "calc"
        assert tc["function"]["arguments"] == '{"x":'
        assert state.active_block == ContentBlockType.TOOL_CALL
        assert state.active_tool_index == 0
        assert 0 in state.tool_calls_seen

    def test_subsequent_arguments_emit_delta(self):
        serving = _serving()
        state = _StreamState()
        # First call: start.
        first = [
            type(
                "Delta",
                (),
                {
                    "index": 0,
                    "id": "c1",
                    "function": type("Fn", (), {"name": "calc", "arguments": ""})(),
                },
            )()
        ]
        serving._handle_tool_call_deltas(state, first)
        # Second call: same index, additional arguments fragment.
        more = [
            type(
                "Delta",
                (),
                {
                    "index": 0,
                    "id": None,
                    "function": type("Fn", (), {"name": None, "arguments": "1}"})(),
                },
            )()
        ]
        events = serving._handle_tool_call_deltas(state, more)
        assert len(events) == 1
        payload = _parse_event(events[0])
        assert payload["type"] == "tool-call-delta"
        assert payload["index"] == 0
        assert (
            payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}"
        )

    def test_new_tool_call_closes_existing_content_block(self):
        serving = _serving()
        state = _StreamState()
        # Open a text block, then start a tool call.
        serving._handle_text_delta(state, "I'll call:")
        deltas = [
            type(
                "Delta",
                (),
                {
                    "index": 0,
                    "id": "c1",
                    "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(),
                },
            )()
        ]
        events = serving._handle_tool_call_deltas(state, deltas)
        types = [_parse_event(ev)["type"] for ev in events]
        assert types == ["content-end", "tool-call-start"]


# ======================================================================
# _handle_citation_deltas
# ======================================================================


class TestHandleCitationDeltas:
    """Test the per-delta emitter. Sources are pre-resolved by the
    reasoning parser via the position map forwarded through
    ``chat_template_kwargs`` (see :func:`_melody_sources_to_vllm` in
    ``vllm/reasoning/cohere_command_reasoning_parser.py``), so the
    handler only has to coerce to the SDK wire shape, apply the
    ``THINKING_CONTENT`` -> ``PLAN`` rewrite, and drop citations with
    no attributable sources.
    """

    @staticmethod
    def _resolved_source(**overrides: Any) -> CitationSource:
        base = dict(
            type="document",
            id="doc_hamlet",
            document={"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
        )
        base.update(overrides)
        return CitationSource(**base)

    def test_resolved_citation_emits_start_and_end_with_real_id(self):
        serving = _serving()
        state = _StreamState()
        citations = [
            VLLMCitation(
                start=0,
                end=11,
                text="Shakespeare",
                sources=[self._resolved_source()],
                type="TEXT_CONTENT",
            )
        ]
        events = serving._handle_citation_deltas(state, citations)
        assert len(events) == 2
        start = _parse_event(events[0])
        end = _parse_event(events[1])
        assert start["type"] == "citation-start"
        assert start["index"] == 0
        cit_payload = start["delta"]["message"]["citations"]
        assert cit_payload["start"] == 0
        assert cit_payload["end"] == 11
        assert cit_payload["text"] == "Shakespeare"
        assert cit_payload["sources"] == [
            {
                "type": "document",
                "id": "doc_hamlet",
                "document": {"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
            }
        ]
        assert end == {"type": "citation-end", "index": 0}

    def test_multiple_sources_survive_to_wire(self):
        serving = _serving()
        state = _StreamState()
        citations = [
            VLLMCitation(
                start=0,
                end=5,
                text="works",
                sources=[
                    self._resolved_source(id="d0", document={"id": "d0"}),
                    self._resolved_source(id="d1", document={"id": "d1"}),
                ],
                type="TEXT_CONTENT",
            )
        ]
        events = serving._handle_citation_deltas(state, citations)
        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
        ids = [s["id"] for s in cit_payload["sources"]]
        assert ids == ["d0", "d1"]

    def test_source_without_id_streamed_with_payload(self):
        # Match the cohere api: id-less tool sources (from text-only
        # tool results) still ride through with their ``tool_output``
        # payload; only the ``id`` field is omitted.
        serving = _serving()
        state = _StreamState()
        citations = [
            VLLMCitation(
                start=0,
                end=5,
                text="hello",
                sources=[CitationSource(type="tool", tool_output={"content": "hi"})],
            )
        ]
        events = serving._handle_citation_deltas(state, citations)
        assert len(events) >= 1
        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
        assert cit_payload["sources"] == [
            {"type": "tool", "tool_output": {"content": "hi"}}
        ]
        assert "id" not in cit_payload["sources"][0]

    def test_citation_with_no_sources_streamed_dropped(self):
        # Citations whose source list is entirely empty are still
        # dropped so we don't emit an unattributed anchor.
        serving = _serving()
        state = _StreamState()
        citations = [
            VLLMCitation(
                start=0,
                end=5,
                text="hello",
                sources=[],
            )
        ]
        events = serving._handle_citation_deltas(state, citations)
        assert events == []

    def test_plan_rewrite_on_non_reasoning_model(self):
        serving = _serving(is_reasoning_model=False)
        state = _StreamState()
        citations = [
            VLLMCitation(
                start=0,
                end=5,
                text="hello",
                sources=[self._resolved_source()],
                type="THINKING_CONTENT",
            )
        ]
        events = serving._handle_citation_deltas(state, citations)
        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
        assert cit_payload["type"] == "PLAN"


# ======================================================================
# _build_message_end_event
# ======================================================================


class TestBuildMessageEndEvent:
    def test_without_usage(self):
        serving = _serving()
        frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop")
        payload = _parse_event(frame)
        assert payload["type"] == "message-end"
        assert payload["id"] == "abc"
        assert payload["delta"]["finish_reason"] == "COMPLETE"
        assert "usage" not in payload["delta"]

    def test_with_usage(self):
        serving = _serving()
        chunk = ChatCompletionStreamResponse.model_validate(
            {
                "id": "abc",
                "object": "chat.completion.chunk",
                "created": 0,
                "model": "m",
                "choices": [],
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                },
            }
        )
        frame = serving._build_message_end_event(
            chunk_id="abc", finish_reason="length", usage_chunk=chunk
        )
        payload = _parse_event(frame)
        assert payload["delta"]["finish_reason"] == "MAX_TOKENS"
        usage = payload["delta"]["usage"]
        assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5}
        assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5}
        assert "cached_tokens" not in usage

    def test_with_cached_tokens(self):
        serving = _serving()
        chunk = ChatCompletionStreamResponse.model_validate(
            {
                "id": "abc",
                "object": "chat.completion.chunk",
                "created": 0,
                "model": "m",
                "choices": [],
                "usage": {
                    "prompt_tokens": 10,
                    "completion_tokens": 5,
                    "total_tokens": 15,
                    "prompt_tokens_details": {"cached_tokens": 3},
                },
            }
        )
        frame = serving._build_message_end_event(
            chunk_id="abc", finish_reason="stop", usage_chunk=chunk
        )
        payload = _parse_event(frame)
        assert payload["delta"]["usage"]["cached_tokens"] == 3


# ======================================================================
# End-to-end: _chat_completion_stream_to_v2
# ======================================================================


class TestChatCompletionStreamToV2:
    """End-to-end stream lifecycle tests."""

    @pytest.mark.asyncio
    async def test_text_only_happy_path_emits_full_lifecycle(self):
        serving = _serving(is_reasoning_model=True)
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(content="Hi"),
            _make_chunk(content=" there"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 4,
                    "completion_tokens": 2,
                    "total_tokens": 6,
                },
            ),
        ]
        frames = await _drain(serving, items)
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        # message-start, content-start, content-delta, content-delta,
        # content-end, message-end, then [DONE] as the last frame.
        assert types == [
            "message-start",
            "content-start",
            "content-delta",
            "content-delta",
            "content-end",
            "message-end",
        ]
        assert frames[-1] == _DONE_FRAME
        # message-end should carry usage stats from the trailing chunk.
        end_payload = _parse_event(frames[-2])
        assert end_payload["delta"]["finish_reason"] == "COMPLETE"
        assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4

    @pytest.mark.asyncio
    async def test_stop_sequence_finish_reason(self):
        serving = _serving()
        frames = await _drain(
            serving,
            [
                _make_chunk(role="assistant"),
                _make_chunk(
                    finish_reason="stop",
                    stop_reason="<END>",
                ),
                _make_chunk(omit_choices=True),
            ],
        )
        end_payload = _parse_event(frames[-2])
        assert end_payload["delta"]["finish_reason"] == "STOP_SEQUENCE"

    @pytest.mark.asyncio
    async def test_text_then_tool_call_closes_text_first(self):
        serving = _serving(is_reasoning_model=True)
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(content="planning..."),
            _make_chunk(
                tool_calls=[
                    {
                        "index": 0,
                        "id": "c1",
                        "type": "function",
                        "function": {"name": "calc", "arguments": "{}"},
                    }
                ]
            ),
            _make_chunk(finish_reason="tool_calls"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        # Text block is opened+delta, then closed before the tool-call
        # opens, and the final close happens on the usage chunk path.
        assert types == [
            "message-start",
            "content-start",
            "content-delta",
            "content-end",
            "tool-call-start",
            "tool-call-end",
            "message-end",
        ]
        # finish_reason captured from the prior chunk.
        end = _parse_event(frames[-2])
        assert end["delta"]["finish_reason"] == "TOOL_CALL"

    @pytest.mark.asyncio
    async def test_thinking_then_text_reasoning_model(self):
        serving = _serving(is_reasoning_model=True)
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(reasoning="thinking..."),
            _make_chunk(content="answer"),
            _make_chunk(finish_reason="stop"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        # Thinking block opens with index 0; text block reopens with index 1.
        assert types == [
            "message-start",
            "content-start",
            "content-delta",
            "content-end",
            "content-start",
            "content-delta",
            "content-end",
            "message-end",
        ]
        thinking_start = _parse_event(frames[1])
        assert thinking_start["delta"]["message"]["content"]["type"] == "thinking"
        text_start = _parse_event(frames[4])
        assert text_start["delta"]["message"]["content"]["type"] == "text"

    @pytest.mark.asyncio
    async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self):
        serving = _serving(is_reasoning_model=False)
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(reasoning="planning"),
            _make_chunk(
                tool_calls=[
                    {
                        "index": 0,
                        "id": "c1",
                        "type": "function",
                        "function": {"name": "f", "arguments": "{}"},
                    }
                ]
            ),
            _make_chunk(finish_reason="tool_calls"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        assert types == [
            "message-start",
            "tool-plan-delta",
            "tool-call-start",
            "tool-call-end",
            "message-end",
        ]
        # No thinking content blocks should be present.
        assert "content-start" not in types
        assert "content-end" not in types

    @pytest.mark.asyncio
    async def test_done_marker_in_middle_closes_open_block(self):
        # Some upstreams send [DONE] without a trailing usage-only chunk.
        # The translator must still emit message-end before [DONE] so
        # Cohere clients don't hang.
        serving = _serving(is_reasoning_model=True)
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(content="Hi"),
            _make_done(),
        ]
        frames = await _drain(serving, items)
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        assert types == [
            "message-start",
            "content-start",
            "content-delta",
            "content-end",
            "message-end",
        ]
        assert frames[-1] == _DONE_FRAME

    @pytest.mark.asyncio
    async def test_skips_empty_and_non_data_lines(self):
        serving = _serving()
        items = [
            "\n",
            "event: ping\n\n",
            _make_chunk(role="assistant"),
            "data: \n\n",  # empty data
            _make_chunk(content="Hi"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        assert frames[-1] == _DONE_FRAME
        # Should still produce a complete lifecycle.
        types = [_parse_event(f)["type"] for f in frames[:-1]]
        assert types[0] == "message-start"
        assert types[-1] == "message-end"

    @pytest.mark.asyncio
    async def test_exception_in_chunk_parsing_emits_error_message_end(self):
        # An invalid JSON payload after the first chunk triggers the
        # error path: a synthetic message-end with finish_reason=ERROR
        # followed by [DONE].
        serving = _serving()
        items = [
            _make_chunk(role="assistant"),
            "data: {not valid json}\n\n",
        ]
        frames = await _drain(serving, items)
        assert frames[-1] == _DONE_FRAME
        # Find the error-shaped message-end.
        error_end = _parse_event(frames[-2])
        assert error_end["type"] == "message-end"
        assert error_end["delta"]["finish_reason"] == "ERROR"
        assert "error" in error_end["delta"]

    @pytest.mark.asyncio
    async def test_error_message_end_is_sanitized(self):
        # The client-visible ``error`` field on the terminal message-end
        # must be routed through ``sanitize_message`` so tracebacks,
        # host filesystem paths, and Python object memory addresses
        # can't leak from a mid-stream exception into the SSE payload.
        serving = _serving()

        async def _raises_mid_stream() -> AsyncGenerator[str, None]:
            yield _make_chunk(role="assistant")
            raise RuntimeError(
                "boom <Foo at 0x7fabc0deface> reading /Users/dev/vllm/x.py"
            )

        request = CohereChatV2Request(
            model="m",
            messages=[{"role": "user", "content": "hi"}],
            stream=True,
        )
        frames = [
            f
            async for f in serving._chat_completion_stream_to_v2(
                _raises_mid_stream(), request
            )
        ]
        assert frames[-1] == _DONE_FRAME
        error_end = _parse_event(frames[-2])
        assert error_end["delta"]["finish_reason"] == "ERROR"
        err_msg = error_end["delta"]["error"]
        # Memory address stripped and path replaced with ``<path>``.
        assert "0x7fabc0deface" not in err_msg
        assert "/Users/dev/vllm/x.py" not in err_msg
        assert "<path>" in err_msg
        # The human-facing prefix should still survive sanitization.
        assert "boom" in err_msg

    @pytest.mark.asyncio
    async def test_citations_in_delta_emit_citation_events(self):
        # End-to-end streaming citation flow. The parser hands us
        # already-resolved sources (see ``_melody_sources_to_vllm``
        # in ``vllm/reasoning/cohere_command_reasoning_parser.py``);
        # the streaming loop just has to coerce and emit.
        serving = _serving()
        items = [
            _make_chunk(role="assistant"),
            _make_chunk(content="hello"),
            _make_chunk(
                citations=[
                    {
                        "start": 0,
                        "end": 5,
                        "text": "hello",
                        "sources": [
                            {
                                "type": "document",
                                "id": "d1",
                                "document": {"id": "d1", "text": "cited"},
                            }
                        ],
                        "type": "TEXT_CONTENT",
                    }
                ]
            ),
            _make_chunk(finish_reason="stop"),
            _make_chunk(
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        parsed = [_parse_event(f) for f in frames[:-1]]
        types = [p["type"] for p in parsed]
        assert "citation-start" in types
        assert "citation-end" in types
        cit_start = next(p for p in parsed if p["type"] == "citation-start")
        assert cit_start["delta"]["message"]["citations"]["sources"] == [
            {
                "type": "document",
                "id": "d1",
                "document": {"id": "d1", "text": "cited"},
            }
        ]

    @pytest.mark.asyncio
    async def test_first_chunk_emits_message_start_with_chunk_id(self):
        serving = _serving()
        items = [
            _make_chunk(chunk_id="my-id", role="assistant"),
            _make_chunk(chunk_id="my-id", content="hi"),
            _make_chunk(
                chunk_id="my-id",
                omit_choices=True,
                usage={
                    "prompt_tokens": 1,
                    "completion_tokens": 1,
                    "total_tokens": 2,
                },
            ),
        ]
        frames = await _drain(serving, items)
        ms = _parse_event(frames[0])
        assert ms["type"] == "message-start"
        assert ms["id"] == "my-id"
        assert ms["delta"]["message"]["role"] == "assistant"