f7f540a58b* 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>
263 lines
10 KiB
Python
263 lines
10 KiB
Python
# 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
|