lightseekorg--tokenspeed/test/runtime/test_inline_detokenizer_receiver.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Inline detokenization receiver tests.

Drive the ``AsyncLLM._inline_detokenize_one`` helper and the
``BatchTokenIDOut`` dispatch branch. They verify:

1. Flag-off regression — ``BatchTokenIDOut`` still flows through the
   raw-token path and produces an out_dict with an empty ``text``.
2. Flag-on inline emit — out_dict gains a ``text`` key populated by
   the per-request ``IncrementalDetokenizer`` and matches the shape
   the ``BatchStrOut`` branch produces byte-for-byte.
3. Per-request lifecycle — the inline detokenizer is lazily created
   per rid, persists across frames for the same rid, and is
   independent between rids.
4. Subprocess-vs-inline text parity — for a given sequence of
   frames, the cumulative ``state.text`` accumulated through the
   inline path equals what ``incremental_decode_batch`` would emit
   character-for-character.
5. Stream vs non-stream ``output_ids`` shape, stop trimming
   pass-through, and finish-reason propagation.

A ``_StubTokenizerManager`` bypasses ZMQ / ModelConfig / HF-tokenizer
bring-up so the tests can exercise the exact production code path
without GPU or network.
"""

from __future__ import annotations

import os
import sys
import types
import unittest
from typing import Any, Dict, List, Optional

# CI registration (AST-parsed, runtime no-op).
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ci_system.ci_register import register_cuda_ci  # noqa: E402

register_cuda_ci(est_time=60, suite="runtime-1gpu")

from transformers import AutoTokenizer  # noqa: E402

from tokenspeed.runtime.engine.async_llm import AsyncLLM  # noqa: E402
from tokenspeed.runtime.engine.collector import (  # noqa: E402
    RequestOutputCollector,
)
from tokenspeed.runtime.engine.detokenizer import (  # noqa: E402
    DecodeStatus,
    IncrementalDetokenizer,
    incremental_decode_batch,
)
from tokenspeed.runtime.engine.io_struct import BatchTokenIDOut  # noqa: E402
from tokenspeed.runtime.engine.output_processor import (  # noqa: E402
    OutputProcessor,
    ReqState,
)

_GPT2_TOKENIZER = "gpt2"


# ---------------------------------------------------------------------------
# Stubs
# ---------------------------------------------------------------------------


class _StubTokenizerManager(AsyncLLM):
    """Bypass ZMQ + ModelConfig + HF bring-up for unit tests.

    We only need the pieces touched by
    ``OutputProcessor.handle_batch_output`` and
    ``_inline_detokenize_one``: ``server_args``, ``tokenizer``, the
    ``rid_to_state`` map, and a handful of flags. Everything else
    that the real ``__init__`` populates (metrics, sockets, model
    config) is untouched because these tests never reach those
    paths.
    """

    def __init__(
        self,
        tokenizer: Any,
        *,
        enable_inline_detokenizer: bool = True,
        stream_output: bool = True,
        speculative_algorithm: Any = None,
    ) -> None:
        self.tokenizer = tokenizer
        self.processor = None
        self.rid_to_state: Dict[str, ReqState] = {}
        self.enable_metrics = False
        self.dump_requests_folder = False
        self.log_requests = False
        # Build a tiny ServerArgs-shaped object so the branch conditions in
        # ``handle_batch_output`` keep working without loading the real
        # ServerArgs dataclass (which pulls torch through ModelConfig).
        self.server_args = types.SimpleNamespace(
            enable_inline_detokenizer=enable_inline_detokenizer,
            stream_output=stream_output,
            speculative_algorithm=speculative_algorithm,
            skip_tokenizer_init=False,
        )
        # OutputProcessor holds a back-reference to this stub via
        # ``engine``.
        self.output_processor = OutputProcessor(self)


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


def _gpt2_tokenizer() -> Any:
    return AutoTokenizer.from_pretrained(_GPT2_TOKENIZER)


def _batch_token_id_out(
    rids: List[str],
    *,
    decode_ids: List[List[int]],
    decoded_texts: Optional[List[str]] = None,
    read_offsets: Optional[List[int]] = None,
    finished_reasons: Optional[List[Optional[Dict[str, Any]]]] = None,
    no_stop_trim: Optional[List[bool]] = None,
    skip_special_tokens: Optional[List[bool]] = None,
    spaces_between_special_tokens: Optional[List[bool]] = None,
    **overrides: Any,
) -> BatchTokenIDOut:
    """Build a ``BatchTokenIDOut`` with safe defaults."""
    n = len(rids)
    defaults: Dict[str, Any] = {
        "output_ids": None,
        "output_multi_ids": None,
        "prompt_tokens": [0] * n,
        "completion_tokens": [0] * n,
        "cached_tokens": [0] * n,
        "spec_verify_ct": [0] * n,
        "input_token_logprobs_val": [],
        "input_token_logprobs_idx": [],
        "output_token_logprobs_val": [],
        "output_token_logprobs_idx": [],
        "input_top_logprobs_val": [],
        "input_top_logprobs_idx": [],
        "output_top_logprobs_val": [],
        "output_top_logprobs_idx": [],
        "input_token_ids_logprobs_val": [],
        "input_token_ids_logprobs_idx": [],
        "output_token_ids_logprobs_val": [],
        "output_token_ids_logprobs_idx": [],
        "output_hidden_states": [[] for _ in range(n)],
        "batch_accept_draft_tokens": [],
        "output_extra_infos": [],
        "generated_time": 0,
    }
    defaults.update(overrides)
    return BatchTokenIDOut(
        rids=rids,
        finished_reasons=(
            finished_reasons if finished_reasons is not None else [None] * n
        ),
        decoded_texts=decoded_texts if decoded_texts is not None else [""] * n,
        decode_ids=decode_ids,
        read_offsets=read_offsets if read_offsets is not None else [0] * n,
        skip_special_tokens=(
            skip_special_tokens if skip_special_tokens is not None else [True] * n
        ),
        spaces_between_special_tokens=(
            spaces_between_special_tokens
            if spaces_between_special_tokens is not None
            else [True] * n
        ),
        no_stop_trim=no_stop_trim if no_stop_trim is not None else [False] * n,
        **defaults,
    )


class _StubReqObj:
    """Minimal stand-in for GenerateReqInput used by ``_handle_batch_output``."""

    def __init__(
        self,
        *,
        stream: bool = True,
        return_logprob: bool = False,
        rid: str = "r1",
        log_metrics: bool = False,
    ) -> None:
        self.stream = stream
        self.return_logprob = return_logprob
        self.rid = rid
        self.log_metrics = log_metrics
        # Fields normally consumed only when return_logprob=True — provide
        # benign defaults so the attribute access never raises.
        self.top_logprobs_num = []
        self.token_ids_logprob = []
        self.return_text_in_logprobs = False


def _mk_state(*, stream: bool = True, rid: str = "r1") -> ReqState:
    return ReqState(
        RequestOutputCollector(),
        False,
        __import__("asyncio").Event(),
        _StubReqObj(stream=stream, rid=rid),
        created_time=0.0,
    )


def _register(manager: _StubTokenizerManager, state: ReqState) -> None:
    manager.rid_to_state[state.obj.rid] = state


# ---------------------------------------------------------------------------
# Flag-off regression: BatchTokenIDOut still takes the raw-token path.
# ---------------------------------------------------------------------------


class TestFlagOffRegression(unittest.TestCase):
    """Flag off → inline path stays dormant. We verify this at the receiver
    level: when a ``BatchTokenIDOut`` reaches ``_handle_batch_output`` with
    the flag off, the inline helper is never invoked and no
    ``inline_detokenizer`` is lazily created on the request state.

    (The pre-existing raw-token path for ``--skip-tokenizer-init`` requires
    ``recv_obj.output_ids`` to be populated by the scheduler; we don't
    exercise that path here — it isn't changed by this PR.)
    """

    def test_flag_off_receiver_does_not_take_inline_branch(self):
        tok = _gpt2_tokenizer()
        mgr = _StubTokenizerManager(tok, enable_inline_detokenizer=False)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        # Populate output_ids so the raw-token fallback path doesn't crash;
        # we only care that the inline branch is NOT taken.
        tokens = tok.encode("hello world")
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[tokens],
            output_ids=[tokens],
            batch_accept_draft_tokens=[1.5],
        )
        mgr.output_processor.handle_batch_output(recv)

        out = state.collector.take()
        self.assertIsNotNone(out)
        # The inline detokenizer does NOT run on this path (the assertion
        # this test exists for). What's emitted is the raw-token out_dict,
        # which since the D.1-regression hotfix carries an empty ``text``
        # key — matching the pre-D.1 BatchStrOut shape that subprocess
        # conversion used to guarantee. The state machine that would have
        # populated ``state.text`` never ran, so the value is "".
        self.assertEqual(out["text"], "")
        self.assertEqual(out["meta_info"]["accept_draft_tokens"], 1.5)
        self.assertIsNone(state.inline_detokenizer)
        self.assertEqual(state.text, "")


# ---------------------------------------------------------------------------
# Inline path basics.
# ---------------------------------------------------------------------------


class TestInlineBasicEmit(unittest.TestCase):
    """Flag-on path produces a BatchStrOut-shape out_dict with ``text``."""

    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def test_single_frame_populates_text_and_output_ids(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        source = "The quick brown fox"
        ids = self.tok.encode(source)
        # Emit all tokens as one finished frame so we can assert the
        # final text without partial-UTF-8 deferral buffering.
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            finished_reasons=[{"type": "stop", "matched": None}],
        )
        mgr.output_processor.handle_batch_output(recv)

        out = state.collector.take()
        self.assertIn("text", out)
        self.assertEqual(out["text"], source)
        self.assertEqual(out["output_ids"], ids)
        self.assertIs(out["meta_info"]["id"], "r1")
        # Inline detokenizer instantiated and still reachable on state.
        self.assertIsInstance(state.inline_detokenizer, IncrementalDetokenizer)

    def test_second_frame_reuses_per_request_detokenizer(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        first_ids = self.tok.encode("Hello ")
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(["r1"], decode_ids=[first_ids])
        )
        det_after_first = state.inline_detokenizer
        self.assertIsNotNone(det_after_first)

        second_ids = self.tok.encode("world")
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["r1"],
                decode_ids=[second_ids],
                finished_reasons=[{"type": "stop", "matched": None}],
            )
        )
        self.assertIs(state.inline_detokenizer, det_after_first)

        # Final state.text must carry the full cumulative decoded string.
        self.assertEqual(state.text, self.tok.decode(first_ids + second_ids))

    def test_meta_info_includes_finish_and_prompt_tokens(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        ids = self.tok.encode("end.")
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            finished_reasons=[{"type": "stop", "matched": None}],
            prompt_tokens=[7],
            completion_tokens=[len(ids)],
            cached_tokens=[0],
        )
        mgr.output_processor.handle_batch_output(recv)
        out = state.collector.take()
        self.assertEqual(out["meta_info"]["prompt_tokens"], 7)
        self.assertEqual(out["meta_info"]["completion_tokens"], len(ids))
        self.assertEqual(
            out["meta_info"]["finish_reason"], {"type": "stop", "matched": None}
        )
        self.assertTrue(state.finished)


# ---------------------------------------------------------------------------
# output_ids stream-vs-non-stream shape parity.
# ---------------------------------------------------------------------------


class TestOutputIdsShape(unittest.TestCase):
    """Inline path must match BatchStrOut branch's output_ids contract."""

    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def test_stream_mode_emits_delta_output_ids(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        ids_a = self.tok.encode("foo ")
        ids_b = self.tok.encode("bar")

        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(["r1"], decode_ids=[ids_a])
        )
        out1 = state.collector.take()
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["r1"],
                decode_ids=[ids_b],
                finished_reasons=[{"type": "stop", "matched": None}],
            )
        )
        out2 = state.collector.take()

        # Deltas: first frame is ids_a, second is ids_b.
        self.assertEqual(out1["output_ids"], ids_a)
        self.assertEqual(out2["output_ids"], ids_b)
        # Cumulative state carries the full list.
        self.assertEqual(state.output_ids, ids_a + ids_b)

    def test_non_stream_mode_emits_full_cumulative_output_ids(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=False, rid="r1")
        _register(mgr, state)

        ids_a = self.tok.encode("foo ")
        ids_b = self.tok.encode("bar")
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(["r1"], decode_ids=[ids_a])
        )
        out1 = state.collector.take()
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["r1"],
                decode_ids=[ids_b],
                finished_reasons=[{"type": "stop", "matched": None}],
            )
        )
        out2 = state.collector.take()

        # Full copy every frame in non-stream mode.
        self.assertEqual(out1["output_ids"], ids_a)
        self.assertEqual(out2["output_ids"], ids_a + ids_b)


