[Bugfix][Frontend] Count reasoning tokens for the Muse-Glimmer parser

tests/reasoning/test_muse_glimmer_count_reasoning_tokens.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import string
+
+import pytest
+
+from vllm.reasoning.muse_glimmer_reasoning_parser import MuseGlimmerReasoningParser
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class CharTokenizer:
+    """Character-level tokenizer: MuseGlimmer's ATEM markers are multi-token,
+    matching the property the parser is written against."""
+
+    def __init__(self):
+        self._chars = {c: i for i, c in enumerate(string.printable)}
+        self._ids = {i: c for c, i in self._chars.items()}
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> list[int]:
+        return [self._chars[c] for c in text]
+
+    def decode(self, token_ids) -> str:
+        return "".join(self._ids[i] for i in token_ids)
+
+
+@pytest.fixture
+def parser_and_tokenizer():
+    tokenizer = CharTokenizer()
+    return MuseGlimmerReasoningParser(tokenizer), tokenizer
+
+
+def test_counts_closed_reasoning_span(parser_and_tokenizer):
+    parser, tok = parser_and_tokenizer
+    text = "to=self<|message|>think hard<|eom|>to=user<|message|>hi<|eot|>"
+    assert parser.count_reasoning_tokens(tok.encode(text)) == len("think hard")
+
+
+def test_counts_multiple_reasoning_spans(parser_and_tokenizer):
+    parser, tok = parser_and_tokenizer
+    text = (
+        "to=self<|message|>abc<|eom|>"
+        "to=user<|message|>answer<|eot|>"
+        "to=self<|message|>de<|eom|>"
+    )
+    assert parser.count_reasoning_tokens(tok.encode(text)) == len("abc") + len("de")
+
+
+def test_counts_open_reasoning_span_mid_stream(parser_and_tokenizer):
+    parser, tok = parser_and_tokenizer
+    text = "to=self<|message|>partial thought"
+    assert parser.count_reasoning_tokens(tok.encode(text)) == len("partial thought")
+
+
+def test_content_only_counts_zero(parser_and_tokenizer):
+    parser, tok = parser_and_tokenizer
+    text = "to=user<|message|>plain answer<|eot|>"
+    assert parser.count_reasoning_tokens(tok.encode(text)) == 0
+
+
+def test_empty_counts_zero(parser_and_tokenizer):
+    parser, _ = parser_and_tokenizer
+    assert parser.count_reasoning_tokens([]) == 0
+
+
+def test_mid_generation_turn_reopen(parser_and_tokenizer):
+    """Regression: real generations re-open the assistant turn before the
+    answer channel; the reasoning before it must still be counted."""
+    parser, tok = parser_and_tokenizer
+    text = (
+        " to=self<|message|>Simple. 2+2 = 4<|eom|>"
+        "<|start|>assistant to=user<|message|>2 + 2 = 4"
+    )
+    assert parser.count_reasoning_tokens(tok.encode(text)) == len("Simple. 2+2 = 4")
+
+
+def test_tool_channel_not_counted(parser_and_tokenizer):
+    parser, tok = parser_and_tokenizer
+    text = (
+        "to=self<|message|>pick a tool<|eom|>"
+        'to=functions.get_weather<|message|>{"city": "SF"}<|eom|>'
+    )
+    assert parser.count_reasoning_tokens(tok.encode(text)) == len("pick a tool")

vllm/reasoning/muse_glimmer_reasoning_parser.py

@@ -161,6 +161,34 @@ def extract_content_ids(self, input_ids: list[int]) -> list[int]:
         # path uses extract_reasoning() for the final split.
         return []
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        """Count generated tokens that belong to ``to=self`` reasoning bodies.
+
+        MuseGlimmer's framing markers are not guaranteed to be single vocab
+        tokens, so the marker-token counting used by ``BaseThinkingReasoningParser``
+        does not apply. Decode the generated ids, reuse the channel
+        classification that already drives extraction, and re-encode the
+        reasoning body. Re-encoding approximates the original token count at
+        span boundaries, which is acceptable for usage accounting.
+        """
+        if not token_ids:
+            return 0
+        try:
+            text = self.model_tokenizer.decode(token_ids)
+        except Exception:
+            return 0
+        # token_ids are generation-only, so no prompt anchoring is needed.
+        # Do NOT scope to _current_assistant_turn here: the model re-opens the
+        # turn mid-generation (``<|eom|><|start|>assistant to=user<|message|>``),
+        # and anchoring on the LAST turn-open would slice off the reasoning
+        # that precedes it.
+        reasoning_body, _ = self._classify_bodies(text)
+        if not reasoning_body:
+            return 0
+        return len(
+            self.model_tokenizer.encode(reasoning_body, add_special_tokens=False)
+        )
+
     @classmethod
     def _scoped_turn(cls, text: str) -> str:
         """Current assistant turn with reasoning spans removed."""