diff --git a/pyproject.toml b/pyproject.toml
index 09b8e23ca..757f038af 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -292,5 +292,9 @@ fullset = "fullset"
Vai = "Vai"
tockens = "tockens"
CANN = "CANN"
+# Step-Audio reasoning parser streams `` tokens split across chunks
+# (e.g. ""). 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"
diff --git a/tests/reasoning/__init__.py b/tests/reasoning/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/tests/reasoning/test_step_audio_reasoning_parser.py b/tests/reasoning/test_step_audio_reasoning_parser.py
new file mode 100644
index 000000000..7eb9fa17a
--- /dev/null
+++ b/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**: ```` and ````
+ (multi-token sequences for , 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 #
+TE = StepAudioReasoningParser.THINK_END_TEXT #
+TSS = StepAudioReasoningParser.THINK_START_SPECIAL # <|THINK_START|>
+TES = StepAudioReasoningParser.THINK_END_SPECIAL # <|THINK_END|>
+
+
+# ---------------------------------------------------------------------------
+# Mock tokenizers
+# ---------------------------------------------------------------------------
+
+
+class SingleTokenTokenizer:
+ """Tokenizer where both ```` and ```` 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, #
+ TE: 101, #
+ TSS: 151669, # <|THINK_START|>
+ TES: 151670, # <|THINK_END|>
+ "Hello": 1,
+ " world": 2,
+ "reasoning": 3,
+ " content": 4,
+ "\n": 5,
+ "": 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 ```` is NOT a single token.
+
+ ```` is in the vocab (ID 100), but ```` is NOT.
+ Instead, ```` encodes as [""] → [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, # - single token
+ # is intentionally NOT in the vocab
+ "": 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 ```` and ```` 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, #
+ TE: 101, #
+ "Hello": 1,
+ " world": 2,
+ "reasoning": 3,
+ " content": 4,
+ "\n": 5,
+ "": 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):
+ """ is a single token; 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 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):
+ """ = [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 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: ... " + " " + ", 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 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 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: reasoning final
+ if parser_name == "multi_token_parser":
+ deltas = [TS, "reasoning", TE, " content"]
+ token_ids_per_delta = [
+ [100], #
+ [3], # reasoning
+ [522, 26865], #
+ [4], # content
+ ]
+ else:
+ deltas = [TS, "reasoning", TE, " content"]
+ token_ids_per_delta = [
+ [100], #
+ [3], # reasoning
+ [101], #
+ [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: is split across streaming chunks.
+
+ Previous chunk ends with "reasoning\\nfinal answer". This is the real-world scenario when
+ ```` is tokenized as two separate tokens.
+ """
+ deltas_and_ids = [
+ ("reasoning", [3]),
+ ("\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 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):
+ """ split across 3 chunks: ."""
+ deltas_and_ids = [
+ ("some reasoning", [3]),
+ ("", [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 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 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
+ # "\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):
+ """ split across 5 tiny chunks."""
+ # Split: each chunk is one character of
+ 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=, 3=reasoning, 101=, 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):
+ # encodes as [522, 26865]
+ # Full sequence: reasoning 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 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=, 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):
+ # = [522, 26865]
+ # [3, 522, 26865, 4] → content after
+ 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
diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py
index ac52de953..148412012 100644
--- a/vllm_omni/engine/arg_utils.py
+++ b/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):
diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py
index 32c92ff68..eb989db0e 100644
--- a/vllm_omni/engine/stage_engine_core_proc.py
+++ b/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:
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 67a639b15..0164e584a 100755
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/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:
diff --git a/vllm_omni/reasoning/__init__.py b/vllm_omni/reasoning/__init__.py
new file mode 100644
index 000000000..37cdcf0fd
--- /dev/null
+++ b/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()
diff --git a/vllm_omni/reasoning/step_audio_reasoning_parser.py b/vllm_omni/reasoning/step_audio_reasoning_parser.py
new file mode 100644
index 000000000..81ef63496
--- /dev/null
+++ b/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. ```` → [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_END_TEXT = ""
+
+ # 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. "").
+ # 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. " 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. " 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. "" → [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 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 "..."). 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. " 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