# ---------------------------------------------------------------------------
# Stop trimming passes through unchanged.
# ---------------------------------------------------------------------------


class TestStopTrimmingPassThrough(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def test_matched_string_is_trimmed_in_text(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        source = "hello STOP world"
        ids = self.tok.encode(source)
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            finished_reasons=[{"type": "stop", "matched": "STOP"}],
        )
        mgr.output_processor.handle_batch_output(recv)
        out = state.collector.take()
        # Matched stop string and everything after must be trimmed.
        self.assertEqual(out["text"], "hello ")

    def test_no_stop_trim_flag_preserves_matched_content(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        source = "hello STOP world"
        ids = self.tok.encode(source)
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            finished_reasons=[{"type": "stop", "matched": "STOP"}],
            no_stop_trim=[True],
        )
        mgr.output_processor.handle_batch_output(recv)
        out = state.collector.take()
        self.assertEqual(out["text"], source)

    def test_matched_int_with_no_stop_trim_preserves_last_token(self):
        # Gap fill for ``matched=int`` (stop-token) case. The two places
        # ``trim_matched_stop`` fires inside the inline branch both take a
        # different code path for ``matched=int`` than for ``matched=str``:
        #   1. ``read_ids = trim_matched_stop(s.decode_ids[surr:], finish,
        #      no_stop_trim)`` — on the id-list side, matched=int drops the
        #      last token from ``read_ids`` before ``batch_decode``, which
        #      shortens the resulting text by whatever that token contributes.
        #   2. ``trim_matched_stop(s.decoded_text + new_text, ...)`` — on
        #      the string side, matched=int falls through and returns the
        #      output unchanged.
        # When ``no_stop_trim=True``, both short-circuit to "return output
        # unchanged", so the decoded text must include every token in
        # ``decode_ids``. This test locks that interaction on the inline
        # path so a later refactor cannot accidentally collapse the
        # matched=int case into the matched=str case.
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        source = "Hello world"
        ids = self.tok.encode(source)
        self.assertGreaterEqual(len(ids), 2)

        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            finished_reasons=[{"type": "stop", "matched": ids[-1]}],
            no_stop_trim=[True],
        )
        mgr.output_processor.handle_batch_output(recv)
        out = state.collector.take()
        # Full source preserved — matched=int + no_stop_trim means no
        # token is dropped from read_ids and no string trimming occurs.
        self.assertEqual(out["text"], self.tok.decode(ids))
        self.assertEqual(state.text, self.tok.decode(ids))


