From bb363db9a5ec2edc7b39e99b00af363a89d1fb81 Mon Sep 17 00:00:00 2001
From: JasonCohere <jasonozuzu@cohere.com>
Date: Thu, 3 Sep 2026 13:02:49 +0100
Subject: [PATCH] feat: Add support for reasoning_token_count to reasoning
 parser (#54982)

Signed-off-by: Jason Ozuzu <jasonozuzu@cohere.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../test_cohere_command_reasoning_parser.py   | 19 +++++++++++++++++++
 .../cohere_command_reasoning_parser.py        | 15 +++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/tests/reasoning/test_cohere_command_reasoning_parser.py b/tests/reasoning/test_cohere_command_reasoning_parser.py
index 5a1359567e..ea9ab63aef 100644
--- a/tests/reasoning/test_cohere_command_reasoning_parser.py
+++ b/tests/reasoning/test_cohere_command_reasoning_parser.py
@@ -253,6 +253,25 @@ class TestIsReasoningEnd:
         )
 
 
+@pytest.mark.parametrize(
+    "parser_cls",
+    [
+        CohereCommand3ReasoningParser,
+        CohereCommand4ReasoningParser,
+    ],
+    ids=["cmd3", "cmd4"],
+)
+def test_count_reasoning_tokens(tokenizer, parser_cls):
+    parser = parser_cls(tokenizer)
+    start = parser.start_token_id
+    end = parser.end_token_id
+
+    assert parser.count_reasoning_tokens([99, start, 11, 12, end, 100]) == 2
+    assert parser.count_reasoning_tokens([end, start, 1, start, 2, end, 3, end]) == 3
+    assert parser.count_reasoning_tokens([start, 1, 2]) == 2
+    assert parser.count_reasoning_tokens([1, 2, end]) == 0
+
+
 SCHEMA_A = {"type": "object", "properties": {"a": {"type": "string"}}}
 SCHEMA_B = {"type": "object", "properties": {"b": {"type": "number"}}}
 GET_WEATHER_TOOL = {
diff --git a/vllm/reasoning/cohere_command_reasoning_parser.py b/vllm/reasoning/cohere_command_reasoning_parser.py
index 55beaa6514..4dada76ee7 100644
--- a/vllm/reasoning/cohere_command_reasoning_parser.py
+++ b/vllm/reasoning/cohere_command_reasoning_parser.py
@@ -628,6 +628,21 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
 
         return has_end_token
 
+    def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
+        count = 0
+        depth = 0
+        for token_id in token_ids:
+            if token_id == self.start_token_id:
+                depth += 1
+                continue
+            if token_id == self.end_token_id:
+                if depth > 0:
+                    depth -= 1
+                continue
+            if depth > 0:
+                count += 1
+        return count
+
     def adjust_request(
         self, request: ChatCompletionRequest | ResponsesRequest
     ) -> ChatCompletionRequest | ResponsesRequest: