eb1ef44255* Quote-aware Gemma strip, symmetric unstarted cleanup, ReDoS anchor
Address review findings on the tool-strip and streaming paths:
- strip_tool_call_markup stripped Gemma-native spans with a plain regex that
stops at the first <tool_call|>, so a literal close marker inside a
<|"|>-quoted argument truncated the span and leaked its suffix into visible
text. A brace/quote-aware _strip_gemma_native_spans now removes complete
spans (keeping an incomplete one unless final), matching the parser's own
balance logic.
- The Gemma close pattern this PR added (<\|tool_call>.*?<tool_call\|>) had no
\Z fallback, so a run of unclosed markers backtracked from every open
position (quadratic, and the streaming stripper re-scans per token). It is
now anchored to (?:<tool_call|>|\Z) like routes/inference.py's _TOOL_XML_RE,
linear with identical output on well-formed input.
- _SameTaskStreamingResponse added unstarted_cleanup for the OpenAI passthrough,
but the local GGUF/safetensors streams that enter _TrackedCancel before
returning only unregister in the generator finally, which never runs if the
client disconnects before the body iterator starts, leaking cancel-registry
entries. Each such stream now passes unstarted_cleanup to exit its tracker.
- __call__ reads _unstarted_cleanup via getattr so a response built through
__new__ (the cancel-timing test) without __init__ does not raise
AttributeError; the test also sets the attribute explicitly.
- Document that the verbatim /v1/chat/completions passthrough delegates
<think>/<|tool_call> splitting to llama-server (--jinja, --reasoning-format
auto) and is intentionally not re-parsed locally, noting the llama.cpp
dependency.
Adds a regression test for the close-marker-inside-quoted-argument strip.
* Tighten comments on the tool-strip and streaming paths
Compress the verbose comment blocks added with the Gemma tool-call / streaming
work to crisp one or two liners, drop restatements of obvious code, and shorten
docstrings, keeping the load-bearing rationale (ReDoS anchor, quote-aware strip,
unstarted-cleanup, llama.cpp passthrough dependency). Code is unchanged
(verified comment-only via AST/ast signature, docstrings stripped).
* Harden Gemma parse/strip: span-aware XML fallback and quote-aware streaming
- Security: the XML fallback in parse_tool_calls_from_text scanned the whole
content for <function=...> markers and only skipped those inside an open XML
parameter, not those inside a collected JSON/Gemma candidate span. A balanced
but unparsable Gemma call whose argument data contained XML tool markup
(<|tool_call>call:outer{code:<function=terminal>...}<tool_call|>) therefore
fell through to the fallback and returned an executable terminal call. The
fallback now also excludes <function=> markers inside any candidate span,
including ones that failed to parse.
- strip_tool_call_markup no longer skips the generic Gemma regex after running
the quote-aware _strip_gemma_native_spans, so a closed Gemma span the helper
cannot match (malformed, e.g. <|tool_call>{"name":"x"}<tool_call|>) is still
stripped instead of leaking its opener and payload into visible text.
- _strip_gemma_native_spans stops at the first unbalanced start instead of
re-scanning every later start to EOF, keeping it linear on a run of unclosed
markers rather than quadratic.
- The GGUF and safetensors streaming strippers run _strip_gemma_native_spans
before the regex patterns, so a well-formed streamed call whose quoted
argument contains a literal close marker no longer leaks its suffix into
incremental display.
Adds regression tests for the nested-XML escape and the malformed-span strip.
* Avoid remainder copy in _strip_gemma_native_spans
Match the Gemma close marker with re pos directly on the buffer instead
of slicing tail = text[brace_end + 1:] on every span. The streaming
strippers re-scan a growing cumulative buffer per token, so the per-span
remainder copy was quadratic. Behavior is unchanged.
* Exclude unclosed Gemma/JSON starts from the XML tool-call fallback
The nested-XML guard only skipped <function=> markers inside recorded
candidate spans, but a span is recorded only when the braces balance. An
unbalanced call such as <|tool_call>call:outer{code:<function=terminal>...
recorded no span, so the fallback still promoted the inner <function=> to
an executable terminal call. Treat unclosed JSON/Gemma starts as exclusion
spans through EOF before scanning. Standalone <function=> calls with no
preceding unclosed start still parse. Regression tests added.
* Skip doomed tool-strip passes to avoid quadratic rescans
The lazy closed-pair strip patterns (<tool_call>.*?</tool_call>,
<function=...>.*?</function>) rescan to EOF from every opener when their
close token is absent, which is O(n^2) and re-runs per streamed token. Add
strip_tool_patterns, which skips a pass whose close token is not present in
the text; output is identical to the per-pattern loop (verified by fuzz),
and a degenerate run drops from ~minutes to milliseconds. Used by
strip_tool_call_markup and the GGUF/safetensors streaming strippers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use full tool-call envelopes to close nested-XML escape variants
Key the parser and stripper off the full <|tool_call>...<tool_call|> /
<tool_call>...</tool_call> envelope (start to close marker, searched after
the braces; EOF if unclosed) instead of just the braces:
- XML between the closing brace and the close marker
(call:outer{broken:{x}}<function=terminal>...<tool_call|>) is now inside
the envelope, so the fallback no longer promotes it to a tool call.
- A balanced inner call inside an unclosed outer
(call:outer{code:<|tool_call>call:terminal{...}<tool_call|>) is skipped
via the envelope nested check, not just the XML fallback.
- strip_tool_call_markup searches for the close marker after the braces, so
junk before <tool_call|> is stripped through the close and text after it is
preserved instead of truncated to EOF; a no-close run stops early (linear).
Regression tests added; standalone XML and well-formed calls unaffected.
* Fix non-final Gemma strip and missing-close recovery for PR #6611
Split the nested-skip from the XML fallback exclusion: nesting is decided by
each marker's brace region, so a balanced call after one with a missing close
marker is recovered instead of being swallowed to EOF. Only the XML fallback
keeps the search-to-close envelope, so trailing nested markup still cannot
escape as an executable call.
Use a closed-only Gemma pattern in the non-final strip list so an incomplete
block is preserved (matching the JSON and function paths); the final list keeps
the close-or-EOF Gemma pattern in its original position, so streaming display
output is byte-for-byte unchanged.
Add regression tests for both cases.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Block gap-nested tool markers and fix XML strip order for PR #6611
Decide candidate nesting by a per-marker coverage region paired with a
per-format stack (a close after the braces pops the nearest still-open marker
of that format). A closed outer call now covers up to its own close marker, so a
JSON or Gemma tool marker smuggled between the outer braces and that close is
treated as data instead of being executed. An outer that balances but has no
close of its own covers only its brace region, so a later sibling after an
omitted close marker is still recovered (adjacent calls use an exclusive end
bound so the next call is not misread as nested).
Strip every closed pair (JSON, Gemma, function) before any to-EOF sweep, so a
closed function call whose parameter text contains a bare Gemma opener is
removed as a unit and the to-EOF sweep can no longer drop the visible text after
the close.
Add regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Strip closed tool blocks before the Gemma final sweep for PR #6611
The final display strip ran the quote-aware Gemma helper before the closed
JSON/function patterns. A closed <tool_call>...</tool_call> or
<function=...>...</function> block whose argument data held a call-form Gemma
opener (e.g. a "<|tool_call>call:t{" string) was read as an incomplete Gemma
span and truncated to EOF, dropping the block's close and any visible text after
it.
Strip closed JSON/function blocks first, so such a block is removed as a unit
before the helper runs. Centralize the final strip order in a shared
strip_tool_markup_final so strip_tool_call_markup and both streaming display
wrappers (safetensors, llama_cpp) stay in sync, and apply the same closed-block
pre-pass to the non-final path.
Add regression tests for the JSON and function variants.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Recover XML/JSON siblings after a close-less tool marker for PR #6611
Two fixes so the XML fallback and marker coverage recover a later valid call
after an earlier marker omits its close, matching the candidate loop:
Reuse the candidate marker-coverage in the XML fallback instead of a separate
search-to-close-or-EOF envelope. A balanced but close-less marker now covers
only its brace region there too, so a following <function=...> sibling is
recovered rather than filtered as nested data; an unbalanced marker still covers
to EOF and a closed one still covers through its close, so nested XML stays
blocked.
Ignore a close token that falls inside another call's balanced braces when
pairing closes in _marker_coverage. Such a token is that call's quoted argument
data, so it no longer pops an earlier close-less marker and extends its coverage
over a later valid sibling.
Add regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make the closed-block strip pre-pass Gemma-span-aware
The final display strip ran the closed JSON/function regex pre-pass before
removing Gemma-native spans, so a literal <function=...> quoted inside a Gemma
argument plus any later </function> (a real call's close or even prose) was
deleted across the Gemma boundary. That mangled the Gemma close marker, the
quote-aware helper then saw an unclosed opener, and the whole visible tail
after the call was truncated.
The pre-pass now skips matches that start inside a complete Gemma span (that
text is the span's argument data) and resumes scanning at the end of the
covering span, so a real function-XML call after the Gemma call is still
stripped. The original ordering rationale is preserved: a Gemma opener inside
a JSON or function argument still cannot truncate that block, covered by
regression tests for both directions.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Trim comments in the Gemma streaming and strip pipeline to essentials
* Tighten comments in the Gemma strip and streaming disconnect paths
* Fold marker-collection comment to two lines
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""strip_tool_patterns must match the plain per-pattern loop while skipping the
|
|
quadratic no-match rescan of a closed-pair sweep whose close token is absent."""
|
|
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
|
|
if str(_BACKEND_ROOT) not in sys.path:
|
|
sys.path.insert(0, str(_BACKEND_ROOT))
|
|
|
|
from core.tool_healing import (
|
|
_TOOL_ALL_PATS,
|
|
_TOOL_CLOSED_PATS,
|
|
strip_tool_call_markup,
|
|
strip_tool_patterns,
|
|
)
|
|
|
|
|
|
def _naive(text, patterns):
|
|
for pat in patterns:
|
|
text = pat.sub("", text)
|
|
return text
|
|
|
|
|
|
_TOKENS = [
|
|
"<tool_call>",
|
|
"</tool_call>",
|
|
"<|tool_call>",
|
|
"<tool_call|>",
|
|
"<function=x>",
|
|
"<function=mcp__s__a-b>",
|
|
"</function>",
|
|
"<parameter=p>",
|
|
"</parameter>",
|
|
"call:fn{",
|
|
"}",
|
|
"{",
|
|
'<|"|>',
|
|
"A",
|
|
" ",
|
|
"\n",
|
|
"id",
|
|
"x:1",
|
|
"</tool",
|
|
"call>",
|
|
]
|
|
|
|
|
|
def test_guard_matches_plain_loop_on_fuzz():
|
|
rng = random.Random(1234)
|
|
for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS):
|
|
for _ in range(20000):
|
|
s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10)))
|
|
assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns)
|
|
|
|
|
|
def test_strip_markup_representative_cases_unchanged():
|
|
assert strip_tool_call_markup("a <tool_call>{}</tool_call> b") == "a b"
|
|
assert strip_tool_call_markup("a <function=x><parameter=p>1</parameter></function> b") == "a b"
|
|
# Non-final keeps an unclosed block; final strips it to EOF.
|
|
assert strip_tool_call_markup("a <tool_call>{partial") == "a <tool_call>{partial"
|
|
assert strip_tool_call_markup("a <tool_call>{partial", final = True) == "a"
|
|
|
|
|
|
def test_no_quadratic_blowup_on_unclosed_markers():
|
|
# Unguarded, this took minutes.
|
|
big = "<tool_call>" * 20000 + "<function=x>" * 20000
|
|
t0 = time.perf_counter()
|
|
out = strip_tool_call_markup(big, final = True)
|
|
assert time.perf_counter() - t0 < 2.0
|
|
assert out == ""
|