# ---------------------------------------------------------------------------
# Per-request independence.
# ---------------------------------------------------------------------------


class TestPerRequestIndependence(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def test_two_rids_share_one_manager_without_cross_talk(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        s_a = _mk_state(stream=True, rid="a")
        s_b = _mk_state(stream=True, rid="b")
        _register(mgr, s_a)
        _register(mgr, s_b)

        text_a = "apple pie"
        text_b = "blueberry muffin"
        ids_a = self.tok.encode(text_a)
        ids_b = self.tok.encode(text_b)

        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["a", "b"],
                decode_ids=[ids_a, ids_b],
                finished_reasons=[
                    {"type": "stop", "matched": None},
                    {"type": "stop", "matched": None},
                ],
            )
        )
        out_a = s_a.collector.take()
        out_b = s_b.collector.take()

        self.assertEqual(out_a["text"], text_a)
        self.assertEqual(out_b["text"], text_b)
        self.assertIsNot(s_a.inline_detokenizer, s_b.inline_detokenizer)


# ---------------------------------------------------------------------------
# Parity: inline text matches subprocess incremental_decode_batch output
# character-for-character for the same frame sequence.
# ---------------------------------------------------------------------------


def _run_subprocess_path(
    tokenizer: Any,
    rid: str,
    frames: List[Dict[str, Any]],
) -> List[str]:
    """Drive ``incremental_decode_batch`` with a single-request sequence.

    Returns the per-frame output_strs list.

    ``incremental_decode_batch`` aliases ``recv_obj.decode_ids[i]`` as the
    per-request ``DecodeStatus.decode_ids`` list and extends it in place on
    subsequent frames. Deep-copy each frame's ``decode_ids`` so the caller's
    frame fixture isn't corrupted across runs.
    """
    safe_frames = [
        {**frame, "decode_ids": list(frame["decode_ids"])} for frame in frames
    ]
    status: Dict[str, DecodeStatus] = {}
    pieces: List[str] = []
    for frame in safe_frames:
        recv = _batch_token_id_out(
            [rid],
            decode_ids=[frame["decode_ids"]],
            decoded_texts=[frame.get("decoded_text", "")],
            read_offsets=[frame.get("read_offset", 0)],
            finished_reasons=[frame.get("finished_reason")],
            no_stop_trim=[frame.get("no_stop_trim", False)],
        )
        out_strs = incremental_decode_batch(tokenizer, status, recv)
        pieces.append(out_strs[0])
    return pieces


