Step audio R1 reasoning parser (#2846)

Signed-off-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: Mike Qiu <qdy220091330@gmail.com>
Co-authored-by: Mike_Qiu <qiudayu.qdy@antgroup.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

pyproject.toml

@@ -292,5 +292,9 @@ fullset = "fullset"
 Vai = "Vai"
 tockens = "tockens"
 CANN = "CANN"
+# Step-Audio reasoning parser streams `</think>` tokens split across chunks
+# (e.g. "</" + "thi" + "nk>"). The fragment "thi" appears literally in its
+# parser logic and tests and is not a typo of "the"/"this".
+thi = "thi"
 ND = "ND"
 byt5 = "byt5"

tests/reasoning/test_step_audio_reasoning_parser.py

@@ -0,0 +1,918 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Tests for StepAudioReasoningParser.
+
+Step-Audio supports two representations of thinking markers:
+
+1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>``
+   (single-token IDs in vocab).
+
+2. **Text markers**: ``<think>`` and ``</think>``
+   (multi-token sequences for </think>, e.g. [522, 26865]).
+
+These tests use standalone mock objects so they can run without
+installing the full vLLM/torch stack.
+"""
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Bootstrap minimal vLLM stubs so the parser can be imported
+# ---------------------------------------------------------------------------
+
+
+def _install_vllm_stubs() -> dict[str, ModuleType | None]:
+    """Temporarily install minimal stub modules for the vllm imports the
+    parser needs, so it can be imported without the full vLLM/torch stack.
+
+    Returns a snapshot of the original ``sys.modules`` entries that were
+    replaced. The caller MUST pass it to ``_restore_vllm_stubs`` once the
+    parser is imported. These stubs are process-global: if they leak, they
+    shadow the real vllm modules for every other test in the same worker
+    (and any forked server subprocess), e.g. replacing
+    ``vllm.entrypoints.openai.engine.protocol`` with a stub that lacks
+    ``ErrorResponse`` and breaking unrelated server-startup tests.
+    """
+    stubs: dict[str, ModuleType] = {}
+
+    # DeltaMessage stub
+    delta_mod = ModuleType("vllm.entrypoints.openai.engine.protocol")
+
+    class DeltaMessage:
+        __slots__ = ("reasoning", "content")
+
+        def __init__(self, reasoning=None, content=None):
+            self.reasoning = reasoning
+            self.content = content
+
+    delta_mod.DeltaMessage = DeltaMessage
+    stubs["vllm.entrypoints.openai.engine.protocol"] = delta_mod
+
+    # ReasoningParser stub
+    reasoning_mod = ModuleType("vllm.reasoning")
+
+    class ReasoningParser:
+        def __init__(self, tokenizer, *args, **kwargs):
+            self.model_tokenizer = tokenizer
+
+        @property
+        def vocab(self):
+            if not self.model_tokenizer:
+                return {}
+            return self.model_tokenizer.get_vocab()
+
+    class _ReasoningParserManager:
+        reasoning_parsers = {}
+        lazy_parsers = {}
+
+        @classmethod
+        def register_lazy_module(cls, name, module_path, class_name):
+            cls.lazy_parsers[name] = (module_path, class_name)
+
+    reasoning_mod.ReasoningParser = ReasoningParser
+    reasoning_mod.ReasoningParserManager = _ReasoningParserManager
+    stubs["vllm.reasoning"] = reasoning_mod
+    stubs["vllm.reasoning.abs_reasoning_parsers"] = reasoning_mod
+
+    # TokenizerLike stub
+    tokenizers_mod = ModuleType("vllm.tokenizers")
+    tokenizers_mod.TokenizerLike = object
+    stubs["vllm.tokenizers"] = tokenizers_mod
+
+    # ChatCompletionRequest / ResponsesRequest stubs
+    chat_mod = ModuleType("vllm.entrypoints.openai.chat_completion.protocol")
+    chat_mod.ChatCompletionRequest = MagicMock
+    stubs["vllm.entrypoints.openai.chat_completion.protocol"] = chat_mod
+    responses_mod = ModuleType("vllm.entrypoints.openai.responses.protocol")
+    responses_mod.ResponsesRequest = MagicMock
+    stubs["vllm.entrypoints.openai.responses.protocol"] = responses_mod
+
+    # Parent packages: only stub the ones not already importable, so we never
+    # clobber a real (installed) vllm package tree.
+    for mod_path in [
+        "vllm",
+        "vllm.entrypoints",
+        "vllm.entrypoints.openai",
+        "vllm.entrypoints.openai.engine",
+        "vllm.entrypoints.openai.chat_completion",
+        "vllm.entrypoints.openai.responses",
+    ]:
+        if mod_path not in sys.modules and mod_path not in stubs:
+            stubs[mod_path] = ModuleType(mod_path)
+
+    saved = {name: sys.modules.get(name) for name in stubs}
+    sys.modules.update(stubs)
+    return saved
+
+
+def _restore_vllm_stubs(saved: dict[str, ModuleType | None]) -> None:
+    """Undo ``_install_vllm_stubs`` so ``sys.modules`` is left untouched."""
+    for name, original in saved.items():
+        if original is None:
+            sys.modules.pop(name, None)
+        else:
+            sys.modules[name] = original
+
+
+# Install stubs only for the duration of the import. The parser binds these
+# stub classes into its own module namespace at import time, so it keeps using
+# them at runtime even after sys.modules is restored below.
+_saved_modules = _install_vllm_stubs()
+try:
+    from vllm_omni.reasoning.step_audio_reasoning_parser import (  # noqa: E402
+        StepAudioReasoningParser,
+    )
+finally:
+    _restore_vllm_stubs(_saved_modules)
+
+# Short aliases for marker constants
+TS = StepAudioReasoningParser.THINK_START_TEXT  # <think>
+TE = StepAudioReasoningParser.THINK_END_TEXT  # </think>
+TSS = StepAudioReasoningParser.THINK_START_SPECIAL  # <|THINK_START|>
+TES = StepAudioReasoningParser.THINK_END_SPECIAL  # <|THINK_END|>
+
+
+# ---------------------------------------------------------------------------
+# Mock tokenizers
+# ---------------------------------------------------------------------------
+
+
+class SingleTokenTokenizer:
+    """Tokenizer where both ``<think>`` and ``</think>`` are single tokens.
+
+    This exercises the fast (token-ID) path.
+    Also includes special-token forms ``<|THINK_START|>`` and ``<|THINK_END|>``.
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think>
+            TE: 101,  # </think>
+            TSS: 151669,  # <|THINK_START|>
+            TES: 151670,  # <|THINK_END|>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "</": 6,
+            "think>": 7,
+            "some text": 8,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+class MultiTokenTokenizer:
+    """Tokenizer where ``</think>`` is NOT a single token.
+
+    ``<think>`` is in the vocab (ID 100), but ``</think>`` is NOT.
+    Instead, ``</think>`` encodes as ["</", "think>"] → [522, 26865].
+
+    Both special tokens ``<|THINK_START|>`` and ``<|THINK_END|>`` ARE
+    present as single tokens.
+
+    This exercises the text-based path (with special-token fallback).
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think> - single token
+            # </think> is intentionally NOT in the vocab
+            "</": 522,
+            "think>": 26865,
+            TSS: 151669,  # <|THINK_START|>
+            TES: 151670,  # <|THINK_END|>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "some text": 8,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+class TextOnlyTokenizer:
+    """Tokenizer with only text-form markers (no special tokens at all).
+
+    Both ``<think>`` and ``</think>`` are single tokens.
+    ``<|THINK_START|>`` and ``<|THINK_END|>`` are NOT in the vocab.
+
+    This tests the pure text-based fallback paths.
+    """
+
+    def __init__(self):
+        self._vocab = {
+            TS: 100,  # <think>
+            TE: 101,  # </think>
+            "Hello": 1,
+            " world": 2,
+            "reasoning": 3,
+            " content": 4,
+            "\n": 5,
+            "</": 6,
+            "think>": 7,
+        }
+        self._reverse = {v: k for k, v in self._vocab.items()}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self._reverse.get(i, "") for i in ids)
+
+    def tokenize(self, text):
+        if not text:
+            return []
+        return [text]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _stream_extract(parser, deltas_and_ids):
+    """Run streaming extraction over a list of (delta_text, delta_ids) tuples.
+    Returns (reasoning_text, content_text)."""
+    prev_text = ""
+    prev_ids = []
+    all_r = []
+    all_c = []
+
+    for delta_text, delta_ids in deltas_and_ids:
+        cur_text = prev_text + delta_text
+        cur_ids = prev_ids + delta_ids
+        result = parser.extract_reasoning_streaming(prev_text, cur_text, delta_text, prev_ids, cur_ids, delta_ids)
+        if result is not None:
+            if result.reasoning:
+                all_r.append(result.reasoning)
+            if result.content:
+                all_c.append(result.content)
+        prev_text = cur_text
+        prev_ids = cur_ids
+
+    return "".join(all_r), "".join(all_c)
+
+
+# ---------------------------------------------------------------------------
+# Fixtures
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def single_token_parser():
+    return StepAudioReasoningParser(SingleTokenTokenizer())
+
+
+@pytest.fixture
+def multi_token_parser():
+    return StepAudioReasoningParser(MultiTokenTokenizer())
+
+
+@pytest.fixture
+def text_only_parser():
+    return StepAudioReasoningParser(TextOnlyTokenizer())
+
+
+@pytest.fixture
+def request_obj():
+    return MagicMock(messages=[], model="test-model")
+
+
+# ---------------------------------------------------------------------------
+# Initialization tests
+# ---------------------------------------------------------------------------
+
+
+class TestInit:
+    def test_single_token_ids(self, single_token_parser):
+        """Text-form markers are in vocab as single tokens."""
+        assert single_token_parser.think_start_text_id == 100
+        assert single_token_parser.think_end_text_id == 101
+
+    def test_single_token_special_ids(self, single_token_parser):
+        """Special-token markers are also in vocab."""
+        assert single_token_parser.think_start_special_id == 151669
+        assert single_token_parser.think_end_special_id == 151670
+
+    def test_single_token_has_ids(self, single_token_parser):
+        assert single_token_parser._has_start_token_id is True
+        assert single_token_parser._has_end_token_id is True
+
+    def test_multi_token_text_ids(self, multi_token_parser):
+        """<think> is a single token; </think> is NOT."""
+        assert multi_token_parser.think_start_text_id == 100
+        assert multi_token_parser.think_end_text_id == -1
+
+    def test_multi_token_special_ids(self, multi_token_parser):
+        """Special tokens are still in vocab."""
+        assert multi_token_parser.think_start_special_id == 151669
+        assert multi_token_parser.think_end_special_id == 151670
+
+    def test_multi_token_has_ids(self, multi_token_parser):
+        """Even without </think> as single token, special tokens provide IDs."""
+        assert multi_token_parser._has_start_token_id is True
+        assert multi_token_parser._has_end_token_id is True
+
+    def test_text_only_ids(self, text_only_parser):
+        """Only text-form markers, no special tokens."""
+        assert text_only_parser.think_start_text_id == 100
+        assert text_only_parser.think_end_text_id == 101
+        assert text_only_parser.think_start_special_id == -1
+        assert text_only_parser.think_end_special_id == -1
+        assert text_only_parser._has_start_token_id is True
+        assert text_only_parser._has_end_token_id is True
+
+    def test_missing_tokenizer_raises(self):
+        with pytest.raises(ValueError, match="model tokenizer must be passed"):
+            StepAudioReasoningParser(None)
+
+
+# ---------------------------------------------------------------------------
+# extract_reasoning (non-streaming)
+# ---------------------------------------------------------------------------
+
+
+class TestExtractReasoning:
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_basic_extraction(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"some reasoning{TE}Hello world"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "some reasoning"
+        assert content == "Hello world"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_no_end_token(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = "still thinking..."
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "still thinking..."
+        assert content is None
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_with_start_token_prefix(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TS}my reasoning{TE}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "my reasoning"
+        assert content == "answer"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_strips_newline_after_end(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"reasoning{TE}\nHello"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "Hello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_empty_reasoning(self, request, parser_name, request_obj):
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TE}Hello"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning is None
+        assert content == "Hello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_special_end_token(self, request, parser_name, request_obj):
+        """Test with <|THINK_END|> special token form."""
+        parser = request.getfixturevalue(parser_name)
+        output = f"reasoning{TES}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "answer"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_special_start_token(self, request, parser_name, request_obj):
+        """Test with <|THINK_START|> special token form."""
+        parser = request.getfixturevalue(parser_name)
+        output = f"{TSS}reasoning{TE}answer"
+        reasoning, content = parser.extract_reasoning(output, request_obj)
+        assert reasoning == "reasoning"
+        assert content == "answer"
+
+
+# ---------------------------------------------------------------------------
+# is_reasoning_end
+# ---------------------------------------------------------------------------
+
+
+class TestIsReasoningEnd:
+    def test_single_token_text_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 101]) is True
+
+    def test_single_token_special_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 151670]) is True
+
+    def test_single_token_no_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_multi_token_text_end_decoded(self, multi_token_parser):
+        """</think> = [522, 26865] detected via decode."""
+        assert multi_token_parser.is_reasoning_end([1, 522, 26865]) is True
+
+    def test_multi_token_special_end(self, multi_token_parser):
+        """<|THINK_END|> detected as single token ID."""
+        assert multi_token_parser.is_reasoning_end([1, 151670]) is True
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_multi_token_partial_marker(self, multi_token_parser):
+        """Only </ without think> should NOT count."""
+        assert multi_token_parser.is_reasoning_end([1, 522, 3]) is False
+
+    def test_multi_token_empty_input(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end([]) is False
+
+    def test_text_only_end(self, text_only_parser):
+        assert text_only_parser.is_reasoning_end([1, 2, 101]) is True
+
+    def test_text_only_no_end(self, text_only_parser):
+        assert text_only_parser.is_reasoning_end([1, 2, 3]) is False
+
+    def test_start_after_end_returns_false(self, single_token_parser):
+        """Multi-turn prompt: start marker AFTER end marker → reasoning
+        is still active (the generation prompt's start marker)."""
+        # IDs: [start=100, end=101, start=100] → last marker is start
+        assert single_token_parser.is_reasoning_end([100, 101, 100]) is False
+
+    def test_end_is_last_marker_returns_true(self, single_token_parser):
+        """End marker is the last marker → reasoning has ended."""
+        # IDs: [start=100, end=101] → last marker is end
+        assert single_token_parser.is_reasoning_end([100, 101]) is True
+
+    def test_start_only_no_end(self, single_token_parser):
+        """Only a start marker, no end marker → reasoning is active."""
+        assert single_token_parser.is_reasoning_end([100, 1, 2]) is False
+
+    def test_multi_turn_prompt_start_after_end(self, multi_token_parser):
+        """Simulated multi-turn prompt with previous assistant response.
+
+        Prompt contains: ... </think ... <think ...
+        The last marker is <think (start), so is_reasoning_end
+        should return False.
+        """
+        # Decode of these IDs: "reasoning" +  "</ " + "think>" + " " + "<think"
+        # The last marker is the start marker.
+        # Using: [3=reasoning, 522="</", 26865="think>", 100=TS]
+        ids = [3, 522, 26865, 100]
+        assert multi_token_parser.is_reasoning_end(ids) is False
+
+    def test_multi_turn_prompt_end_is_last(self, multi_token_parser):
+        """Previous turn ends with </think, no new start → ended."""
+        # [3=reasoning, 522="</", 26865="think>"]
+        ids = [3, 522, 26865]
+        assert multi_token_parser.is_reasoning_end(ids) is True
+
+    def test_special_start_after_end(self, multi_token_parser):
+        """Special token start after end marker → reasoning still active."""
+        # [151670=<|THINK_END|>, 151669=<|THINK_START|>]
+        ids = [151670, 151669]
+        assert multi_token_parser.is_reasoning_end(ids) is False
+
+
+# ---------------------------------------------------------------------------
+# is_reasoning_end_streaming
+# ---------------------------------------------------------------------------
+
+
+class TestIsReasoningEndStreaming:
+    def test_single_token_delta_has_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end_streaming([1, 2], [101]) is True
+
+    def test_single_token_delta_no_end(self, single_token_parser):
+        assert single_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
+
+    def test_multi_token_delta_has_both_parts(self, multi_token_parser):
+        """Full </think> in delta → [522, 26865]."""
+        assert multi_token_parser.is_reasoning_end_streaming([1], [522, 26865]) is True
+
+    def test_multi_token_marker_spans_boundary(self, multi_token_parser):
+        """</ in previous, think> in delta."""
+        assert multi_token_parser.is_reasoning_end_streaming([1, 522], [26865]) is True
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        assert multi_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
+
+
+# ---------------------------------------------------------------------------
+# extract_reasoning_streaming
+# ---------------------------------------------------------------------------
+
+
+class TestExtractReasoningStreaming:
+    """Test streaming extraction for both single-token and multi-token modes."""
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_simple_streaming(self, request_, parser_name):
+        parser = request_.getfixturevalue(parser_name)
+        # Simulate: <think> reasoning </think>  final
+        if parser_name == "multi_token_parser":
+            deltas = [TS, "reasoning", TE, " content"]
+            token_ids_per_delta = [
+                [100],  # <think>
+                [3],  # reasoning
+                [522, 26865],  # </think>
+                [4],  #  content
+            ]
+        else:
+            deltas = [TS, "reasoning", TE, " content"]
+            token_ids_per_delta = [
+                [100],  # <think>
+                [3],  # reasoning
+                [101],  # </think>
+                [4],  #  content
+            ]
+
+        r, c = _stream_extract(parser, list(zip(deltas, token_ids_per_delta)))
+        assert r == "reasoning"
+        assert c == " content"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+            "text_only_parser",
+        ],
+    )
+    def test_no_end_token_streaming(self, request_, parser_name):
+        parser = request_.getfixturevalue(parser_name)
+        deltas_and_ids = [
+            (TS, [100]),
+            ("thinking", [3]),
+            (" more", [2]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "thinking more"
+        assert c == ""
+
+    def test_multi_token_end_split_across_chunks(self, multi_token_parser):
+        """The key test: </think> is split across streaming chunks.
+
+        Previous chunk ends with "reasoning</" and delta starts with
+        "think>\\nfinal answer".  This is the real-world scenario when
+        ``</think>`` is tokenized as two separate tokens.
+        """
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            ("</", [522]),
+            ("think>\nHello", [26865, 1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_multi_token_start_token_in_delta(self, multi_token_parser):
+        """Test that <think> start token is skipped in streaming output."""
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+            ("Hello", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_multi_token_three_way_split(self, multi_token_parser):
+        """</think> split across 3 chunks: </ + thi + nk>."""
+        deltas_and_ids = [
+            ("some reasoning", [3]),
+            ("</", [522]),
+            ("thi", [8]),
+            ("nk>", [8]),
+            ("\nHello", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "some reasoning"
+        assert c == "Hello"
+
+    def test_special_end_token_in_stream(self, multi_token_parser):
+        """<|THINK_END|> special token in streaming (decoded as text)."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TES, [151670]),
+            ("answer", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_no_start_token_in_stream(self, multi_token_parser):
+        """Model starts directly with reasoning (no <think> prefix)."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+            ("answer", [1]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_empty_reasoning(self, single_token_parser):
+        """Start immediately followed by end → empty reasoning."""
+        deltas_and_ids = [
+            (TS, [100]),
+            (TE, [101]),
+            ("Hello", [1]),
+        ]
+        r, c = _stream_extract(single_token_parser, deltas_and_ids)
+        assert r in ("", None)
+        assert c == "Hello"
+
+    def test_end_marker_no_content_after(self, multi_token_parser):
+        """End marker with no content following it."""
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            (TE, [522, 26865]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == ""
+
+    def test_delta_is_just_start_token(self, multi_token_parser):
+        """Delta consisting only of the start token should be skipped."""
+        deltas_and_ids = [
+            (TS, [100]),
+        ]
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == ""
+        assert c == ""
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_only_first_newline_stripped_after_end(self, request_, parser_name):
+        """Only the first newline after </think> should be stripped."""
+        parser = request_.getfixturevalue(parser_name)
+        if parser_name == "multi_token_parser":
+            deltas_and_ids = [
+                ("reasoning", [3]),
+                (TE, [522, 26865]),
+                ("\n\nHello", [5, 1]),
+            ]
+        else:
+            deltas_and_ids = [
+                ("reasoning", [3]),
+                (TE, [101]),
+                ("\n\nHello", [5, 1]),
+            ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "\nHello"
+
+    @pytest.mark.parametrize(
+        "parser_name",
+        [
+            "single_token_parser",
+            "multi_token_parser",
+        ],
+    )
+    def test_content_after_end_multiple_deltas(self, request_, parser_name):
+        """Multiple content deltas after end marker."""
+        parser = request_.getfixturevalue(parser_name)
+        if parser_name == "multi_token_parser":
+            deltas_and_ids = [
+                (TS, [100]),
+                ("reasoning", [3]),
+                (TE, [522, 26865]),
+                ("Hello", [1]),
+                (" world", [2]),
+            ]
+        else:
+            deltas_and_ids = [
+                (TS, [100]),
+                ("reasoning", [3]),
+                (TE, [101]),
+                ("Hello", [1]),
+                (" world", [2]),
+            ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello world"
+
+    def test_pending_buffer_resolved_on_non_marker(self, multi_token_parser):
+        """When buffered text turns out not to be a marker prefix, it is
+        emitted as reasoning on the next call."""
+        # "reason<" is buffered because "<" could start a marker.
+        # Next delta ">" doesn't form a marker, so the buffer is
+        # flushed as reasoning.
+        parser = multi_token_parser
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reason", [3]),
+            ("<", [8]),  # "<" is buffered (prefix of markers)
+            ("> more", [8]),  # ">" doesn't complete any marker
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reason> more"
+        assert c == ""
+
+    def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser):
+        """Reasoning text before a buffered prefix should be emitted."""
+        # "reasoning<" → "reasoning" emitted, "<" buffered
+        # "</think>\nanswer" → buffer resolves as end marker
+        parser = multi_token_parser
+        deltas_and_ids = [
+            ("reasoning", [3]),
+            ("<", [8]),
+            (f"/{TE[2:]}\nanswer", [522, 26865, 1]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "answer"
+
+    def test_multi_token_five_way_split(self, multi_token_parser):
+        """</think> split across 5 tiny chunks."""
+        # Split: each chunk is one character of </think>
+        chunks = list(TE)
+        deltas_and_ids = [
+            ("reasoning", [3]),
+        ]
+        for ch in chunks:
+            deltas_and_ids.append((ch, [8]))
+        deltas_and_ids.append(("\nHello", [1]))
+        r, c = _stream_extract(multi_token_parser, deltas_and_ids)
+        assert r == "reasoning"
+        assert c == "Hello"
+
+    def test_single_token_parser_with_pending(self, single_token_parser):
+        """Single-token parser: ambiguous '<' prefix buffered then resolved."""
+        # Even with single-token markers, the text-based path buffers
+        # '<' because it could be a prefix of a marker.
+        parser = single_token_parser
+        deltas_and_ids = [
+            (TS, [100]),
+            ("reason", [3]),
+            ("<", [8]),
+            ("> not a marker", [8]),
+        ]
+        r, c = _stream_extract(parser, deltas_and_ids)
+        assert r == "reason> not a marker"
+        assert c == ""
+
+
+# ---------------------------------------------------------------------------
+
+
+class TestCountReasoningTokens:
+    def test_single_token_basic(self, single_token_parser):
+        # IDs: [100=<think>, 3=reasoning, 101=</think>, 4=content]
+        assert single_token_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
+
+    def test_single_token_no_end(self, single_token_parser):
+        # All tokens after start are reasoning
+        assert single_token_parser.count_reasoning_tokens([100, 3, 2]) == 2
+
+    def test_multi_token_basic(self, multi_token_parser):
+        # </think> encodes as [522, 26865]
+        # Full sequence: <think> reasoning </think> content
+        # IDs: [100, 3, 522, 26865, 4]
+        result = multi_token_parser.count_reasoning_tokens([100, 3, 522, 26865, 4])
+        assert result > 0
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        # No end token → all tokens inside <think> are counted.
+        result = multi_token_parser.count_reasoning_tokens([100, 3, 2])
+        assert result > 0
+
+    def test_text_only_basic(self, text_only_parser):
+        assert text_only_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
+
+    def test_text_only_no_end(self, text_only_parser):
+        assert text_only_parser.count_reasoning_tokens([100, 3, 2]) == 2
+
+
+# ---------------------------------------------------------------------------
+# extract_content_ids
+# ---------------------------------------------------------------------------
+
+
+class TestExtractContentIds:
+    def test_single_token_basic(self, single_token_parser):
+        # [1, 101=</think>, 2, 3] → content starts after ID 101
+        result = single_token_parser.extract_content_ids([1, 101, 2, 3])
+        assert result == [2, 3]
+
+    def test_single_token_no_end(self, single_token_parser):
+        result = single_token_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+    def test_multi_token_basic(self, multi_token_parser):
+        # </think> = [522, 26865]
+        # [3, 522, 26865, 4] → content after </think>
+        result = multi_token_parser.extract_content_ids([3, 522, 26865, 4])
+        assert 4 in result
+
+    def test_multi_token_no_end(self, multi_token_parser):
+        result = multi_token_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+    def test_text_only_basic(self, text_only_parser):
+        result = text_only_parser.extract_content_ids([1, 101, 2, 3])
+        assert result == [2, 3]
+
+    def test_text_only_no_end(self, text_only_parser):
+        result = text_only_parser.extract_content_ids([1, 2, 3])
+        assert result == []
+
+
+# ---------------------------------------------------------------------------
+# Fixture request workaround (pytest "request" fixture shadowed by parametrize)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def request_(request):
+    """Proxy for the pytest ``request`` fixture to avoid name collisions
+    with the ``request_obj`` fixture used in extract_reasoning tests."""
+    return request

vllm_omni/engine/arg_utils.py

@@ -84,6 +84,9 @@ def register_omni_models_to_vllm():
         if arch not in supported_archs:
             ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
 
+    # Register omni-specific reasoning parsers (e.g., step_audio).
+    import vllm_omni.reasoning  # noqa: F401
+
 
 @dataclass
 class OmniEngineArgs(EngineArgs):

vllm_omni/engine/stage_engine_core_proc.py

@@ -89,6 +89,20 @@ class StageEngineCoreProc(EngineCoreProc):
         signal_callback: SignalCallback | None = None
         maybe_register_config_serialize_by_value()
 
+        # Register vllm-omni reasoning parsers (e.g. step_audio) in this
+        # subprocess so they are available when the engine core resolves
+        # ``--reasoning-parser``.  The main process already registered them
+        # at import time, but the forked subprocess starts with a fresh
+        # ReasoningParserManager.
+        try:
+            import vllm_omni.reasoning  # noqa: F401
+        except ImportError:
+            logger.warning(
+                "Failed to import vllm_omni.reasoning in subprocess; "
+                "custom reasoning parsers (e.g. step_audio) will not be "
+                "available."
+            )
+
         engine_core: StageEngineCoreProc | None = None
         coord_client: OmniCoordClientForStage | None = None
         try:

vllm_omni/entrypoints/openai/api_server.py

@@ -425,6 +425,15 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin)
 
+    # In vllm-omni's multi-process architecture, create_engine_config() runs
+    # in the subprocess (StageEngineCoreProc), not in the API server process.
+    # Propagate reasoning_parser from CLI args into structured_outputs_config
+    # here so that OmniOpenAIServingChat receives the correct value.
+    # (Upstream vLLM does this in EngineArgs.create_engine_config().)
+    if hasattr(args, "reasoning_parser") and args.reasoning_parser:
+        if not args.structured_outputs_config.reasoning_parser:
+            args.structured_outputs_config.reasoning_parser = args.reasoning_parser
+
     # Load logging config for uvicorn if specified
     log_config = get_uvicorn_log_config(args)
     if log_config is not None:

vllm_omni/reasoning/__init__.py

@@ -0,0 +1,24 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from vllm.reasoning import ReasoningParserManager
+
+# Step-Audio specific reasoning parser.
+# Registered at import time so it is available before the engine
+# resolves ``--reasoning-parser step_audio``.
+
+_OMNI_REASONING_PARSERS = {
+    "step_audio": (
+        "vllm_omni.reasoning.step_audio_reasoning_parser",
+        "StepAudioReasoningParser",
+    ),
+}
+
+
+def register_omni_reasoning_parsers():
+    """Register vllm-omni reasoning parsers with vLLM's ReasoningParserManager."""
+    for name, (module_path, class_name) in _OMNI_REASONING_PARSERS.items():
+        ReasoningParserManager.register_lazy_module(name, module_path, class_name)
+
+
+register_omni_reasoning_parsers()

vllm_omni/reasoning/step_audio_reasoning_parser.py

@@ -0,0 +1,574 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections.abc import Iterable, Sequence
+from itertools import islice
+from typing import TYPE_CHECKING
+
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.logger import init_logger
+from vllm.reasoning import ReasoningParser
+from vllm.tokenizers import TokenizerLike
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+    from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
+
+
+class StepAudioReasoningParser(ReasoningParser):
+    """Reasoning parser for Step-Audio models.
+
+    Step-Audio supports two representations of thinking markers:
+
+    1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>``
+       (single-token IDs, e.g. 151669 and 151670).
+
+    2. **Text markers**: ```` and ``````
+       (multi-token sequences, e.g. ``</think`` → [522, 26865],
+       ``>`` → [29]).
+
+    The chat template and typical model output use the **text form**
+    (multi-token).  The parser therefore always uses **text-based**
+    matching for streaming extraction, but also checks single-token
+    IDs when available for methods that only receive token IDs
+    (``is_reasoning_end``, ``extract_content_ids``, etc.).
+    """
+
+    # Text-form think markers (used in chat template & typical model output).
+    THINK_START_TEXT = "<think>"
+    THINK_END_TEXT = "</think>"
+
+    # Special-token-form think markers (single-token IDs in vocab).
+    THINK_START_SPECIAL = "<|THINK_START|>"
+    THINK_END_SPECIAL = "<|THINK_END|>"
+
+    def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
+        super().__init__(tokenizer, *args, **kwargs)
+
+        if not self.model_tokenizer:
+            raise ValueError(
+                "The model tokenizer must be passed to the ReasoningParser constructor during construction."
+            )
+
+        # --- Resolve think markers ----------------------------------------
+        # Special-token form (single tokens in vocab for fast ID checks).
+        self.think_start_special_id: int = self.vocab.get(self.THINK_START_SPECIAL, -1)
+        self.think_end_special_id: int = self.vocab.get(self.THINK_END_SPECIAL, -1)
+
+        # Text-form markers in vocab (usually NOT single tokens for
+        # Step-Audio, but checked for completeness).
+        self.think_start_text_id: int = self.vocab.get(self.THINK_START_TEXT, -1)
+        self.think_end_text_id: int = self.vocab.get(self.THINK_END_TEXT, -1)
+
+        # Set canonical token IDs for the ReasoningParser base contract.
+        # Prefer special-token IDs; fall back to text-form IDs.
+        self.think_start_token_id: int = (
+            self.think_start_special_id if self.think_start_special_id != -1 else self.think_start_text_id
+        )
+        self.think_end_token_id: int = (
+            self.think_end_special_id if self.think_end_special_id != -1 else self.think_end_text_id
+        )
+
+        # Whether single-token ID matching is possible for start/end.
+        self._has_start_token_id = self.think_start_token_id != -1
+        self._has_end_token_id = self.think_end_token_id != -1
+
+        # For text-based matching we must recognise BOTH forms.
+        # self.think_start_token / self.think_end_token are used by
+        # some helper methods and for the text-based path.
+        self.think_start_token = self.THINK_START_TEXT
+        self.think_end_token = self.THINK_END_TEXT
+
+        # Tracks whether we have already emitted the end of reasoning
+        # in streaming mode.
+        self._reasoning_ended = False
+
+        # Tracks whether reasoning JUST ended (within the last call),
+        # so we can strip a leading newline from the next content delta.
+        self._just_ended = False
+
+        # Buffer for text that might be the start of a multi-token
+        # think marker (e.g. "</" could be the beginning of "</think>").
+        # We hold this text until the next delta disambiguates it.
+        self._pending = ""
+
+        # Complete think marker strings used for ambiguous-prefix detection.
+        # When streaming, the tail of delta_text might be the start of one
+        # of these markers (e.g. "</" could be the beginning of ATHINK_END_AT).
+        # We buffer such ambiguous text to avoid leaking partial markers.
+        self._all_markers = (
+            self.THINK_START_TEXT,
+            self.THINK_END_TEXT,
+            self.THINK_START_SPECIAL,
+            self.THINK_END_SPECIAL,
+        )
+
+    # ------------------------------------------------------------------
+    # Token-ID helpers
+    # ------------------------------------------------------------------
+
+    def _has_end_token_in_ids(self, token_ids: Sequence[int]) -> bool:
+        """Check for end marker in token IDs (single-token forms only)."""
+        if self.think_end_special_id != -1 and self.think_end_special_id in token_ids:
+            return True
+        if self.think_end_text_id != -1 and self.think_end_text_id in token_ids:
+            return True
+        return False
+
+    def _has_start_token_in_ids(self, token_ids: Sequence[int]) -> bool:
+        """Check for start marker in token IDs (single-token forms only)."""
+        if self.think_start_special_id != -1 and self.think_start_special_id in token_ids:
+            return True
+        if self.think_start_text_id != -1 and self.think_start_text_id in token_ids:
+            return True
+        return False
+
+    # ------------------------------------------------------------------
+    # Text-based helpers (handles both text-form and special-form markers)
+    # ------------------------------------------------------------------
+
+    def _has_end_token_in_text(self, text: str) -> bool:
+        return self.THINK_END_TEXT in text or self.THINK_END_SPECIAL in text
+
+    def _has_start_token_in_text(self, text: str) -> bool:
+        return self.THINK_START_TEXT in text or self.THINK_START_SPECIAL in text
+
+    def _find_end_token_in_text(self, text: str) -> tuple[int, int]:
+        """Find the first end token in text. Returns (start, end) indices
+        such that text[start:end] is the end token.  Returns (-1, -1) if
+        not found."""
+        idx = text.find(self.THINK_END_TEXT)
+        if idx != -1:
+            return idx, idx + len(self.THINK_END_TEXT)
+        idx = text.find(self.THINK_END_SPECIAL)
+        if idx != -1:
+            return idx, idx + len(self.THINK_END_SPECIAL)
+        return -1, -1
+
+    def _find_start_token_in_text(self, text: str) -> tuple[int, int]:
+        """Find the first start token in text. Returns (start, end) indices
+        such that text[start:end] is the start token.  Returns (-1, -1) if
+        not found."""
+        idx = text.find(self.THINK_START_TEXT)
+        if idx != -1:
+            return idx, idx + len(self.THINK_START_TEXT)
+        idx = text.find(self.THINK_START_SPECIAL)
+        if idx != -1:
+            return idx, idx + len(self.THINK_START_SPECIAL)
+        return -1, -1
+
+    def _decode_ids(self, token_ids: Sequence[int]) -> str:
+        """Decode a sequence of token IDs to text using the model tokenizer."""
+        return self.model_tokenizer.decode(token_ids, skip_special_tokens=False)
+
+    # ------------------------------------------------------------------
+    # ReasoningParser interface
+    # ------------------------------------------------------------------
+
+    def is_reasoning_end(self, input_ids: Sequence[int]) -> bool:
+        """Check if reasoning has ended in the given token sequence.
+
+        When called with **prompt** token IDs (by the serving layer), the
+        prompt may contain think markers from previous assistant turns.
+        In multi-turn conversations the prompt can include both start
+        and end markers, with the *last* marker being a start marker
+        (from the generation prompt).  In that case reasoning is NOT
+        ended — the model is about to generate inside a new think block.
+
+        We therefore find the **last** think marker (start or end) in
+        the decoded text and only return True if it is an end marker.
+        """
+        try:
+            text = self._decode_ids(input_ids)
+        except Exception:
+            # Fall back to token-ID-only check.
+            result = self._has_end_token_in_ids(input_ids)
+            logger.debug(
+                "StepAudio is_reasoning_end: decode failed, fallback result=%s, num_ids=%d",
+                result,
+                len(input_ids),
+            )
+            return result
+
+        # Find the position of the LAST think marker (start or end).
+        last_start = max(
+            text.rfind(self.THINK_START_TEXT),
+            text.rfind(self.THINK_START_SPECIAL),
+        )
+        last_end = max(
+            text.rfind(self.THINK_END_TEXT),
+            text.rfind(self.THINK_END_SPECIAL),
+        )
+
+        if last_end == -1:
+            # No end marker at all → reasoning not ended.
+            logger.debug(
+                "StepAudio is_reasoning_end: no end marker found, text_tail=%r (len=%d), returning False",
+                text[-100:] if len(text) > 100 else text,
+                len(text),
+            )
+            return False
+
+        if last_start > last_end:
+            # Last marker is a START marker (e.g. the generation prompt's
+            # ``).  Reasoning is still active.
+            logger.debug(
+                "StepAudio is_reasoning_end: last_start=%d > last_end=%d, reasoning still active, returning False",
+                last_start,
+                last_end,
+            )
+            return False
+
+        # Last marker is an END marker → reasoning has ended.
+        logger.debug(
+            "StepAudio is_reasoning_end: last_end=%d >= last_start=%d, reasoning ended, returning True",
+            last_end,
+            last_start,
+        )
+        return True
+
+    def is_reasoning_end_streaming(self, input_ids: Sequence[int], delta_ids: Iterable[int]) -> bool:
+        delta_list = tuple(delta_ids)
+        # Fast check: single-token markers in delta.
+        if self._has_end_token_in_ids(delta_list):
+            return True
+        # Fallback: decode delta + tail of input_ids and check text.
+        # The end marker may span the boundary between previous tokens
+        # and new delta tokens when it is encoded as multiple tokens
+        # (e.g. "</think" = [522, 26865]).
+        try:
+            delta_text = self._decode_ids(delta_list)
+            if self._has_end_token_in_text(delta_text):
+                return True
+            overlap_ids = list(input_ids[-8:]) + list(delta_list)
+            overlap_text = self._decode_ids(overlap_ids)
+            return self._has_end_token_in_text(overlap_text)
+        except Exception:
+            return False
+
+    def extract_content_ids(self, input_ids: list[int]) -> list[int]:
+        # Fast check: single end-token in IDs.
+        if self._has_end_token_id:
+            # Check for special-token or text-form single token.
+            for end_id in (self.think_end_special_id, self.think_end_text_id):
+                if end_id != -1 and end_id in islice(input_ids, 0, max(0, len(input_ids) - 1)):
+                    idx = input_ids.index(end_id)
+                    return input_ids[idx + 1 :]
+
+        # Fallback: decode to text and find multi-token end marker.
+        try:
+            text = self._decode_ids(input_ids)
+            end_pos, end_len = self._find_end_token_in_text(text)
+            if end_pos == -1:
+                return []
+            content_start = end_len
+            if content_start < len(text) and text[content_start] == "\n":
+                content_start += 1
+            if content_start >= len(text):
+                return []
+            for i in range(len(input_ids)):
+                prefix_text = self._decode_ids(input_ids[: i + 1])
+                if len(prefix_text) > content_start:
+                    return input_ids[i:]
+            return []
+        except Exception:
+            return []
+
+    def extract_reasoning(
+        self,
+        model_output: str,
+        request: "ChatCompletionRequest | ResponsesRequest",
+    ) -> tuple[str | None, str | None]:
+        logger.debug(
+            "StepAudio extract_reasoning: model_output=%r (len=%d)",
+            model_output[:200],
+            len(model_output),
+        )
+        # Strip leading start token if present (text or special form).
+        for start_tok in (self.THINK_START_TEXT, self.THINK_START_SPECIAL):
+            if model_output.startswith(start_tok):
+                model_output = model_output[len(start_tok) :]
+                logger.debug("StepAudio extract_reasoning: stripped start_tok=%r", start_tok)
+                break
+
+        # Find the end token (text or special form).
+        end_pos, end_len = self._find_end_token_in_text(model_output)
+        if end_pos == -1:
+            # No end token found — everything is reasoning.
+            logger.debug(
+                "StepAudio extract_reasoning: no end token, reasoning=%r (len=%d), content=None",
+                model_output[:200],
+                len(model_output),
+            )
+            return model_output or None, None
+
+        reasoning = model_output[:end_pos]
+        content = model_output[end_len:]
+        # Strip leading newline that models often emit right after end token.
+        content = content.lstrip("\n")
+        logger.debug(
+            "StepAudio extract_reasoning: end_pos=%d, reasoning=%r (len=%d), content=%r (len=%d)",
+            end_pos,
+            reasoning[:200],
+            len(reasoning),
+            content[:200],
+            len(content),
+        )
+        return reasoning or None, content or None
+
+    def extract_reasoning_streaming(
+        self,
+        previous_text: str,
+        current_text: str,
+        delta_text: str,
+        previous_token_ids: Sequence[int],
+        current_token_ids: Sequence[int],
+        delta_token_ids: Sequence[int],
+    ) -> DeltaMessage | None:
+        # Step-Audio typically uses multi-token text-form markers
+        # (e.g. "</think" → [522, 26865] + ">" → [29]).
+        # We always use the text-based path because it handles both
+        # single-token special markers and multi-token text markers.
+        # The text-based path is also efficient since previous_text
+        # and delta_text are already available.
+        return self._extract_streaming_text(
+            previous_text,
+            delta_text,
+        )
+
+    # ------------------------------------------------------------------
+    # Streaming: text-based (handles both text and special token forms)
+    # ------------------------------------------------------------------
+
+    def _ends_with_marker_prefix(self, text: str) -> str:
+        """If text ends with a prefix of any think marker, return the
+        ambiguous suffix that should be buffered.
+
+        For example:
+        - text="reasoning</" → marker ATHINK_END_AT starts with "</" → return "</"
+        - text="reasoning<" → markers start with "<" → return "<"
+        - text="reasoning<th" → ATHINK_START starts with "<th" → return "<th"
+          (but "</think" also starts with "<", so we buffer from the longest)
+        - text="reasoning</thi" → ATHINK_END_AT starts with "</thi" → return "</thi"
+        """
+        if not text:
+            return ""
+
+        # Find the longest suffix of text that is a prefix of any marker.
+        best = ""
+        for marker in self._all_markers:
+            for i in range(1, len(marker) + 1):
+                prefix = marker[:i]
+                if text.endswith(prefix) and len(prefix) > len(best):
+                    best = prefix
+        return best
+
+    # Counter for debug logging (to avoid flooding logs)
+    _stream_call_count: int = 0
+
+    def _extract_streaming_text(
+        self,
+        previous_text: str,
+        delta_text: str,
+    ) -> DeltaMessage | None:
+        # Reset _just_ended flag at the start of each call.
+        # It will be set to True if this call ends reasoning.
+        just_ended_before = self._just_ended
+        self._just_ended = False
+
+        # Debug logging for first few calls
+        StepAudioReasoningParser._stream_call_count += 1
+        call_num = StepAudioReasoningParser._stream_call_count
+        if call_num <= 20:
+            logger.debug(
+                "StepAudio stream #%d: prev_text=%r (len=%d), delta_text=%r (len=%d), _reasoning_ended=%s, _pending=%r",
+                call_num,
+                previous_text[:100],
+                len(previous_text),
+                delta_text[:100],
+                len(delta_text),
+                self._reasoning_ended,
+                self._pending[:50] if self._pending else "",
+            )
+
+        # Prepend any pending text from a previous call.
+        # The pending text is already part of previous_text (the framework
+        # accumulates all generated text), so we must also trim it from
+        # previous_text to avoid double-counting when we form `combined`.
+        if self._pending:
+            n = len(self._pending)
+            if previous_text.endswith(self._pending):
+                previous_text = previous_text[:-n]
+            delta_text = self._pending + delta_text
+            self._pending = ""
+
+        combined = previous_text + delta_text
+        has_end = self._has_end_token_in_text(combined)
+        has_end_in_previous = self._has_end_token_in_text(previous_text)
+        has_end_in_delta = self._has_end_token_in_text(delta_text)
+
+        if call_num <= 20:
+            logger.debug(
+                "StepAudio stream #%d: has_end=%s, has_end_in_previous=%s, has_end_in_delta=%s, combined=%r (len=%d)",
+                call_num,
+                has_end,
+                has_end_in_previous,
+                has_end_in_delta,
+                combined[:150],
+                len(combined),
+            )
+
+        if has_end_in_previous or self._reasoning_ended:
+            # Already past reasoning — everything is content.
+            self._reasoning_ended = True
+            # Strip at most one leading newline right after the end marker.
+            if delta_text.startswith("\n") and just_ended_before:
+                return DeltaMessage(content=delta_text[1:] or None)
+            return DeltaMessage(content=delta_text)
+
+        if has_end_in_delta:
+            # End token fully contained in the delta — split reasoning
+            # and content.
+            end_index, end_len = self._find_end_token_in_text(delta_text)
+            reasoning = delta_text[:end_index]
+            content = delta_text[end_len:]
+            # Strip leading newline after end token.
+            if content.startswith("\n"):
+                content = content[1:]
+            self._reasoning_ended = True
+            self._just_ended = True
+            return DeltaMessage(
+                reasoning=reasoning or None,
+                content=content or None,
+            )
+
+        if has_end and not has_end_in_previous and not has_end_in_delta:
+            # End token spans the boundary between previous_text and
+            # delta_text (e.g. previous ends with "</" and delta
+            # starts with "think>...").  Find the split point in the
+            # combined text.
+            end_index, end_len = self._find_end_token_in_text(combined)
+            marker_start_in_prev = end_index - len(previous_text)
+            if marker_start_in_prev > 0:
+                # Some of previous_text's trailing chars are part of
+                # the marker — they were already emitted as reasoning
+                # in a previous call.  The delta begins the content
+                # after the marker.
+                content_after = combined[end_len:]
+                if content_after.startswith("\n"):
+                    content_after = content_after[1:]
+                self._reasoning_ended = True
+                self._just_ended = True
+                return DeltaMessage(content=content_after or None)
+            else:
+                # Edge case: the entire marker is in previous_text
+                # but was not detected because of how text was split
+                # across streaming chunks.
+                self._reasoning_ended = True
+                self._just_ended = True
+                full_content = combined[end_len:]
+                if full_content.startswith("\n"):
+                    full_content = full_content[1:]
+                return DeltaMessage(content=full_content or None)
+
+        # No end token seen yet — everything is reasoning.
+        # Skip the start token prefix if present in delta.
+        start_idx, start_len = self._find_start_token_in_text(delta_text)
+        if start_idx != -1 and delta_text.strip() == delta_text[start_idx:start_len]:
+            # Delta is only the start token — skip it entirely.
+            return None
+        if start_idx == 0:
+            # Delta starts with the start token — skip past it.
+            after_start = delta_text[start_len:]
+            return DeltaMessage(reasoning=after_start or None)
+
+        # Check if the delta ends with a partial marker prefix (e.g. "</"
+        # which could be the beginning of the end marker).  Buffer the
+        # ambiguous text to avoid leaking partial markers into reasoning.
+        ambiguous_suffix = self._ends_with_marker_prefix(delta_text)
+        if ambiguous_suffix:
+            if ambiguous_suffix == delta_text:
+                # Entire delta is a possible marker prefix — buffer it all.
+                self._pending = delta_text
+                return None
+            else:
+                # Emit the non-ambiguous part as reasoning, buffer the rest.
+                to_emit = delta_text[: -len(ambiguous_suffix)]
+                self._pending = ambiguous_suffix
+                return DeltaMessage(reasoning=to_emit) if to_emit else None
+
+        return DeltaMessage(reasoning=delta_text)
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        """Count tokens within thinking spans."""
+        # Fast path: if single-token markers are present, count using IDs.
+        if self._has_start_token_id and self._has_end_token_id:
+            count = 0
+            depth = 0
+            for token_id in token_ids:
+                if token_id == self.think_start_token_id:
+                    depth += 1
+                    continue
+                if token_id == self.think_end_token_id:
+                    if depth > 0:
+                        depth -= 1
+                    continue
+                if depth > 0:
+                    count += 1
+            # Only trust the fast-path result if we actually found markers.
+            if depth == 0 and count > 0:
+                return count
+
+        # Text-based fallback: handles both text-form multi-token markers
+        # and special-token single-token markers in decoded text.
+        try:
+            text = self._decode_ids(token_ids)
+            if not text:
+                return 0
+            n_tokens = len(token_ids)
+            if n_tokens == 0:
+                return 0
+
+            total_chars = len(text)
+            if total_chars == 0:
+                return 0
+            chars_per_token = total_chars / n_tokens
+
+            count = 0
+            depth = 0
+            i = 0
+            while i < len(text):
+                found_start = False
+                for start_tok in (self.THINK_START_TEXT, self.THINK_START_SPECIAL):
+                    if text[i:].startswith(start_tok):
+                        depth += 1
+                        i += len(start_tok)
+                        found_start = True
+                        break
+                if found_start:
+                    continue
+
+                found_end = False
+                for end_tok in (self.THINK_END_TEXT, self.THINK_END_SPECIAL):
+                    if text[i:].startswith(end_tok):
+                        if depth > 0:
+                            depth -= 1
+                        i += len(end_tok)
+                        found_end = True
+                        break
+                if found_end:
+                    continue
+
+                if depth > 0:
+                    count += 1
+                i += 1
+
+            # Scale per-character count to approximate token count.
+            if chars_per_token > 1:
+                count = max(1, round(count / chars_per_token))
+
+            return count
+        except Exception:
+            return 0