Studio: strip orphan tool_call XML leaking into visible content (#5735)
* Studio: strip orphan tool_call XML from streamed visible content
The speculative-buffer state machine in
`studio/backend/core/inference/llama_cpp.py` can slice a tool_call XML
block between the silent DRAINING path and the user-visible
content_accum, depending on when in the model's emission the BUFFERING
-> STREAMING -> DRAINING transitions fire. Three leak shapes were
observed in a 2026-05-22 sweep of 900 Qwen3.5 / Qwen3.6 GGUF runs:
Pre-fix XML leak rate: 20/900 (2.22%), concentrated 6.7% on the
larger Q8 / MTP configs:
Qwen3.6-35B-A3B Q8_0 4/60 (6.7%)
Qwen3.6-35B-A3B-MTP Q4 4/60 (6.7%)
Qwen3.5-35B-A3B Q8_0 3/60 (5.0%)
Qwen3.6-27B Q8_0 3/60 (5.0%)
The existing `_TOOL_XML_RE` only matched well-formed
`<tool_call>...</tool_call>` and `<function=...></function>` pairs, so
unterminated openings (close was DRAINED) and orphan closes (opening
was DRAINED) survived the strip and reached the user.
Fix relaxes the regex to also strip:
1. Orphan opening up to end-of-string: `(?:</tool_call>|\Z)`
2. Orphan closing tag: bare `</tool_call>` / `</function>`
Verified on the full sweep: 20/900 -> 0/900 (100% of detected leaks
eliminated). 16 unit tests in `test_tool_xml_strip.py` pin all three
leak shapes plus the well-formed cases, plus parametrised checks on
the 5 actual real-world leak samples from the sweep data.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: strip tail-only </parameter> orphan + tighten regex
The 2026-05-22 gdpval sweep surfaced a 4th XML-leak shape not caught
by the earlier regex: a bare `</parameter>\n\n` at end-of-buffer (7
of 192 trials, all Qwen3.5-27B + a few Qwen3.6-27B). The model emits
the full `<tool_call><function=...><parameter=...>...content...
</parameter></function></tool_call>` envelope, the speculative buffer
DRAINS the opening tags as intended, but EOS (max_tokens cutoff)
truncates the outer `</function></tool_call>` close, leaving just
`</parameter>` as the visible tail.
We strip this ONLY when end-anchored (`\s*\Z`) so legitimate
mid-text uses (user code samples, documentation discussing the
Qwen tool-call XML shape) survive. Verified on the 192-trial
gdpval corpus: before=7, after=0.
While at it, fold the five top-level alternations into three by
sharing tag-name and prefix subgroups:
<tool_call>... + <function=\w+>... + --> <(?:tool_call|function=\w+)>...
</tool_call> | </function> --> </(?:tool_call|function)>
Semantically identical (verified by replay over the 192-trial
corpus + adversarial inputs, 0 diffs) and 1.34x faster on real
workloads. Backtracking-safety pinned by two new perf guards
(256KB '<' spam, 1000x orphan opens).
Tests: 16 -> 28 (6 new functional + 4 well-formed-vs-orphan +
2 perf guards).
* Tighten comments in XML-strip regex and tests
Code says what it does; comments were repeating it. Strip the verbose
explanations down to the WHY-only bits (engine quirk, tail-anchor
rationale, real-world source of each test sample). No code changes.
inference.py: 21 -> 12 lines around _TOOL_XML_RE
test_tool_xml_strip.py: 343 -> 259 lines (-84)
Tests: 28/28 still pass.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
56e9046b2f
commit
f7f540a58b
2 changed files with 273 additions and 2 deletions
|
|
@ -427,9 +427,17 @@ _TOOL_ACTION_NUDGE = (
|
|||
" Do NOT output code blocks -- use the python tool instead."
|
||||
)
|
||||
|
||||
# Regex for stripping leaked tool-call XML from assistant messages/stream
|
||||
# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py
|
||||
# split across the visible/DRAIN boundary. Four leak shapes:
|
||||
# 1. well-formed `<tool_call>...</tool_call>` / `<function=...>...</function>`
|
||||
# 2. orphan opening to EOF (close was DRAINED)
|
||||
# 3. bare orphan close (open was DRAINED)
|
||||
# 4. tail-only `</parameter>` (outer close truncated by EOS); anchored to
|
||||
# `\Z` so mid-text `<parameter>` in user code samples survives.
|
||||
_TOOL_XML_RE = _re.compile(
|
||||
r"<tool_call>.*?</tool_call>|<function=\w+>.*?</function>",
|
||||
r"<(?:tool_call|function=\w+)>.*?(?:</(?:tool_call|function)>|\Z)"
|
||||
r"|</(?:tool_call|function)>"
|
||||
r"|</parameter>\s*\Z",
|
||||
_re.DOTALL,
|
||||
)
|
||||
logger = get_logger(__name__)
|
||||
|
|
|
|||
263
studio/backend/tests/test_tool_xml_strip.py
Normal file
263
studio/backend/tests/test_tool_xml_strip.py
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Tests for `_TOOL_XML_RE` (routes/inference.py) -- strips tool-call
|
||||
XML that leaks past the speculative buffer in core/inference/llama_cpp.py
|
||||
when the open/close pair is split across the visible/DRAIN boundary.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import types as _types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
||||
if _BACKEND_DIR not in sys.path:
|
||||
sys.path.insert(0, _BACKEND_DIR)
|
||||
|
||||
# Extract the regex from source (routes module needs heavy stubbing to import).
|
||||
import re as _re
|
||||
|
||||
_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text()
|
||||
_m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL)
|
||||
assert _m, "could not extract _TOOL_XML_RE source"
|
||||
_ns = {"_re": _re}
|
||||
exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
|
||||
_TOOL_XML_RE = _ns["_TOOL_XML_RE"]
|
||||
|
||||
|
||||
# ── Well-formed pairs ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_strips_well_formed_tool_call():
|
||||
text = (
|
||||
"Let me search.\n"
|
||||
"<tool_call>\n"
|
||||
"<function=web_search>\n"
|
||||
"<parameter=query>\nBillboard 2015\n</parameter>\n"
|
||||
"</function>\n"
|
||||
"</tool_call>\n"
|
||||
"Here are the songs:"
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<tool_call>" not in cleaned
|
||||
assert "<function=" not in cleaned
|
||||
assert "</tool_call>" not in cleaned
|
||||
assert "</function>" not in cleaned
|
||||
assert "Here are the songs:" in cleaned, "non-XML content must survive"
|
||||
assert "Let me search." in cleaned
|
||||
|
||||
|
||||
def test_strips_function_only_well_formed():
|
||||
text = "Setup.\n<function=python>\n<parameter=code>\nprint(1)\n</parameter>\n</function>\nDone."
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<function=" not in cleaned
|
||||
assert "Setup." in cleaned
|
||||
assert "Done." in cleaned
|
||||
|
||||
|
||||
# ── Orphan openings ───────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_strips_orphan_tool_call_no_close():
|
||||
text = (
|
||||
"Reasoning.\n</think>"
|
||||
"<tool_call>\n"
|
||||
"<function=web_search>\n"
|
||||
"<parameter=query>\nBillboard 2015\n</parameter>\n"
|
||||
"</function"
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<tool_call>" not in cleaned
|
||||
assert "<function=" not in cleaned
|
||||
assert "Reasoning." in cleaned
|
||||
|
||||
|
||||
def test_strips_orphan_function_no_close():
|
||||
text = (
|
||||
"I'll call python:\n<function=python>\n<parameter=code>\nprint(1)\n</parameter>"
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<function=" not in cleaned
|
||||
assert "I'll call python:" in cleaned
|
||||
|
||||
|
||||
def test_strips_orphan_only_opening_tag():
|
||||
cleaned = _TOOL_XML_RE.sub("", "Search starting.\n<tool_call>")
|
||||
assert "<tool_call>" not in cleaned
|
||||
assert "Search starting." in cleaned
|
||||
|
||||
|
||||
def test_strips_multiple_orphans():
|
||||
text = (
|
||||
"First call:\n<tool_call>\n<function=python>\n<parameter=code>\nx=1\n"
|
||||
"Second call:\n<function=web_search>\n<parameter=query>\nhi\n"
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<tool_call>" not in cleaned
|
||||
assert "<function=" not in cleaned
|
||||
|
||||
|
||||
# ── Orphan closes ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_strips_orphan_closing_tag():
|
||||
# Real shape from Qwen3.6-27B Q8 sweep (open got DRAINED, close leaked).
|
||||
text = "...the table rows directly.\n</parameter>\n</function>\n</tool_call><think>Continuing</think>"
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "</tool_call>" not in cleaned
|
||||
assert "</function>" not in cleaned
|
||||
# Mid-string </parameter> intentionally preserved (see preserve test).
|
||||
|
||||
|
||||
# ── Tail-only </parameter> (PR #5735 follow-up) ───────────────────
|
||||
|
||||
|
||||
def test_strips_tail_only_parameter_orphan():
|
||||
# Outer </function></tool_call> truncated by EOS, inner <parameter=...> DRAINED.
|
||||
cleaned = _TOOL_XML_RE.sub("", "and the text is not readable.\n</parameter>\n\n")
|
||||
assert "</parameter>" not in cleaned
|
||||
assert "and the text is not readable." in cleaned
|
||||
|
||||
|
||||
def test_strips_tail_only_parameter_orphan_single_newline():
|
||||
cleaned = _TOOL_XML_RE.sub("", "Global Economic Prospects\n</parameter>\n")
|
||||
assert "</parameter>" not in cleaned
|
||||
assert "Global Economic Prospects" in cleaned
|
||||
|
||||
|
||||
def test_strips_tail_only_parameter_orphan_no_trailing_ws():
|
||||
cleaned = _TOOL_XML_RE.sub("", "Final answer.</parameter>")
|
||||
assert "</parameter>" not in cleaned
|
||||
assert "Final answer." in cleaned
|
||||
|
||||
|
||||
def test_preserves_mid_string_parameter_in_code_sample():
|
||||
# Tail-anchor on `</parameter>` is required so doc/example prose survives.
|
||||
text = (
|
||||
"Here is the Qwen tool-call format:\n"
|
||||
"```xml\n"
|
||||
"<tool_call><function=foo><parameter=arg>value</parameter></function></tool_call>\n"
|
||||
"```\n"
|
||||
"Note the closing </parameter> sits inside <function>."
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "Note the closing </parameter> sits inside" in cleaned
|
||||
|
||||
|
||||
def test_strips_well_formed_then_orphan():
|
||||
text = (
|
||||
"Round one:\n<tool_call>\n<function=python>\n<parameter=code>\n1\n"
|
||||
"</parameter>\n</function>\n</tool_call>\n"
|
||||
"Now round two:\n<tool_call>\n<function=web_search>\n<parameter=query>\n"
|
||||
"what is X\n</parameter>\n</function"
|
||||
)
|
||||
cleaned = _TOOL_XML_RE.sub("", text)
|
||||
assert "<tool_call>" not in cleaned
|
||||
assert "<function=" not in cleaned
|
||||
assert "Round one:" in cleaned
|
||||
assert "Now round two:" in cleaned
|
||||
|
||||
|
||||
# ── Preservation (no false positives) ────────────────────────────
|
||||
|
||||
|
||||
def test_preserves_plain_text():
|
||||
text = "1. Animals — Maroon 5\n2. Take Me to Church — Hozier"
|
||||
assert _TOOL_XML_RE.sub("", text) == text
|
||||
|
||||
|
||||
def test_preserves_code_fences():
|
||||
text = "```python\nimport sys\nprint(sys.version)\n```"
|
||||
assert _TOOL_XML_RE.sub("", text) == text
|
||||
|
||||
|
||||
def test_preserves_html_in_prose():
|
||||
text = "Use the <html> tag for documents."
|
||||
assert _TOOL_XML_RE.sub("", text) == text
|
||||
|
||||
|
||||
# ── Real-world leak samples from the 2026-05-22 sweep ────────────
|
||||
|
||||
|
||||
REAL_LEAKS = [
|
||||
# Qwen3.5-35B-A3B UD-Q4_K_XL billboard s22 -- orphan open
|
||||
'rectly.\n\nLet me try searching for Wikipedia pages that might have weekly chart data for 2015.\n</think><tool_call>\n<function=web_search>\n<parameter=query>\n"Billboard Hot 100" "2015" "weekly" "chart" "position" "3"\n</parameter>\n</function',
|
||||
# Qwen3.6-27B UD-Q2_K_XL billboard s14 -- orphan open
|
||||
'arch `site:wikipedia.org "peaked at number 3" "2015" Billboard`\nI\'ll do a quick web search.\n</think><tool_call>\n<function=web_search>\n<parameter=query>\n"peaked at number 3" Billboard Hot 100 2015 list\n</parameter>\n</function',
|
||||
# Qwen3.6-27B UD-Q2_K_XL billboard s15 -- orphan open
|
||||
'rd Hot 100 top-ten singles in 2015".\nI\'ll use web_search to find this exact Wikipedia page.\n</think><tool_call>\n<function=web_search>\n<parameter=query>\n"List of Billboard Hot 100 top-ten singles in 2015" wikipedia\n</parameter>\n</function',
|
||||
# Qwen3.6-27B Q8_0 billboard s02 -- orphan close
|
||||
"the table rows directly.\n</parameter>\n</function>\n</tool_call><think>The user wants me to list and categorize all songs that charted #3 on the Billboard Hot 100 in 2015. I have been trying to get this data",
|
||||
# Qwen3.6-35B-A3B Q8_0 billboard s21 -- orphan close
|
||||
"parse it more carefully.\n</parameter>\n</function>\n</tool_call><think>The user wants a list of songs that charted #3 on the Billboard Hot 100 in 2015, categorized.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"leak", REAL_LEAKS, ids = [f"sweep_sample_{i}" for i in range(len(REAL_LEAKS))]
|
||||
)
|
||||
def test_real_world_sweep_leaks_get_stripped(leak):
|
||||
cleaned = _TOOL_XML_RE.sub("", leak)
|
||||
assert "<tool_call>" not in cleaned, f"leak survived: {cleaned!r}"
|
||||
assert "<function=" not in cleaned, f"leak survived: {cleaned!r}"
|
||||
|
||||
|
||||
# ── Real-world tail-only </parameter> from gdpval sweep ──────────
|
||||
|
||||
|
||||
# All end-anchored: outer </function></tool_call> truncated by EOS,
|
||||
# inner <parameter=...> open DRAINED, leaving bare </parameter> tail.
|
||||
GDPVAL_PARAMETER_LEAKS = [
|
||||
# Qwen3.5-27B Q8_0 / worldbank s00
|
||||
"the page contains image data and the text is not readable.\n</parameter>\n\n",
|
||||
# Qwen3.5-27B Q8_0 / worldbank s42 (preceded by mojibake)
|
||||
"...some mojibake content here...\n</parameter>\n\n",
|
||||
# Qwen3.5-27B UD-Q4_K_XL / coppa s07
|
||||
"blocked, while others may still be in effect. The law is currently under further review by the Ninth Circuit.\n</parameter>\n\n",
|
||||
# Qwen3.5-27B UD-Q4_K_XL / police_training s00
|
||||
"comprehensive training report\n</parameter>\n\n",
|
||||
# Qwen3.5-27B UD-Q4_K_XL / worldbank s00
|
||||
"Global Economic Prospects\nJune 2025\nGlobal Economic Prospects\n</parameter>\n",
|
||||
# Qwen3.6-27B Q8_0 / overpass s07
|
||||
"Let me create a comprehensive query and instructions document.\n</parameter>\n\n",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"leak",
|
||||
GDPVAL_PARAMETER_LEAKS,
|
||||
ids = [f"gdpval_param_orphan_{i}" for i in range(len(GDPVAL_PARAMETER_LEAKS))],
|
||||
)
|
||||
def test_gdpval_parameter_orphans_get_stripped(leak):
|
||||
cleaned = _TOOL_XML_RE.sub("", leak)
|
||||
assert "</parameter>" not in cleaned, f"leak survived: {cleaned!r}"
|
||||
|
||||
|
||||
# ── Backtracking guards ──────────────────────────────────────────
|
||||
|
||||
|
||||
def test_no_catastrophic_backtracking_on_open_bracket_spam():
|
||||
# 256KB of '<' must fail fast (literal mismatch char 2), not backtrack.
|
||||
import time
|
||||
|
||||
adv = "<" * (1024 * 256) + "X"
|
||||
t0 = time.perf_counter()
|
||||
_TOOL_XML_RE.sub("", adv)
|
||||
elapsed = time.perf_counter() - t0
|
||||
assert elapsed < 0.5, f"regex took {elapsed*1000:.0f}ms on 256KB '<' spam"
|
||||
|
||||
|
||||
def test_no_catastrophic_backtracking_on_orphan_opening_spam():
|
||||
# 1000 unclosed openings: first alt must consume them all greedily.
|
||||
import time
|
||||
|
||||
adv = "<tool_call>X" * 1000
|
||||
t0 = time.perf_counter()
|
||||
cleaned = _TOOL_XML_RE.sub("", adv)
|
||||
elapsed = time.perf_counter() - t0
|
||||
assert elapsed < 0.1, f"regex took {elapsed*1000:.0f}ms on 1000x orphan opens"
|
||||
assert "<tool_call>" not in cleaned
|
||||
Loading…
Add table
Add a link
Reference in a new issue