def _run_inline_path(
    tokenizer: Any,
    rid: str,
    frames: List[Dict[str, Any]],
) -> List[str]:
    """Drive the inline receiver for a single request and collect the
    incremental text fragments that would have appeared in the stream.

    We read ``state.text`` before and after each frame and use the delta.
    This matches what an OpenAI streaming client would observe.

    Same deep-copy dance as ``_run_subprocess_path``: the state machine
    aliases the frame's ``decode_ids`` list on the first frame and extends
    it in place afterward, so shield the caller's fixture from mutation.
    """
    safe_frames = [
        {**frame, "decode_ids": list(frame["decode_ids"])} for frame in frames
    ]
    mgr = _StubTokenizerManager(tokenizer, enable_inline_detokenizer=True)
    state = _mk_state(stream=True, rid=rid)
    _register(mgr, state)

    pieces: List[str] = []
    prev_text = ""
    for idx, frame in enumerate(safe_frames):
        seed_text = frame.get("decoded_text", "") if idx == 0 else ""
        seed_offset = frame.get("read_offset", 0) if idx == 0 else 0
        recv = _batch_token_id_out(
            [rid],
            decode_ids=[frame["decode_ids"]],
            decoded_texts=[seed_text],
            read_offsets=[seed_offset],
            finished_reasons=[frame.get("finished_reason")],
            no_stop_trim=[frame.get("no_stop_trim", False)],
        )
        mgr.output_processor.handle_batch_output(recv)
        state.collector.take()  # drain
        pieces.append(state.text[len(prev_text) :])
        prev_text = state.text
    return pieces


