From 3c809f44383bcebc318b8cddd8c6d11aaa861e9b Mon Sep 17 00:00:00 2001
From: Tarek Dakhran <tarek@liquid.ai>
Date: Fri, 30 Oct 2026 01:27:50 +0100
Subject: [PATCH] common : end reasoning budget forcing on an end tag from the
 prompt (#31374)

---
 common/reasoning-budget.cpp     |  3 ++-
 tests/test-reasoning-budget.cpp | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp
index 4884299f3..81b7c3e56 100644
--- a/common/reasoning-budget.cpp
+++ b/common/reasoning-budget.cpp
@@ -136,7 +136,8 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to
             // track the end sequence within forced_tokens so it is also reported on DONE
             const int32_t match = ctx->end_matcher.advance(token);
             ctx->force_pos++;
-            if (ctx->force_pos >= ctx->forced_tokens.size()) {
+            // an end tag that is already in the prompt (e.g. "<think></think>") also ends forcing
+            if (match >= 0 || ctx->force_pos >= ctx->forced_tokens.size()) {
                 ctx->state = REASONING_BUDGET_DONE;
                 ctx->end_match = match;
                 COM_TRC("%s", "forced sequence complete, done\n");
diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp
index 3bcc77e17..a50a10824 100644
--- a/tests/test-reasoning-budget.cpp
+++ b/tests/test-reasoning-budget.cpp
@@ -316,6 +316,23 @@ static void test_reasoning_budget_end_match() {
         llama_sampler_free(sampler);
     }
 
+    // budget=0 and the prompt already closes the block: the end tag ends forcing before the forced sequence
+    {
+        const std::vector<llama_tokens> end_nl = {{103, 101}, {101}};
+
+        auto * sampler = common_reasoning_budget_init(nullptr, start, end_nl, {103, 101}, 0, REASONING_BUDGET_IDLE);
+
+        llama_sampler_accept(sampler, 100); // FORCING
+        llama_sampler_accept(sampler, 101); // end tag from the prompt, DONE
+        GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE);
+
+        const llama_tokens * matched = common_reasoning_budget_get_end_match(sampler);
+        GGML_ASSERT(matched != nullptr);
+        GGML_ASSERT(*matched == llama_tokens({101}));
+
+        llama_sampler_free(sampler);
+    }
+
     // forced_tokens not ending with a known end sequence records nothing
     {
         auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING);