From 1f97870fa9ab3a9b6022262131039c8c2c3a3517 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Sun, 21 Jun 2026 19:58:33 +0800
Subject: [PATCH] Fix think tag streaming in agent output

---
 app/agent/__init__.py              |  9 +++++----
 tests/test_agent_tool_streaming.py | 26 ++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/app/agent/__init__.py b/app/agent/__init__.py
index 5c3f0055..6713c9c1 100644
--- a/app/agent/__init__.py
+++ b/app/agent/__init__.py
@@ -192,10 +192,11 @@ class _ThinkTagStripper:
                             self.buffer = self.buffer[-i:]
                             partial_match = True
                             break
-                    if not partial_match:
-                        on_output(self.buffer)
-                        emitted = True
-                        self.buffer = ""
+                    if partial_match:
+                        break
+                    on_output(self.buffer)
+                    emitted = True
+                    self.buffer = ""
             else:
                 end_idx = self.buffer.find("</think>")
                 if end_idx != -1:
diff --git a/tests/test_agent_tool_streaming.py b/tests/test_agent_tool_streaming.py
index f7adc414..d408b35a 100644
--- a/tests/test_agent_tool_streaming.py
+++ b/tests/test_agent_tool_streaming.py
@@ -7,6 +7,7 @@ import langchain.agents as langchain_agents
 if not hasattr(langchain_agents, "create_agent"):
     langchain_agents.create_agent = lambda *args, **kwargs: None
 
+from app.agent import _ThinkTagStripper
 from app.agent.callback import StreamingHandler
 from app.agent.middleware.subagents import is_subagent_stream_metadata
 from app.agent.tools.base import MoviePilotTool
@@ -17,6 +18,31 @@ from app.schemas.message import MessageResponse
 from app.schemas.types import MessageChannel, NotificationType
 
 
+def test_think_tag_stripper_waits_for_partial_open_tag():
+    """普通文本后出现不完整 think 开始标签时不应进入死循环。"""
+    stripper = _ThinkTagStripper()
+    outputs = []
+
+    emitted = stripper.process("你好<", outputs.append)
+    emitted_next = stripper.process("世界", outputs.append)
+
+    assert emitted is True
+    assert emitted_next is True
+    assert outputs == ["你好", "<世界"]
+
+
+def test_think_tag_stripper_hides_split_think_tag_content():
+    """think 标签被拆分到多个 token 时应继续隐藏思考内容。"""
+    stripper = _ThinkTagStripper()
+    outputs = []
+
+    stripper.process("回答前<", outputs.append)
+    stripper.process("thi", outputs.append)
+    stripper.process("nk>隐藏内容</think>回答后", outputs.append)
+
+    assert outputs == ["回答前", "回答后"]
+
+
 class DummyTool(MoviePilotTool):
     """用于流式输出测试的固定结果工具。"""