class TestSubprocessVsInlineParity(unittest.TestCase):
    """For identical frame sequences the two paths must emit identical text.

    The per-frame emits are compared byte-for-byte; the cumulative text is
    compared too. Every drift between inline and subprocess behavior would
    surface here.
    """

    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def _assert_parity(self, frames: List[Dict[str, Any]]) -> None:
        subp = _run_subprocess_path(self.tok, "r1", frames)
        inline = _run_inline_path(self.tok, "r1", frames)
        self.assertEqual(
            inline,
            subp,
            msg=f"inline={inline!r} subprocess={subp!r}",
        )

    def test_parity_single_frame_ascii(self):
        self._assert_parity(
            [
                {
                    "decode_ids": self.tok.encode("Hello, world!"),
                    "finished_reason": {"type": "stop", "matched": None},
                }
            ]
        )

    def test_parity_per_token_ascii(self):
        source = "streaming tokens one at a time"
        ids = self.tok.encode(source)
        frames = [{"decode_ids": [tid]} for tid in ids[:-1]]
        frames.append(
            {
                "decode_ids": [ids[-1]],
                "finished_reason": {"type": "stop", "matched": None},
            }
        )
        self._assert_parity(frames)

    def test_parity_two_frame_split(self):
        source = "hello there friend"
        ids = self.tok.encode(source)
        mid = len(ids) // 2
        self._assert_parity(
            [
                {"decode_ids": ids[:mid]},
                {
                    "decode_ids": ids[mid:],
                    "finished_reason": {"type": "stop", "matched": None},
                },
            ]
        )

    def test_parity_cjk_per_token(self):
        # CJK exercises partial-UTF-8 deferral + find_printable_text.
        source = "你好世界"
        ids = self.tok.encode(source)
        frames = [{"decode_ids": [tid]} for tid in ids[:-1]]
        frames.append(
            {
                "decode_ids": [ids[-1]],
                "finished_reason": {"type": "stop", "matched": None},
            }
        )
        self._assert_parity(frames)

    def test_parity_finish_with_matched_stop_string(self):
        source = "keep this STOP drop this"
        ids = self.tok.encode(source)
        self._assert_parity(
            [
                {
                    "decode_ids": ids,
                    "finished_reason": {"type": "stop", "matched": "STOP"},
                }
            ]
        )

    def test_parity_unfinished_streaming_does_not_trim(self):
        # Streaming (no finish yet) must NOT apply stop trimming.
        source = "prefix STOP more"
        ids = self.tok.encode(source)
        self._assert_parity(
            [
                {"decode_ids": ids[:3]},
                {"decode_ids": ids[3:]},  # no finish — still streaming
            ]
        )

    def test_parity_emoji_per_token(self):
        # 4-byte UTF-8 emoji per-token streaming. Exercises a different
        # partial-byte shape than CJK: the emoji codepoint is split
        # across ~4 byte-level BPE tokens (one per UTF-8 byte) and
        # ``find_printable_text`` has to defer every intermediate frame
        # until the full codepoint arrives. Any drift between the inline
        # path's offset bookkeeping and the subprocess path's surfaces
        # here because the defer-then-commit timing has to match
        # byte-for-byte.
        source = "a🌟b"
        ids = self.tok.encode(source)
        frames: List[Dict[str, Any]] = [{"decode_ids": [tid]} for tid in ids[:-1]]
        frames.append(
            {
                "decode_ids": [ids[-1]],
                "finished_reason": {"type": "stop", "matched": None},
            }
        )
        self._assert_parity(frames)

    def test_parity_finish_on_later_frame_with_prior_unfinished_commits(self):
        # The highest-risk state-machine branch: finish arrives on a
        # non-first frame after one or more unfinished commits have
        # already landed. On the finished frame the commit block is
        # skipped and ``trim_matched_stop`` runs on
        # ``s.decoded_text + new_text`` using the offsets accumulated
        # from prior commits. This mirrors the batch detokenizer coverage in
        # ``test_finish_arrives_on_later_frame_applies_trim_at_finish_only``
        # against the inline receiver path so any inline-specific drift shows up.
        source = "Hello world the long sentence STOP tail text"
        ids = self.tok.encode(source)
        self.assertGreaterEqual(len(ids), 6)

        # Pick the largest split whose decoded prefix is still
        # STOP-free so frame 1 commits a clean prefix.
        split: Optional[int] = None
        for candidate in range(len(ids) - 1, 0, -1):
            if "STOP" not in self.tok.decode(ids[:candidate]):
                split = candidate
                break
        self.assertIsNotNone(split, "need a STOP-free prefix split")

        self._assert_parity(
            [
                {"decode_ids": ids[:split]},  # unfinished commit
                {
                    "decode_ids": ids[split:],
                    "finished_reason": {"type": "stop", "matched": "STOP"},
                },
            ]
        )


