[Bugfix][Frontend] Count reasoning tokens for Harmony, DeepSeek-V3 and Step3 parsers (#58626)

Signed-off-by: Samyabrata Maji <116789799+sammaji@users.noreply.github.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>

tests/parser/test_harmony.py

@@ -873,6 +873,67 @@ class TestProcessChunk:
         ]
 
 
+class TestCountReasoningTokens:
+    def test_matches_process_chunk(self, harmony_parser, gpt_oss_tokenizer):
+        """Chat usage must report the same count as the Responses API path."""
+        token_ids = get_model_output_tokens(
+            [
+                assistant("Let me check the weather.", "analysis"),
+                tool_call("functions.get_weather", '{"location": "SF"}'),
+            ]
+        )
+        count = harmony_parser.count_reasoning_tokens(token_ids)
+
+        fresh = HarmonyParser(gpt_oss_tokenizer)
+        assert count == fresh.process_chunk(token_ids).reasoning_token_count
+        assert count > 0
+
+    def test_excludes_harmony_control_tokens(self, harmony_parser):
+        text = "Let me think about this."
+        token_ids = get_model_output_tokens(
+            [assistant(text, "analysis"), assistant("Hi", "final")]
+        )
+        expected = len(get_encoding().encode(text))
+        assert harmony_parser.count_reasoning_tokens(token_ids) == expected
+
+    def test_final_only_has_no_reasoning(self, harmony_parser):
+        token_ids = get_model_output_tokens([assistant("Hello", "final")])
+        assert harmony_parser.count_reasoning_tokens(token_ids) == 0
+
+    def test_streaming_reuses_running_total(
+        self, harmony_parser, gpt_oss_tokenizer, chat_request, monkeypatch
+    ):
+        """Called per streamed chunk with all tokens so far; must not replay."""
+        token_ids = get_model_output_tokens(
+            [assistant("Think", "analysis"), assistant("Answer", "final")]
+        )
+        expected = HarmonyParser(gpt_oss_tokenizer).count_reasoning_tokens(token_ids)
+        harmony_parser.parse_delta("", token_ids[:3], chat_request, finished=False)
+
+        def fail():
+            raise AssertionError("count_reasoning_tokens replayed the output")
+
+        monkeypatch.setattr(
+            "vllm.parser.harmony.get_streamable_parser_for_assistant", fail
+        )
+        for end in range(3, len(token_ids)):
+            harmony_parser.process_chunk(token_ids[end : end + 1])
+            harmony_parser.count_reasoning_tokens(token_ids[: end + 1])
+        assert harmony_parser.count_reasoning_tokens(token_ids) == expected
+
+    def test_does_not_disturb_stream_state(self, harmony_parser, chat_request):
+        """Streaming counts after each chunk, so it must not consume parser state."""
+        token_ids = get_model_output_tokens(
+            [assistant("Think", "analysis"), assistant("Answer", "final")]
+        )
+        harmony_parser.count_reasoning_tokens(token_ids)
+
+        reasoning, content, _ = harmony_parser.parse(
+            "", chat_request, model_output_token_ids=token_ids
+        )
+        assert (reasoning, content) == ("Think", "Answer")
+
+
 class TestAdjustRequest:
     REQUEST_TEXT = "Hello"
     TOOL_TYPE = "function"

tests/reasoning/test_deepseekv3_reasoning_parser.py

@@ -87,3 +87,14 @@ def test_identity_reasoning_parser_basic(tokenizer):
         delta_token_ids=[],
     )
     assert result_none is None
+
+
+@pytest.mark.parametrize("thinking,expected", [(True, 3), (False, 0)])
+def test_count_reasoning_tokens_delegates(tokenizer, thinking, expected):
+    parser = DeepSeekV3ReasoningParser(
+        tokenizer, chat_template_kwargs={"thinking": thinking}
+    )
+    think = tokenizer.convert_tokens_to_ids(["<think>", "</think>"])
+    token_ids = [think[0], 11, 12, 13, think[1], 14]
+
+    assert parser.count_reasoning_tokens(token_ids) == expected

