use reasoning state from chat-parser, thus removing duplicated logic

aviallon · aviallon · commit 9c82848b538f · 2025-12-09T13:30:41.000+01:00
diff --git a/common/chat-parser-xml-toolcall.cpp b/common/chat-parser-xml-toolcall.cpp
@@ -705,6 +705,9 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
 
     // Parse content
     bool reasoning_unclosed = builder.syntax().thinking_forced_open;
+    if (reasoning_unclosed) {
+        builder.mark_reasoning_active(end_think);
+    }
     std::string unclosed_reasoning_content("");
     for (;;) {
         auto tc = try_find_2_literal_splited_by_spaces(builder, form.scope_start, form.tool_start);
@@ -730,6 +733,7 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
                 }
             } else {
                 reasoning_unclosed = false;
+                builder.mark_reasoning_closed();
                 std::string reasoning_content;
                 if (pos == std::string::npos) {
                     reasoning_content = std::move(content);
@@ -766,13 +770,15 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
         bool toolcall_in_think = false;
         for (auto think_start = content.find(start_think); think_start != std::string::npos; think_start = content.find(start_think, think_start)) {
             if (auto think_end = content.find(end_think, think_start + start_think.size()); think_end != std::string::npos) {
+                builder.mark_reasoning_active(end_think);
                 if (builder.syntax().reasoning_format != COMMON_REASONING_FORMAT_NONE && !builder.syntax().reasoning_in_content) {
                     auto reasoning_content = content.substr(think_start + start_think.size(), think_end - think_start - start_think.size());
                     builder.add_reasoning_content(reasoning_content);
                     think_start = erase_spaces(content, think_start, think_end + end_think.size() - 1);
                 } else {
                     think_start = think_end + end_think.size() - 1;
                 }
+                builder.mark_reasoning_closed();
             } else {
                 // This <tool_call> start is in thinking block, skip this tool call
                 // This <tool_call> start is in thinking block
@@ -782,6 +788,7 @@ inline void parse_msg_with_xml_tool_calls(common_chat_msg_parser & builder, cons
                     unclosed_reasoning_content = content.substr(think_start + start_think.size()) + tool_call_start;
                 }
                 reasoning_unclosed = true;
+                builder.mark_reasoning_active(end_think);
                 content.resize(think_start);
                 toolcall_in_think = true;
             }
diff --git a/common/chat-parser.cpp b/common/chat-parser.cpp
@@ -156,6 +156,20 @@ void common_chat_msg_parser::add_reasoning_content(const std::string &reasoning_
     result_.reasoning_content += reasoning_content;
 }
 
+void common_chat_msg_parser::mark_reasoning_active(const std::string & end_tag) {
+    result_.reasoning_status.detected = true;
+    result_.reasoning_status.active   = true;
+    if (!end_tag.empty()) {
+        result_.reasoning_status.end_tag = end_tag;
+    }
+}
+
+void common_chat_msg_parser::mark_reasoning_closed() {
+    if (result_.reasoning_status.detected) {
+        result_.reasoning_status.active = false;
+    }
+}
+
 bool common_chat_msg_parser::add_tool_call(const std::string & name, const std::string & id, const std::string & arguments) {
     if (name.empty()) {
         return false;
@@ -329,11 +343,13 @@ bool common_chat_msg_parser::try_parse_reasoning(const std::string & start_think
     const size_t saved_pos = pos_;
     const size_t saved_content_size = result_.content.size();
     const size_t saved_reasoning_size = result_.reasoning_content.size();
+    const auto   saved_reasoning_status = result_.reasoning_status;
 
     auto restore_state = [&]() {
         move_to(saved_pos);
         result_.content.resize(saved_content_size);
         result_.reasoning_content.resize(saved_reasoning_size);
+        result_.reasoning_status = saved_reasoning_status;
     };
 
     // Allow leading whitespace to be preserved as content when reasoning is present at the start
@@ -370,9 +386,11 @@ bool common_chat_msg_parser::try_parse_reasoning(const std::string & start_think
         if (whitespace_end > pos_) {
             add_content(input_.substr(pos_, whitespace_end - pos_));
         }
+        mark_reasoning_active(end_think);
         set_reasoning_prefix(cursor);
         cursor += start_think.size();
     } else if (syntax_.thinking_forced_open) {
+        mark_reasoning_active(end_think);
         cursor = whitespace_end;
     } else {
         restore_state();
@@ -398,8 +416,10 @@ bool common_chat_msg_parser::try_parse_reasoning(const std::string & start_think
 
         if (end_pos > cursor) {
             handle_reasoning(input_.substr(cursor, end_pos - cursor), /* closed */ true);
+            mark_reasoning_closed();
         } else {
             handle_reasoning("", /* closed */ true);
+            mark_reasoning_closed();
         }
 
         cursor = end_pos + end_think.size();
@@ -420,6 +440,7 @@ bool common_chat_msg_parser::try_parse_reasoning(const std::string & start_think
                 move_to(input_.size());
                 return true;
             }
+            mark_reasoning_active(end_think);
             set_reasoning_prefix(cursor);
             cursor += start_think.size();
             continue;
diff --git a/common/chat-parser.h b/common/chat-parser.h
@@ -56,6 +56,10 @@ class common_chat_msg_parser {
     // Appends to the result.reasoning_content field
     void add_reasoning_content(const std::string & reasoning_content);
 
+    // Track reasoning status to expose start/end markers to callers
+    void mark_reasoning_active(const std::string & end_tag);
+    void mark_reasoning_closed();
+
     // Adds a tool call to the result. If the tool call is too incomplete (e.g. name empty), it won't add anything.
     bool add_tool_call(const std::string & name, const std::string & id, const std::string & arguments);
 
diff --git a/common/chat.h b/common/chat.h
@@ -22,6 +22,19 @@ struct common_chat_tool_call {
     }
 };
 
+struct common_chat_reasoning_status {
+    bool detected = false;      // a reasoning block start was observed
+    bool active   = false;      // we are currently inside a reasoning block (not closed yet)
+    std::string end_tag;        // closing tag to use when forcing a close
+
+    bool operator==(const common_chat_reasoning_status & other) const {
+        return detected == other.detected && active == other.active && end_tag == other.end_tag;
+    }
+    bool operator!=(const common_chat_reasoning_status & other) const {
+        return !(*this == other);
+    }
+};
+
 struct common_chat_msg_content_part {
     std::string type;
     std::string text;
@@ -37,6 +50,7 @@ struct common_chat_msg {
     std::vector<common_chat_msg_content_part> content_parts;
     std::vector<common_chat_tool_call> tool_calls;
     std::string reasoning_content;
+    common_chat_reasoning_status reasoning_status;
     std::string tool_name;
     std::string tool_call_id;
 
@@ -63,6 +77,7 @@ struct common_chat_msg {
             && content_parts == other.content_parts
             && tool_calls == other.tool_calls
             && reasoning_content == other.reasoning_content
+            && reasoning_status == other.reasoning_status
             && tool_name == other.tool_name
             && tool_call_id == other.tool_call_id;
     }
diff --git a/tests/test-chat-parser.cpp b/tests/test-chat-parser.cpp
@@ -119,6 +119,9 @@ static void test_reasoning() {
     auto msg = common_chat_parse(input, false, syntax);
     assert_equals(variant, std::string("Pense"), msg.reasoning_content);
     assert_equals(variant, std::string("Bonjour"), msg.content);
+    assert_equals(variant, true, msg.reasoning_status.detected);
+    assert_equals(variant, false, msg.reasoning_status.active);
+    assert_equals(variant, std::string("</think>"), msg.reasoning_status.end_tag);
   }
   {
     const std::string variant("llama_3_inline_think");
@@ -133,6 +136,9 @@ static void test_reasoning() {
     auto msg = common_chat_parse(input, false, syntax);
     assert_equals(variant, std::string("Plan"), msg.reasoning_content);
     assert_equals(variant, std::string("Réponse"), msg.content);
+    assert_equals(variant, true, msg.reasoning_status.detected);
+    assert_equals(variant, false, msg.reasoning_status.active);
+    assert_equals(variant, std::string("</think>"), msg.reasoning_status.end_tag);
   }
   // Test DeepSeek V3.1 parsing - reasoning content followed by "</think>" and then regular content
   {
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
@@ -34,11 +34,6 @@ using json = nlohmann::ordered_json;
 
 constexpr int HTTP_POLLING_SECONDS = 1;
 
-static const std::vector<std::pair<std::string, std::string>> kReasoningThinkMarkers = {
-    {"<think>", "</think>"},
-    {"<|START_THINKING|>", "<|END_THINKING|>"},
-};
-
 // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
 enum slot_state {
     SLOT_STATE_IDLE,
@@ -1109,15 +1104,8 @@ struct server_context_impl {
         // Initialize reasoning tracking
         slot.forced_tokens.clear();
         slot.n_reasoning_tokens = 0;
-
-        const bool thinking_forced_open = slot.task->params.oaicompat_chat_syntax.thinking_forced_open;
-        slot.reasoning = thinking_forced_open ? REASONING_STATE_REASONING : REASONING_STATE_NONE;
-
-        if (thinking_forced_open) {
-            slot.reasoning_end_tag = kReasoningThinkMarkers.front().second;
-        } else {
-            slot.reasoning_end_tag.clear();
-        }
+        slot.reasoning = REASONING_STATE_NONE;
+        slot.reasoning_end_tag.clear();
 
 
         SLT_INF(slot, "%s", "processing task\n");
@@ -1198,65 +1186,62 @@ struct server_context_impl {
         const int32_t reasoning_budget = (slot.task ? slot.task->params.reasoning_budget : params_base.reasoning_budget);
 
         // check reasoning budget limit
-        // Track reasoning tokens when we're inside thinking blocks (<think>...</think> or similar)
+        // Track reasoning tokens using the chat parser to detect reasoning segments consistently across formats
         // When the budget is exceeded we enqueue the closing tag tokens so they get sent to the client
         // and fed back into the model before continuing normal generation
         if (slot.has_next_token && reasoning_budget > 0 && slot.reasoning != REASONING_STATE_FINISHED)  {
-            // Check if we've entered or exited reasoning mode
-            if (slot.reasoning == REASONING_STATE_NONE) {
-                for (const auto & [start_tag, end_tag] : kReasoningThinkMarkers) {
-                    size_t start_pos = slot.generated_text.rfind(start_tag);
-                    if (start_pos != std::string::npos) {
-                        SLT_DBG(slot, "detected reasoning start with '%s'\n", start_tag.c_str());
-                        slot.reasoning = REASONING_STATE_REASONING;
-                        slot.reasoning_end_tag = end_tag;
-                        slot.n_reasoning_tokens = 0;
-                        break;
-                    }
-                }
-            } else if (slot.reasoning == REASONING_STATE_REASONING) {
-                size_t end_pos = slot.generated_text.rfind(slot.reasoning_end_tag);
-                if (end_pos != std::string::npos) {
-                    SLT_DBG(slot, "detected reasoning end with '%s'\n", slot.reasoning_end_tag.c_str());
-                    slot.reasoning = REASONING_STATE_FINISHED;
+            std::vector<common_chat_msg_diff> reasoning_diffs;
+            const auto & parsed_msg = slot.update_chat_msg(reasoning_diffs, /* compute_diffs = */ false);
+            const auto & rstatus = parsed_msg.reasoning_status;
+
+            if (rstatus.active && slot.reasoning != REASONING_STATE_PENDING_FORCE_CLOSE) {
+                if (slot.reasoning != REASONING_STATE_REASONING) {
+                    SLT_DBG(slot, "detected reasoning start via parser%s\n", "");
+                    slot.reasoning = REASONING_STATE_REASONING;
+                    slot.reasoning_end_tag = rstatus.end_tag;
                     slot.n_reasoning_tokens = 0;
-                } else {
-                    // If actively reasoning (and we haven't already scheduled a forced close) count this token
-                    slot.n_reasoning_tokens++;
+                }
+            } else if (!rstatus.active && slot.reasoning == REASONING_STATE_REASONING) {
+                SLT_DBG(slot, "detected reasoning end '%s' via parser\n", rstatus.end_tag.c_str());
+                slot.reasoning = REASONING_STATE_FINISHED;
+                slot.n_reasoning_tokens = 0;
+            }
 
-                    if (slot.n_reasoning_tokens >= reasoning_budget) {
-                        SLT_INF(slot, "reasoning budget exceeded, forcing close with '%s', n_reasoning_tokens = %d, reasoning_budget = %d\n",
-                            slot.reasoning_end_tag.c_str(), slot.n_reasoning_tokens, reasoning_budget);
+            if (slot.reasoning == REASONING_STATE_REASONING) {
+                slot.n_reasoning_tokens++;
 
-                        auto fail_close = [&](const char * reason) {
-                            SLT_WRN(slot, "failed to inject reasoning close tag (%s) -> stopping generation\n", reason);
-                            slot.stop           = STOP_TYPE_LIMIT;
-                            slot.has_next_token = false;
-                        };
+                if (slot.n_reasoning_tokens >= reasoning_budget) {
+                    SLT_INF(slot, "reasoning budget exceeded, forcing close with '%s', n_reasoning_tokens = %d, reasoning_budget = %d\n",
+                        slot.reasoning_end_tag.c_str(), slot.n_reasoning_tokens, reasoning_budget);
 
-                        if (slot.reasoning_end_tag.empty()) {
-                            fail_close("no closing tag detected");
-                        } else {
-                            const std::string forced_message = slot.task->params.reasoning_force_close_message.empty()
-                                ? std::string(COMMON_DEFAULT_REASONING_FORCE_CLOSE_MESSAGE)
-                                : slot.task->params.reasoning_force_close_message;
-                            const std::string forced_injection = forced_message + slot.reasoning_end_tag;
-
-                            llama_tokens closing_tokens;
-                            try {
-                                closing_tokens = common_tokenize(ctx, forced_injection, /*add_special=*/false, /*parse_special=*/true);
-                            } catch (const std::exception & err) {
-                                SLT_WRN(slot, "tokenization error while forcing reasoning close: %s\n", err.what());
-                                fail_close("tokenization error");
-                                closing_tokens.clear();
-                            }
+                    auto fail_close = [&](const char * reason) {
+                        SLT_WRN(slot, "failed to inject reasoning close tag (%s) -> stopping generation\n", reason);
+                        slot.stop           = STOP_TYPE_LIMIT;
+                        slot.has_next_token = false;
+                    };
 
-                            if (!closing_tokens.empty()) {
-                                slot.forced_tokens.insert(slot.forced_tokens.end(), closing_tokens.begin(), closing_tokens.end());
-                                slot.reasoning = REASONING_STATE_PENDING_FORCE_CLOSE;
-                            } else if (slot.has_next_token) {
-                                fail_close("closing tag produced no tokens");
-                            }
+                    if (slot.reasoning_end_tag.empty()) {
+                        fail_close("no closing tag detected");
+                    } else {
+                        const std::string forced_message = slot.task->params.reasoning_force_close_message.empty()
+                            ? std::string(COMMON_DEFAULT_REASONING_FORCE_CLOSE_MESSAGE)
+                            : slot.task->params.reasoning_force_close_message;
+                        const std::string forced_injection = forced_message + slot.reasoning_end_tag;
+
+                        llama_tokens closing_tokens;
+                        try {
+                            closing_tokens = common_tokenize(ctx, forced_injection, /*add_special=*/false, /*parse_special=*/true);
+                        } catch (const std::exception & err) {
+                            SLT_WRN(slot, "tokenization error while forcing reasoning close: %s\n", err.what());
+                            fail_close("tokenization error");
+                            closing_tokens.clear();
+                        }
+
+                        if (!closing_tokens.empty()) {
+                            slot.forced_tokens.insert(slot.forced_tokens.end(), closing_tokens.begin(), closing_tokens.end());
+                            slot.reasoning = REASONING_STATE_PENDING_FORCE_CLOSE;
+                        } else if (slot.has_next_token) {
+                            fail_close("closing tag produced no tokens");
                         }
                     }
                 }

Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,9 @@ static void test_reasoning() {`
`119`	`119`	`auto msg = common_chat_parse(input, false, syntax);`
`120`	`120`	`assert_equals(variant, std::string("Pense"), msg.reasoning_content);`
`121`	`121`	`assert_equals(variant, std::string("Bonjour"), msg.content);`
	`122`	`+ assert_equals(variant, true, msg.reasoning_status.detected);`
	`123`	`+ assert_equals(variant, false, msg.reasoning_status.active);`
	`124`	`+ assert_equals(variant, std::string("</think>"), msg.reasoning_status.end_tag);`
`122`	`125`	`}`
`123`	`126`	`{`
`124`	`127`	`const std::string variant("llama_3_inline_think");`
`@@ -133,6 +136,9 @@ static void test_reasoning() {`
`133`	`136`	`auto msg = common_chat_parse(input, false, syntax);`
`134`	`137`	`assert_equals(variant, std::string("Plan"), msg.reasoning_content);`
`135`	`138`	`assert_equals(variant, std::string("Réponse"), msg.content);`
	`139`	`+ assert_equals(variant, true, msg.reasoning_status.detected);`
	`140`	`+ assert_equals(variant, false, msg.reasoning_status.active);`
	`141`	`+ assert_equals(variant, std::string("</think>"), msg.reasoning_status.end_tag);`
`136`	`142`	`}`
`137`	`143`	`// Test DeepSeek V3.1 parsing - reasoning content followed by "</think>" and then regular content`
`138`	`144`	`{`