# ---------------------------------------------------------------------------
# Seed handling from BatchTokenIDOut.decoded_texts / read_offsets.
# ---------------------------------------------------------------------------


class TestSeedHandling(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def test_first_frame_honors_decoded_text_seed(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        seed = "<resume> "
        ids = self.tok.encode("continuation")
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[ids],
            decoded_texts=[seed],
            finished_reasons=[{"type": "stop", "matched": None}],
        )
        mgr.output_processor.handle_batch_output(recv)
        out = state.collector.take()
        self.assertEqual(out["text"], seed + self.tok.decode(ids))

    def test_only_first_frame_seeds_detokenizer(self):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_state(stream=True, rid="r1")
        _register(mgr, state)

        ids_first = self.tok.encode("A")
        ids_second = self.tok.encode("B")

        # First frame with seed.
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["r1"],
                decode_ids=[ids_first],
                decoded_texts=["seed-"],
                read_offsets=[0],
            )
        )
        state.collector.take()

        # Second frame passes a misleading decoded_texts which MUST be ignored
        # because the detokenizer is already initialized.
        mgr.output_processor.handle_batch_output(
            _batch_token_id_out(
                ["r1"],
                decode_ids=[ids_second],
                decoded_texts=["ignored-"],
                read_offsets=[999],
                finished_reasons=[{"type": "stop", "matched": None}],
            )
        )
        state.collector.take()
        self.assertEqual(state.text, "seed-" + self.tok.decode(ids_first + ids_second))