tests/reasoning/test_step3_reasoning_parser.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.reasoning import ReasoningParserManager
+from vllm.tokenizers import get_tokenizer
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+REASONING_MODEL_NAME = "stepfun-ai/step3"
+
+
+@pytest.fixture(scope="module")
+def step3_parser():
+    tokenizer = get_tokenizer(tokenizer_name=REASONING_MODEL_NAME)
+    return ReasoningParserManager.get_reasoning_parser("step3")(tokenizer)
+
+
+def test_count_reasoning_tokens_stops_at_think_end(step3_parser):
+    end = step3_parser.think_end_token_id
+    assert step3_parser.count_reasoning_tokens([11, 12, end, 13]) == 2
+
+
+def test_count_reasoning_tokens_without_think_end(step3_parser):
+    """Step3 treats output as reasoning until </think> appears."""
+    assert step3_parser.count_reasoning_tokens([11, 12, 13]) == 3

vllm/parser/harmony.py

@@ -38,6 +38,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
 )
 from vllm.entrypoints.openai.parser.harmony_utils import (
     extract_function_from_recipient,
+    get_encoding,
     get_streamable_parser_for_assistant,
     is_function_recipient,
 )
@@ -114,6 +115,9 @@ class HarmonyParser(DelegatingParser):
         self._next_tool_call_index = 0
         self._num_processed_messages = 0
 
+        self._num_counted_tokens = 0
+        self._num_reasoning_tokens = 0
+
         # For error recovery
         self._current_message_tokens: list[int] = []
 
@@ -348,9 +352,7 @@ class HarmonyParser(DelegatingParser):
             else:
                 self._current_message_tokens.append(token_id)
 
-            if channel == "analysis" or (
-                channel == "commentary" and recipient is not None
-            ):
+            if self._is_reasoning_token(token_id, channel, recipient):
                 reasoning_token_count += 1
 
             segments.append(
@@ -364,11 +366,35 @@ class HarmonyParser(DelegatingParser):
 
             # TODO: Optionally merge and suppress empty Segments
 
+        self._num_counted_tokens += len(token_ids)
+        self._num_reasoning_tokens += reasoning_token_count
         return ChunkResult(
             segments=segments,
             reasoning_token_count=reasoning_token_count,
         )
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        if len(token_ids) == self._num_counted_tokens:
+            return self._num_reasoning_tokens
+
+        parser = get_streamable_parser_for_assistant()
+        count = 0
+        for token_id in token_ids:
+            parser.process(token_id)
+            recipient = self._normalize_recipient(parser.current_recipient)
+            if self._is_reasoning_token(token_id, parser.current_channel, recipient):
+                count += 1
+        return count
+
+    @staticmethod
+    def _is_reasoning_token(
+        token_id: int, channel: str | None, recipient: str | None
+    ) -> bool:
+        is_reasoning_channel = channel == "analysis" or (
+            channel == "commentary" and recipient is not None
+        )
+        return is_reasoning_channel and not get_encoding().is_special_token(token_id)
+
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest:

vllm/reasoning/deepseek_v3_reasoning_parser.py

@@ -55,6 +55,9 @@ class DeepSeekV3ReasoningParser(ReasoningParser):
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
         return self._parser.extract_content_ids(input_ids)
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        return self._parser.count_reasoning_tokens(token_ids)
+
     def extract_reasoning(
         self, model_output: str, request: "ChatCompletionRequest | ResponsesRequest"
     ) -> tuple[str | None, str | None]:

vllm/reasoning/step3_reasoning_parser.py

@@ -121,3 +121,8 @@ class Step3ReasoningParser(ReasoningParser):
             return []
         else:
             return input_ids[input_ids.index(self.think_end_token_id) + 1 :]
+
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        if self.think_end_token_id not in token_ids:
+            return len(token_ids)
+        return list(token_ids).index(self.think_end_token_id)