From a3115bf68643cd28d356f0a8e6c8f99ed4b41067 Mon Sep 17 00:00:00 2001
From: 0xallam <ahmed39652003@gmail.com>
Date: Sun, 3 May 2026 15:55:40 -0700
Subject: [PATCH] strip thinking blocks in place so streaming detection uses
 clean content
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Previously the streaming loop computed a separate check_content for the
end-tag presence check but still resolved end_tag and find() against the
unstripped accumulated, so a <thinking> block containing </function>
while the real output used </invoke> (or vice-versa) caused the wrong
tag to be picked and truncation to land inside the thinking block.

The post-loop strip also ran after _truncate_to_first_function, which
counts every <function= / <invoke name= occurrence — so a thinking
block mentioning a tool call ("I should call <function=read_file>...")
became the "first" function and the real tool call got sliced off.

Move the strip into the streaming loop and apply it to accumulated
itself; the post-loop pipeline now operates on already-clean content
and the redundant trailing re.sub is dropped.
---
 strix/llm/llm.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/strix/llm/llm.py b/strix/llm/llm.py
index d166dfd4..fcd2ae21 100644
--- a/strix/llm/llm.py
+++ b/strix/llm/llm.py
@@ -26,6 +26,8 @@ from strix.utils.resource_paths import get_strix_resource_path
 litellm.drop_params = True
 litellm.modify_params = True
 
+_THINKING_BLOCK_RE = re.compile(r"<thinking[^>]*>.*?</thinking>", re.DOTALL)
+
 
 class LLMRequestFailedError(Exception):
     def __init__(self, message: str, details: str | None = None):
@@ -188,11 +190,8 @@ class LLM:
                 continue
             delta = self._get_chunk_content(chunk)
             if delta:
-                accumulated += delta
-                check_content = re.sub(
-                    r"<thinking[^>]*>.*?</thinking>", "", accumulated, flags=re.DOTALL
-                )
-                if "</function>" in check_content or "</invoke>" in check_content:
+                accumulated = _THINKING_BLOCK_RE.sub("", accumulated + delta)
+                if "</function>" in accumulated or "</invoke>" in accumulated:
                     end_tag = "</function>" if "</function>" in accumulated else "</invoke>"
                     pos = accumulated.find(end_tag)
                     accumulated = accumulated[: pos + len(end_tag)]
@@ -207,8 +206,6 @@ class LLM:
         accumulated = normalize_tool_format(accumulated)
         accumulated = fix_incomplete_tool_call(_truncate_to_first_function(accumulated))
 
-        accumulated = re.sub(r"<thinking[^>]*>.*?</thinking>", "", accumulated, flags=re.DOTALL)
-
         yield LLMResponse(
             content=accumulated,
             tool_invocations=parse_tool_invocations(accumulated),