# ---------------------------------------------------------------------------
# Logprob / meta field pass-through through the inline branch.
# ---------------------------------------------------------------------------


class _LogprobReqObj(_StubReqObj):
    """Request object that asks for output logprobs in a given dialect so
    ``_handle_batch_output`` invokes ``convert_logprob_style`` on the recv_obj.

    ``fmt="vllm"`` sets ``sampling_params["logprobs"]=0`` (the output processor
    renders the vLLM ``logprobs`` dict). ``fmt="sglang"`` leaves sampling_params
    without ``logprobs`` so the SGLang ``return_logprob`` flag drives the
    tuple-list rendering.
    """

    def __init__(
        self,
        *,
        rid: str = "r1",
        fmt: str = "vllm",
    ) -> None:
        super().__init__(stream=True, return_logprob=True, rid=rid)
        self.sampling_params = {"logprobs": 0} if fmt == "vllm" else {}
        self.logprob_format = None  # auto: match the request dialect


def _mk_logprob_state(*, rid: str = "r1", fmt: str = "vllm") -> ReqState:
    return ReqState(
        RequestOutputCollector(),
        False,
        __import__("asyncio").Event(),
        _LogprobReqObj(rid=rid, fmt=fmt),
        created_time=0.0,
    )


class TestInlineLogprobPassThrough(unittest.TestCase):
    """Verify the sampled-token output logprob arrays on a ``BatchTokenIDOut``
    flow through the inline branch of ``_handle_batch_output`` into ``meta_info``
    — in BOTH dialects, selected per request, from the same wire arrays.
    """

    @classmethod
    def setUpClass(cls) -> None:
        cls.tok = _gpt2_tokenizer()

    def _run(self, fmt):
        mgr = _StubTokenizerManager(self.tok, enable_inline_detokenizer=True)
        state = _mk_logprob_state(rid="r1", fmt=fmt)
        _register(mgr, state)
        recv = _batch_token_id_out(
            ["r1"],
            decode_ids=[self.tok.encode("Hello")],
            finished_reasons=[{"type": "stop", "matched": None}],
            input_token_logprobs_val=[[]],
            input_token_logprobs_idx=[[]],
            output_token_logprobs_val=[[-0.5, -0.6]],
            output_token_logprobs_idx=[[1, 2]],
        )
        mgr.output_processor.handle_batch_output(recv)
        return state.collector.take()["meta_info"]

    def test_vllm_format(self):
        meta = self._run("vllm")
        # vLLM shape: "logprobs" is a list[dict[int, Logprob]], one per token.
        self.assertIn("logprobs", meta)
        self.assertNotIn("output_token_logprobs", meta)
        self.assertEqual(len(meta["logprobs"]), 2)
        self.assertEqual(meta["logprobs"][0][1].logprob, -0.5)
        self.assertEqual(meta["logprobs"][0][1].rank, 0)
        self.assertEqual(meta["logprobs"][1][2].logprob, -0.6)
        self.assertAlmostEqual(meta["cumulative_logprob"], -1.1)

    def test_sglang_format(self):
        meta = self._run("sglang")
        # SGLang shape: "output_token_logprobs" is a list of (val, idx, text)
        # tuples (text None when not decoding); no vLLM "logprobs" key.
        self.assertIn("output_token_logprobs", meta)
        self.assertNotIn("logprobs", meta)
        self.assertEqual([e[0] for e in meta["output_token_logprobs"]], [-0.5, -0.6])
        self.assertEqual([e[1] for e in meta["output_token_logprobs"]], [1, 2])


if __name__ == "__main__":
    unittest.main(verbosity=2)