From 50d2d685566b7bc0608f3bfda7821dd21cfde41b Mon Sep 17 00:00:00 2001 From: jiangjiang248 Date: Mon, 28 Sep 2026 16:38:12 +0800 Subject: [PATCH] fix(llm): avoid quadratic think-block stripping on unclosed tags (#5854) * fix(llm): avoid quadratic think-block stripping on unclosed tags * test(llm): pin think tag parsing edge cases --------- Co-authored-by: Willem Jiang --- .../packages/harness/deerflow/utils/AGENTS.md | 6 +++ .../harness/deerflow/utils/llm_text.py | 47 +++++++++++++++---- backend/tests/test_utils_llm_text.py | 28 +++++++++++ 3 files changed, 71 insertions(+), 10 deletions(-) diff --git a/backend/packages/harness/deerflow/utils/AGENTS.md b/backend/packages/harness/deerflow/utils/AGENTS.md index 0bc17cbca8..bc5d586b12 100644 --- a/backend/packages/harness/deerflow/utils/AGENTS.md +++ b/backend/packages/harness/deerflow/utils/AGENTS.md @@ -8,6 +8,12 @@ which takes a whole message and uses different list/mapping semantics. Regression coverage lives in `tests/test_utils_messages.py`, `tests/test_subagent_executor.py`, and `tests/test_task_continuity.py`. +`llm_text.py::strip_think_blocks` scans opening and closing tags forward so +repeated unclosed `` prefixes in model output do not rescan the same +suffix. Keep its case-insensitive tag handling, optional whitespace before +the closing `>`, and the `truncate_unclosed` behavior. Regression coverage +lives in `tests/test_utils_llm_text.py`. + ### Agent / Tool Assembly Off-Load Tool and agent assembly re-enters `get_available_tools()` and may block on MCP discovery, so the four async assembly entry points — `run_agent`'s `agent_factory` call, `task_tool`, durable batch `_execute_item`, and `abuild_checkpoint_state_accessor` — dispatch through `deerflow.utils.assembly_io.run_assembly`, a dedicated ContextVar-preserving bounded executor (`DEER_FLOW_ASSEMBLY_WORKERS`, default 8) rather than the loop's default executor; a hung MCP server therefore parks an assembly worker instead of queueing unrelated default-executor work, and the pool logs a warning when pending assemblies exceed the worker count. `tests/blocking_io/test_tool_assembly_offloop.py` pins all four offloads plus the ContextVar propagation. diff --git a/backend/packages/harness/deerflow/utils/llm_text.py b/backend/packages/harness/deerflow/utils/llm_text.py index 625cb41aa5..9dd1ec24db 100644 --- a/backend/packages/harness/deerflow/utils/llm_text.py +++ b/backend/packages/harness/deerflow/utils/llm_text.py @@ -4,10 +4,31 @@ from __future__ import annotations import re -# Matches a complete ... block (case-insensitive, spans newlines). -_THINK_BLOCK_RE = re.compile(r"]*>.*?", re.IGNORECASE | re.DOTALL) -# Matches a dangling, unclosed (model truncated at max_tokens mid-thought). -_OPEN_THINK_RE = re.compile(r"]*>", re.IGNORECASE) +_THINK_OPEN_PREFIX_RE = re.compile(r" tuple[int, int] | None: + """Find the next complete opening tag without retrying its suffix at each prefix.""" + match = _THINK_OPEN_PREFIX_RE.search(text, start) + if match is None: + return None + end = text.find(">", match.end()) + if end < 0: + return None + return match.start(), end + 1 + + +def _find_think_close(text: str, start: int) -> tuple[int, int] | None: + """Find the first closing tag, allowing whitespace before its final ``>``.""" + while match := _THINK_CLOSE_PREFIX_RE.search(text, start): + end = match.end() + while end < len(text) and text[end].isspace(): + end += 1 + if end < len(text) and text[end] == ">": + return match.start(), end + 1 + start = match.end() + return None def strip_think_blocks(text: str, *, truncate_unclosed: bool = True) -> str: @@ -22,12 +43,18 @@ def strip_think_blocks(text: str, *, truncate_unclosed: bool = True) -> str: draft that mentions the tag) pass ``truncate_unclosed=False`` so the tag is preserved instead of silently discarding the rest of the text. """ - text = _THINK_BLOCK_RE.sub("", text) - if truncate_unclosed: - open_match = _OPEN_THINK_RE.search(text) - if open_match: - text = text[: open_match.start()] - return text.strip() + parts: list[str] = [] + start = 0 + while (opening := _find_think_open(text, start)) is not None: + closing = _find_think_close(text, opening[1]) + if closing is None: + if truncate_unclosed: + return ("".join(parts) + text[start : opening[0]]).strip() + break + parts.append(text[start : opening[0]]) + start = closing[1] + parts.append(text[start:]) + return "".join(parts).strip() def strip_markdown_code_fence(text: str) -> str: diff --git a/backend/tests/test_utils_llm_text.py b/backend/tests/test_utils_llm_text.py index 8a0d72ada2..ea85026b3e 100644 --- a/backend/tests/test_utils_llm_text.py +++ b/backend/tests/test_utils_llm_text.py @@ -2,6 +2,10 @@ from __future__ import annotations +import time + +import pytest + from deerflow.utils.llm_text import ( extract_response_text, strip_markdown_code_fence, @@ -65,6 +69,30 @@ def test_strip_think_blocks_removes_complete_and_keeps_dangling_when_disabled() assert result == "keeptrunc" +def test_strip_think_blocks_does_not_treat_junk_close_prefix_as_closing_tag() -> None: + text = "axb" + assert strip_think_blocks(text) == "a" + assert strip_think_blocks(text, truncate_unclosed=False) == text + + +def test_strip_think_blocks_accepts_multiline_whitespace_before_close() -> None: + assert strip_think_blocks("xok") == "ok" + + +def test_strip_think_blocks_consumes_multiple_opens_through_first_close() -> None: + assert strip_think_blocks("abc") == "c" + + +@pytest.mark.parametrize("tag", ["", " None: + text = "answer" + tag * 12_000 + expected = "answer" if tag.endswith(">") and truncate_unclosed else text + start = time.perf_counter() + assert strip_think_blocks(text, truncate_unclosed=truncate_unclosed) == expected + assert time.perf_counter() - start < 1.0 + + # --------------------------------------------------------------------------- # strip_markdown_code_fence # ---------------------------------------------------------------------------