ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/studio/backend/tests/test_tool_xml_strip.py
Daniel Han f109e7f0e6
 Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704) 
* Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes

Extends the rescue parsers in core/tool_healing.py and
core/inference/tool_call_parser.py to recognise two extra serialisations
local models commonly emit when bypassing native function calling:

* [TOOL_CALLS]name{json_args} (Devstral-Small-2, Mistral-Small-3.x).
* name[ARGS]{json_args} (reasoning-model rehearsal).

Both extractors use a brace-balance scan that honours escapes and
quoted strings so nested JSON args stay intact.

Also pre-strips <think>...</think> and [THINK]...[/THINK] blocks before
matching so calls emitted after a reasoning preamble are recognised
regardless of position.

Streaming gates (TOOL_XML_SIGNALS, llama_cpp.py _TOOL_XML_SIGNALS) and
the SSE strip regex (routes/inference.py _TOOL_XML_RE) gain the new
sentinels so the parser is actually invoked and the raw markup never
leaks to the UI.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Strip unclosed think blocks and catch rehearsal [ARGS] mid-buffer

The pre-existing ``_THINK_TAG_RE`` only matched closed thinking
blocks (``<think>...</think>`` or ``[THINK]...[/THINK]``). During
streaming the model is still inside the open block when the parser
runs, so any tool-shaped markup the model is REHEARSING inside that
block survived the strip and could be executed as a real call.
Switch both copies of the regex (parser + healing) to accept the
trailing block being terminated by end-of-string in addition to
the explicit closer.

The ``_TOOL_XML_SIGNALS`` list on the llama_cpp streaming buffer
included ``[ARGS]`` to catch rehearsal syntax, but the gate used a
``startswith`` check against the buffer head -- rehearsal is shaped
``name[ARGS]{json}``, so the buffer never STARTS with ``[ARGS]``
and the signal had no effect. Add a substring fallback for the
bracket-style signals so the BUFFERING window can still divert the
stream into DRAINING when rehearsal markup arrives mid-buffer.

Adds three regression tests covering rehearsal inside unclosed
``<think>`` / ``[THINK]`` blocks (must yield no calls) and the
positive case after a closed think block (still parsed).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: harden bracket-tag tool-call parsing and streaming strip

Address review findings on the Mistral [TOOL_CALLS] / rehearsal [ARGS] paths:

- Accept hyphenated tool names in the bracket parsers and strip patterns.
  _MISTRAL_BRACKET_RE and _REHEARSAL_RE used \w+, which dropped or truncated
  MCP function names containing dashes (mcp__srv__list-issues). Use [\w-]+ to
  match the XML and Gemma parsers.
- Strip a partial bracket marker streamed before its opening brace. The
  trailing-unclosed patterns required the {, so a [TOOL_CALLS]web_search or
  python[ARGS] split across deltas leaked the raw marker to the UI. Match the
  bare marker to end-of-text, mirroring how the bare open tags are stripped.
  Closed pairs are unchanged so in-progress markup stays buffered until parsed.
- Strip a truncated bracket tail in the route-level display regex. _TOOL_XML_RE
  required a balanced JSON object; a tool call truncated by EOS now strips up
  to \Z, like the orphan-opening XML shapes. Complete calls still strip only
  their balanced JSON so following prose survives.

Add regression tests for hyphenated names, the streaming partial-marker strip,
and the unclosed-tail route strip.

* Studio: preserve XML parameter indentation in tool_healing

The chat template emits <parameter=k>\nVALUE\n</parameter>; the parameter-start
regex consumed the wrapping newline AND the value's first-line indentation via a
trailing \s*, then str.strip() removed the rest, corrupting code/diff arguments.
Narrow the trailing class to horizontal whitespace and trim exactly one wrapping
newline (_trim_param_value), preserving indentation. Matches SGLang's qwen3_coder
detector and the same fix on the multi-format parser. Add a regression test.

* Studio: tighten Mistral/rehearsal tool-call comments

Compress the comments in the Mistral [TOOL_CALLS] / rehearsal [ARGS] healing shim
and its callers to one or two lines, keeping the bracket-tag stripping rationale,
the thinking-block handling note, and the forge attribution intact.

Comment-only: no code or behavior change (verified with comment_tools.py check
--strip-docstrings; tests green).

* Studio: fix think-strip arg corruption and nested bracket-JSON strip

Review follow-up for the Mistral/rehearsal healing shim:

- The <think>/[THINK] strip ran unconditionally over the whole content before
  parsing, so a real tool argument that legitimately contained a <think> /
  [THINK] literal was silently corrupted. Don't delete the blocks: compute the
  reasoning-block spans and skip any tool-call candidate that STARTS inside one,
  across all parse paths (JSON, Gemma, XML, bracket, rehearsal). A rehearsed call
  inside reasoning is still ignored; a real call after </think> still parses.
- The bracket-tag display strip used a fixed one-level-nesting regex, so a call
  with two-level-nested JSON args either leaked raw markup or, in final mode, let
  the catch-all eat the trailing prose. Add a balanced-brace
  _strip_bracket_tag_calls pass (any nesting depth) used by strip_tool_call_markup
  and the route display strip.

Add regressions: <think>/[THINK] literal inside a real argument, rehearsal-inside-
think with a real call after, and two-level-nested bracket/rehearsal strip keeping
trailing prose.

* Studio: correct think-block comments to match span-skip behavior

The think-strip fix replaced the unconditional think-block strip with a
span-skip (the block is kept and any tool-call candidate starting inside it is
ignored), but two comments still described the old strip-first behavior. Update
the _THINK_TAG_RE comment and the parse_tool_calls_from_text docstring.

* Studio: parse Mistral arrays and call-ids, unify bracket parse/strip, keep it linear

- Parse the canonical Mistral array form (TOOL_CALLS followed by a JSON list of
  calls) and emit every call; parse the v11 shape that carries an opaque CALL_ID
  token between the name and ARGS (the function name is the token after
  TOOL_CALLS, never the call-id); and parse a Mistral call plus a rehearsal call
  in one message (the second was dropped yet still stripped from display).
- One shared balanced forward scan (_iter_bracket_spans) backs both the parser
  and the strip path, so they no longer diverge. It is linear: each regex is
  re-searched only once its cached match falls behind the cursor, replacing the
  per-match full-tail re-scan that was O(n^2) (O(n^3) over a stream). A length cap
  before the scan is a backstop.
- strip_tool_call_markup preserves think/reasoning blocks verbatim (the parser
  skips tool markup inside them), stripping only the visible text around them.
- _in_think uses bisect over the sorted think spans (was a linear scan per
  candidate).
- GGUF streaming strip runs the balanced bracket pre-pass before the regex
  patterns so nested-arg calls do not leak or eat trailing prose, and the
  BUFFERING ARGS detector requires the rehearsal name-ARGS shape.
- Tests: canonical array, array string-args, array strip keeps prose, Mistral
  plus rehearsal multi-call, v11 call-id name, think-rehearsal strip
  preservation, and bracket-strip linearity.

* Studio: preserve reasoning blocks in the route and streaming strip paths too

Addresses Gemini/Codex review: making strip_tool_call_markup preserve think
blocks left the route display strip and the GGUF streaming strip inconsistent,
so a rehearsed call inside a reasoning block was still deleted from the visible
text on those paths.

- Extract the think-block segmentation into one shared helper (strip_outside_think)
  and route all three strip paths through it: strip_tool_call_markup,
  _strip_tool_xml_for_display, and the GGUF _strip_tool_markup_streaming closure.
- Add a route-strip regression test that a rehearsal inside a reasoning block is
  preserved while a real call outside it is still stripped.

* Studio: fix bracket-tag strip/buffer review findings

Address the live code-review findings on the Mistral bracket-tag / rehearsal
tool-call rescue path:

- tool_healing: a literal think block inside a tool-call argument is no longer
  treated as a reasoning block. strip_outside_think now excludes think spans
  that sit inside a complete tool-call span, so the call is stripped whole
  instead of the split hiding its open/close pair and leaking the raw call.
- tool_healing: the rehearsal trailing-strip pattern requires a following brace
  or end-of-text, so prose that merely mentions name[ARGS] is not truncated as
  a phantom call. The bracket strip patterns are aligned with the parser
  regexes (whitespace, v11 [CALL_ID]/[ARGS] metadata, and the [CALL_ID]
  lookbehind).
- routes: strip a truncated canonical Mistral array ([TOOL_CALLS] [{... with no
  closing bracket) that the balanced scan cannot remove, align the display
  regex with the parser regexes, and apply the same rehearsal-prose guard.
- safetensors loop: mirror the GGUF [ARGS] rehearsal-substring check during
  BUFFERING so a rehearsal name does not stream before its [ARGS] arrives.

Adds regression tests for each; existing parser suite stays green.

* Studio: hold split rehearsal tool-name prefix in both streaming loops

A reasoning-model rehearsal call can stream the tool name and its [ARGS] arm in
separate chunks (web_search then [ARGS]{...}). The buffering detector only
recognised the rehearsal once [ARGS] was present, so the bare tool name was
emitted as visible content before the call drained and executed.

Add _is_rehearsal_prefix (mirrored in the safetensors loop and the GGUF loop):
when a no-signal buffer is a bare active-tool name -- or a partial prefix of
NAME[ARGS] -- hold it as a prefix instead of streaming it, so the next chunk's
[ARGS] flips it to a drain. A whitespace in the buffer means prose, not a split
call, so ordinary text still streams.

Adds regression tests for the split rehearsal in both loops and a guard that a
plain non-tool word still streams.

* Studio: route Anthropic tool-call cleanup through the protected display strip

The Anthropic stream, non-stream, and passthrough paths cleaned content with raw
_TOOL_XML_RE.sub instead of _strip_tool_xml_for_display, so a rehearsal call
inside <think> was deleted from the reasoning and a nested [TOOL_CALLS] call
dropped its trailing prose (the OpenAI-compatible paths already use the helper).
Route all four sites (prior-assistant cleanup, streaming content events,
non-stream aggregation, passthrough conversion) through the protected helper, and
add a source-level guard test so raw _TOOL_XML_RE.sub stays confined to the
helper itself.

* Studio: stop split rehearsal tool names leaking once streaming, uncapped, or unrestricted

The split-rehearsal guard (NAME in one chunk, [ARGS]{...} in the next) only held
the name in the initial BUFFERING state. Three gaps remained where the bare tool
name still streamed as visible content before the call drained:

- STREAMING: after prose had already streamed, both loops emitted a trailing
  active-tool-name token (and the GGUF/safetensors [ARGS] boundary was not pulled
  back over the name). Hold the trailing rehearsal token and release it on the
  next chunk, with an end-of-stream flush so a plain answer that merely ends on a
  tool-name word is never dropped.
- Buffer cap: a realistic MCP name longer than the 32-char _MAX_BUFFER_CHARS cap
  defeated the BUFFERING hold. A rehearsal prefix is self-bounding (it stops
  matching once it grows past NAME[ARGS]), so the generic cap no longer applies to
  it.
- Unrestricted mode (tools=[]): with no declared tool list, any bare identifier
  may be a NAME[ARGS] rehearsal, so the prefix check now recognises one instead of
  leaking the name and mis-parsing the call.

Regression tests cover the streaming, long-name, and unrestricted cases plus the
plain-prose paths that must not be held or corrupted.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio tools: protect think blocks in safetensors streaming, hold split rehearsal on initial flush, advertise Mistral tools

Pass-3 review follow-ups on the Mistral [TOOL_CALLS] / rehearsal [ARGS] work:

- Safetensors streaming display strip now preserves think / [THINK] reasoning
  verbatim (routes through strip_outside_think like the GGUF path). A call
  rehearsed inside a reasoning block was stripped mid-stream and then restored by
  the final strip, a non-monotonic shrink/grow that corrupted append-by-length
  stream consumers and the visible reasoning.
- The first flush out of BUFFERING (safetensors and GGUF) now applies the same
  trailing-name hold the STREAMING branch uses, so a split rehearsal (prose plus a
  trailing active tool name in one chunk, [ARGS]{...} in the next) no longer leaks
  the bare name before the call drains.
- Safetensors capability gate no longer suppresses tools for Mistral [TOOL_CALLS]
  templates, which the shared bracket-tag parser now handles end to end. Llama
  python_tag stays suppressed (still unparseable).
- Route display strip applies the open-ended / bare-marker tail arms only on the
  segment after the last reasoning block (closed-only regex before it), matching
  strip_tool_call_markup, so a bare foo[ARGS] before a reasoning block is preserved
  while complete calls are still removed in every segment.

Adds regression tests for each and updates the now-stale Mistral capability test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix tool-call think-marker and bracket-wrapper edge cases

Round-1 review follow-ups on the Mistral/rehearsal tool-call healing:

- tool_healing: a reasoning marker that opens INSIDE a tool call's
  arguments is argument data, not a reasoning block. Add
  _think_spans_outside_tool_markup (start-inside test) and use it in
  both parse_tool_calls_from_text and strip_outside_think so a literal
  marker in one call's args no longer hides a later call (parse) or
  leaks the raw markup (strip) when the greedy match runs past the
  call's closer.
- tool_healing: strip the orphan Mistral v11 [/TOOL_CALLS] closer left
  behind after the balanced scan removes the call body. Add a route arm
  for the same closer in _TOOL_XML_RE / _TOOL_XML_CLOSED_RE.
- safetensors + llama_cpp streaming strip: run the open-ended (EOS
  anchored) tail patterns only on the last segment; segments before a
  reasoning block use the closed-only patterns, matching the final
  strip and the route strip. A bare foo[ARGS] before a reasoning block
  is prose, not a truncated call.
- safetensors streaming detector: validate each [ARGS] hit before
  draining. A bare foo[ARGS] in prose (no active tool name in front)
  no longer drains the rest of the turn; a later real NAME[ARGS] call
  is still found and the prose in between is preserved.

Regression tests added for each case across the parser, strip helpers,
and both streaming loops.

* Strip incomplete-XML tool markup with literal think tags; widen render-html detector

Round-2 review follow-ups.

- tool_healing: an UNCLOSED <tool_call> / <function= call that the parser still
  executes via allow_incomplete leaked its markup when an argument contained a
  literal think marker. _tool_call_markup_spans only covered closed calls, so the
  literal was treated as a reasoning block to preserve. Extend it to the
  open-ended XML tail forms (shared as _TOOL_OPEN_XML_TAIL_PATS) so a think marker
  inside an unclosed call is argument data and the call's markup is stripped. A
  complete call's opener stays bounded to its closed span, and a real reasoning
  block with no tool call is still preserved.
- safetensors render-html provisional card: _detect_render_html_tool_start was
  XML-only, so a Mistral [TOOL_CALLS]render_html or rehearsal render_html[ARGS]
  call executed but skipped the early card. Detect the earliest tool-call marker
  across every serialization the loop executes and fire when it is render_html.

Regression tests added for both.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio tools: gate [ARGS] on active tools and skip think-block render_html rehearsal

Round 3 review fixes for the Mistral / rehearsal tool-call parsing path. Both are
asymmetric-fix bugs where one code path applied a guard the analogous paths did not.

- [ARGS] active-tool gating: the streaming state already validates a rehearsal
  NAME[ARGS] against the active tool list before draining, but the BUFFERING
  detection and the end-of-stream safety-net checks (safetensors and GGUF) treated
  any word[ARGS] substring as a tool boundary. An answer containing a literal
  foo[ARGS]{...} in prose, where foo is not an enabled tool, was drained, parsed into
  a disabled foo no-op, and forced an extra generation turn. Gate those checks on the
  active tool name too (unrestricted mode still accepts any name), so inactive-name
  prose is neither drained nor parsed. Adds a shared _has_genuine_tool_signal helper
  (safetensors) and _gguf_rehearsal_signal_pos / _gguf_has_genuine_tool_signal (GGUF).

- render_html provisional card vs think blocks: the parser skips tool candidates that
  start inside a <think>/[THINK] reasoning block, but the provisional render_html
  detector scanned raw content. A render_html rehearsed inside <think> followed by a
  real non-render_html call emitted a provisional render_html tool_start (reusing the
  later call's id) that the loop never executed. Drop candidates that start inside a
  think span and use the first marker of each shape outside the blocks. Also resolve
  the [TOOL_CALLS] [{...}] array shape through the parser so a nested "name" argument
  key no longer fires a false provisional card ahead of the real top-level tool name.

Adds regression tests for both loops: inactive-name foo[ARGS]{...} is not drained into
a disabled no-op or a retry turn, a think-block render_html rehearsal emits no
provisional card, and the array top-level name is read correctly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Gate ambiguous bare-rehearsal parse and strip on the active tool list

A bare NAME[ARGS]{json} is a genuine rehearsal call only when NAME is an
active tool; otherwise it is prose. The earlier round gated only detection
(so an inactive foo[ARGS] no longer drained the buffer or forced a retry
turn), but the parse and strip stayed unrestricted, which produced two
regressions:

1. An inactive foo[ARGS]{...} placed immediately before a real
   web_search[ARGS]{...} in the same content span made the real call fail
   to execute (parse consumed the phantom foo call).
2. An inactive foo[ARGS]{...} in a prose answer had its markup stripped
   from the visible text, corrupting the sentence to " is just syntax."

Thread enabled_tool_names through the shared parser/strip so parse and
strip apply the SAME active-tool gate as detection:

- core/tool_healing.py: _iter_bracket_spans skips an inactive rehearsal
  span; parse_tool_calls_from_text, _strip_bracket_tag_calls,
  _strip_markup_segment and strip_tool_call_markup accept and thread the
  gate; apply_tool_strip_patterns keeps an inactive rehearsal match.
- core/inference/tool_call_parser.py: wrappers forward the gate.
- core/inference/safetensors_agentic.py and core/inference/llama_cpp.py:
  compute the gate from the active tool list (None when unrestricted, to
  keep the legacy strip-all behavior) and thread it into every parse and
  streaming/final strip site.
- routes/inference.py: _strip_tool_xml_for_display accepts the gate and
  keeps an inactive rehearsal via a capture group on its rehearsal arm, so
  the display cleanup does not re-strip the already-correct loop output.
  The [TOOL_CALLS] control-token arms still strip unconditionally. Wire
  the current turn's active tool names into the GGUF and safetensors
  content-display sites.

Tests: parse and strip gate coverage in test_tool_call_parser_strict.py,
test_tool_xml_strip.py and test_safetensors_tool_loop.py; end-to-end GGUF
coverage for the real-call-after-inactive-rehearsal case and a
strengthened assertion that the inactive rehearsal prose survives intact.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: render the reasoning block for safetensors and MLX like GGUF

enable_thinking chat templates (Qwen3/Qwen3.5/GLM) prefill an unclosed <think>
into the generation prompt, so the model emits only the closing </think> then
the answer. The safetensors/MLX chat stream emitted that as plain content, so
the reasoning showed inline with no collapsible thinking block, while GGUF
(which surfaces reasoning via reasoning_content) rendered one. This brings
safetensors and MLX to parity.

- _ResponsesReasoningExtractor gains a reasoning_prefilled mode that starts
  inside the reasoning block and splits on the first </think>; default False
  keeps GGUF and every existing caller byte-identical. It suppresses a stray
  re-emitted <think> and holds partial markers back across chunk boundaries.
- _sf_reasoning_prefill_mode gates the mode on reasoning being enabled for the
  request, an enable_thinking or enable_thinking_effort style, and the template
  actually using the standard <think>/</think> markers. Models with a bespoke
  reasoning channel (e.g. gemma's <|think|>/<|channel>) are excluded so their
  answer is never swallowed; gpt-oss (Harmony) and thinking-off requests are
  excluded too.
- sf_tool_stream and stream_chunks (the latter also serves MLX) feed text
  through the extractor, emitting reasoning_content then content deltas, with a
  per-turn reset in the tool loop and a flush before each tool_start; only the
  visible delta reaches the monitor reply. The two non-streaming drains split
  reasoning_content the same way.
- Tests: extractor prefilled mode (streaming and edge cases), the gate matrix
  including the gemma-style exclusion, and a route-replay of the tool-loop
  reasoning stream.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: skip tool calls rehearsed in prefilled reasoning

Reasoning models (Qwen3.5 enable_thinking) open <think> in the prompt, so the
generated text starts inside the thought and emits only a closing </think> with
no opener. _think_spans_outside_tool_markup only found spans with an explicit
opener, so a NAME[ARGS]{...} or [TOOL_CALLS] call rehearsed in that leading
thought was parsed and executed as a real call.

Add a leading think span (offset 0 through the first close marker) when the
content opens with a bare close, so the rehearsed call is skipped and the
reasoning is preserved by strip_outside_think. Guarded by the existing call-span
check: a literal </think> inside a real call's arguments does not trigger the
span, so a genuine leading call still fires. Tests for both cases.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: do not start prefilled reasoning mode when reasoning_effort is none

enable_thinking_effort models (e.g. GLM-5.2) express thinking-off via
reasoning_effort="none" rather than enable_thinking=False, but
_sf_reasoning_prefill_mode only looked at enable_thinking, so such a request
started the extractor in prefilled mode. With thinking off the model never emits
</think>, so the whole answer was captured as reasoning_content and the visible
content/stream came back empty. Thread reasoning_effort through and return False
when it is "none". Tests for none vs a real effort level.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* studio: only treat a leading bare </think> as prefilled reasoning when a real call follows

The prefilled-reasoning virtual span fired on any unmatched leading close marker,
so a non-prefilled turn that emits a real call before a stray </think> (for
example "Now web_search[ARGS]{...}</think> answer") had the call swallowed by the
span and dropped. Require that a real tool call also appear after the close (the
actual turn that follows the thought) before adding the span, so a stray close in
a normal answer no longer suppresses a genuine leading call. The rehearse-then-
call case still skips the rehearsal. Test for the stray-close case.

* Studio: trim redundant comments (comment-only, AST-verified)

* studio: keep tool_healing importable on Python 3.9

_balanced_json_span was annotated -> int | None. With no
from __future__ import annotations, that PEP 604 union is evaluated at
import time, so on Python 3.9 (which the package still supports,
requires-python >=3.9, and where external inference servers import this
module standalone) the def raises TypeError and the whole module fails
to import before any parsing runs.

Add from __future__ import annotations so annotations stay lazy strings,
matching the prevailing convention across studio/backend. No behavior
change: the module has no runtime annotation introspection.

* Studio: gate the Anthropic tool-stream display strip on declared tools

The Anthropic streaming and non-streaming tool paths called
_strip_tool_xml_for_display without enabled_tool_names, so with the default
strip-all behavior a final answer that literally contains an inactive-name
NAME[ARGS]{json} (prose, not a call) lost those bytes in the delivered text.
The GGUF and safetensors paths already pass _display_tool_name_gate(tools);
these two sites were missed when that gate was threaded through.

Compute the gate from the declared tools and pass it at both sites (threading
openai_tools into _anthropic_tool_non_streaming and its caller), so an
inactive-name rehearsal survives while an active-name one is still stripped.
Add a regression test.

* Studio: hold a split unrestricted rehearsal prefix at the bracket

In unrestricted tool mode (tools=[]) the rehearsal-prefix regex required
[A after the bracket, so a chunk boundary landing right after NAME[ (e.g.
web_search[ then ARGS]{...}) failed the prefix check and streamed the
partial tool markup web_search[ to the client before the call drained.
Restricted mode already holds this via a startswith check. Make the bracket
and each ARGS letter individually optional so NAME[ is held too, matching
the documented intent. Add a regression test.

* Studio: gate rehearsal detection and history strip on the original tool set

Two display/loop gate fixes so a spent one-shot tool is handled consistently:

- Rehearsal DETECTION (safetensors and GGUF loops) now uses the ORIGINAL tool
  list, matching the strip gate, instead of the post-removal active_tools. After a
  one-shot tool (render_html) runs it is dropped from active_tools; a repeat
  render_html[ARGS]{...} while another tool is still active was stripped from
  display yet never detected, so it was not routed to the render_html_repeat no-op
  and the turn ended as a blank continuation. Detection now fires for it.

- The GGUF assistant-history sanitiser forwards the enabled-tool-name gate (like
  the live-response strip), so a prior turn documenting an inactive foo[ARGS]{...}
  shape is preserved in the replayed prompt context instead of being deleted.

Add regression tests for both loops and the history strip.

* Studio: thread the tool-name gate through the remaining rehearsal/history sites

Follow-up to the rehearsal-detection and history-strip gate fixes, covering the
sibling sites that were missed:

- GGUF loop: the rehearsal-prefix and trailing-name hold checks now use the
  original tool list (_detect_tools) like the detection path, so a spent one-shot's
  split repeat (bare render_html then [ARGS]{...}) is held instead of flushed as
  visible text.
- The safetensors and Anthropic assistant-history sanitisers and the Anthropic
  non-streaming passthrough now forward the enabled-tool-name gate to
  _strip_tool_xml_for_display, matching the GGUF history sanitiser and the live
  strips, so a prior turn documenting an inactive foo[ARGS]{...} example is
  preserved in the replayed prompt / final text instead of deleted.

Add regression tests.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tile bracket-call spans per array item and include the v11 closer

Two with_spans fixes for the Mistral bracket parser, both hit through the
client-tool passthrough healers:
- A multi-call [TOOL_CALLS] array carried its whole markup span on the first
  call and zero-width spans after, so a consumer that filters promotions by
  the declared tool set either re-emitted the full raw array as text next to
  the promoted call or silently dropped a filtered call's bytes. The region is
  now tiled across the call-producing items (each call's span covers its own
  JSON object plus the separator bytes before it; the last span runs to the
  region end), so promoted markup strips exactly once and a skipped call's
  bytes stay visible.
- The v11 wrapper closer [/TOOL_CALLS] sat outside the reported span and
  leaked as stray text after promotion; the region now extends over an
  immediately-following closer.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Address review: decouple healer signals from the loop signal set

The passthrough healer buffered on every TOOL_XML_SIGNALS entry, so the bare
[ARGS] rehearsal marker this branch adds for the loops (where it is gated on
active tool names) put legitimate prose like 'Use foo[ARGS] in templates'
into the holding state and stalled the stream until finalization. The healer
can never promote a bare rehearsal call, so it now buffers only on formats
its parser promotes: <tool_call>, <|tool_call>, <function=, [TOOL_CALLS].

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Condense comments in the Mistral tool-call rescue to contract essentials

* verify_import_hoist: exempt __future__ imports and same-diff relocations

Two false positives fired on this PR's refactor. A from __future__ import
is a compiler directive whose name never appears as a runtime load, so
HOISTED-IMPORT-UNUSED can never see it used, yet the file requires it for
PEP 604 annotations on Python 3.9. TARGET-CHANGED flagged the deliberate
move of the strip-pattern constants into core.inference.tool_call_parser
as a silent re-point even though the old module-level target was removed
and the new one added in the same diff. Both get narrow exemptions; a
re-point to a pre-existing target is still caught, and the self-test
negative controls all pass unchanged.

* Drain the whole Mistral [TOOL_CALLS] array in streaming passthrough healing

StreamToolCallHealer._drain promoted only the first parsed call per pass and
dropped the rest of the buffer past that one span. For a well-formed Mistral
parallel-tool-call array streamed through client-tool passthrough
([TOOL_CALLS][{...},{...}]), the per-item spans are contiguous, so after the
first call was promoted the residue began with ,{...}] (no leading signal) and
was flushed as raw text: every call after the first was lost.

_drain now walks the contiguous run of parsed calls (adjacent tiled spans =
one array), promoting each declared call and relaying undeclared ones as data,
and stops at the first gap (prose) or incomplete trailing block so separate
blocks still stream incrementally in document order. This mirrors the
non-streaming heal_openai_message / finalize promote-or-flush loop and the
server-side safetensors loop, which already handled multi-call arrays.

Added regression tests: 2-call array in one feed and char-by-char, an
undeclared middle call kept as text, and an array followed by trailing prose.

* Drain comma-less Mistral tool-call arrays and normalize null arguments

The array branch fed the whole body to a single json.loads, which rejects the
comma-less multi-call form the repo's own Mistral/Ollama templates render (the
range loop in ollama_template_mappers.py emits the objects with no separator) and
so dropped every call. Decode elements individually with the existing
comma-tolerant raw_decode helper, now _decode_array_items, which also returns the
objects, so all calls are recovered while the span tiling is unchanged.

Also normalize a non-object array argument such as arguments null to an empty
object, matching the wrapped tool_call path, instead of serializing None to the
string "null" that auto-heal would turn into a bogus query of "null".

* Gate safetensors reasoning prefill on the rendered generation prompt

reasoning_always_on fires on any paired <think></think> in the template,
including markup that only renders PAST assistant history (Kimi-K2-Thinking)
while the generation prompt opens no <think>. Starting the reasoning extractor
in prefilled mode there captured a normal answer entirely as reasoning_content
and returned blank visible content. Prefill only when rendering the generation
prompt actually leaves <think> open (DeepSeek-R1 / QwQ / Qwen3-Thinking);
history-only templates start the extractor in normal mode and parse the model's
own <think>...</think>. Adds a Kimi-shape regression test.

* Keep bare scalar Mistral array arguments raw instead of double-encoding

A scalar string argument in the canonical Mistral [TOOL_CALLS] array
(for example [TOOL_CALLS][{"name":"web_search","arguments":"weather"}])
was run through json.dumps, turning weather into the JSON string
"weather". The downstream argument healer then wrapped that quoted
form, so a single-string tool like web_search searched for the literal
"weather" with quotes. The <tool_call> path already keeps a scalar
argument raw; mirror it here so only a dict is serialized. Add a
regression test asserting both paths yield the same healed arguments.

* Tighten tool-call rescue and reasoning-prefill comments

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* [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>
2026-07-06 18:52:13 -07:00

916 lines
40 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 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"
# The lazy ``(.*?)\n\)`` could grab a shorter expression if an arm is ever wrapped;
# pin the DeepSeek + bare-Kimi arms so a silent truncation fails loudly here.
assert "_DS_OPEN_SRC" in _m.group(1) and "tool_call_begin" in _m.group(
1
), "extracted _TOOL_XML_RE is missing expected arms (extraction truncated?)"
# The regex reuses the parser's shared DeepSeek opener alternation; provide it so the extracted
# ``_re.compile`` expression resolves the same source.
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
from core.inference.tool_call_parser import (
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
)
from typing import Optional as _Optional
_ns = {
"_re": _re,
"_DS_OPEN_SRC": _DS_OPEN_SRC,
"Optional": _Optional,
"_strip_mistral_closed_calls": _strip_mistral_closed_calls,
"_strip_gemma_wrapperless_calls": _strip_gemma_wrapperless_calls,
"_strip_glm_calls": _strip_glm_calls,
"_strip_function_xml_calls": _strip_function_xml_calls,
}
exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
_TOOL_XML_RE = _ns["_TOOL_XML_RE"]
# The display helper uses the closed-only variant before the last think block; keep it in scope.
_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL)
assert _mc, "could not extract _TOOL_XML_CLOSED_RE source"
exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns)
_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"]
# Signatures may span multiple lines and now carry the enabled_tool_names gate; match
# the whole (possibly multi-line) signature up to ``-> str:`` then the indented body.
_xml_helper = _re.search(
r"def _strip_tool_xml\((?:.|\n)*?\) -> str:\n(?: .+\n)+",
_src,
)
assert _xml_helper, "could not extract _strip_tool_xml source"
assert "_strip_mistral_closed_calls" in _xml_helper.group(
0
), "extracted _strip_tool_xml no longer runs the Mistral balanced strip"
exec(_xml_helper.group(0), _ns)
_strip_tool_xml = _ns["_strip_tool_xml"]
# Extract the gate helper and display strip up to the next top-level ``logger =``.
_helper = _re.search(
r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)",
_src,
_re.DOTALL,
)
assert _helper, "could not extract display strip helper source"
# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before
# ``logger =``); confirm the shared _strip_tool_xml delegate is present.
assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates"
exec(_helper.group(0), _ns)
_strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"]
_display_tool_name_gate = _ns["_display_tool_name_gate"]
_gate_src = _re.search(
r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+",
_src,
)
assert _gate_src, "could not extract _gemma_strip_gate source"
exec(_gate_src.group(0), _ns)
_gemma_strip_gate = _ns["_gemma_strip_gate"]
# ── Well-formed pairs ─────────────────────────────────────────────
def test_route_display_strip_respects_disabled_auto_heal_contract():
text = 'literal <tool_call>{"name":"web_search"}</tool_call> survives'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
assert "<tool_call>" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
def test_route_display_strip_preserves_rehearsal_inside_think():
# A rehearsed bracket call inside think is reasoning: the block is preserved while a real
# call outside it still strips.
text = '<think>plan: search[ARGS]{"q":"x"}</think> answer [TOOL_CALLS]web_search{"q":"y"} tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert '<think>plan: search[ARGS]{"q":"x"}</think>' in out
assert "[TOOL_CALLS]web_search" not in out
assert "answer" in out and "tail" in out
def test_route_display_strip_keeps_bare_args_before_think_block():
# A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on
# the last segment (earlier segments use the closed-only regex).
text = "Please pass foo[ARGS] <think>pause</think> to the template."
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text
def test_route_display_strip_removes_complete_call_before_think_block():
# A complete bracket call before a think block still strips (balanced scan runs on every segment).
text = 'before search[ARGS]{"q":"x"} <think>pause</think> after'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "search[ARGS]" not in out
assert "<think>pause</think>" in out
assert "before" in out and "after" in out
def test_route_display_strip_removes_closed_xml_before_think_block():
# A closed <tool_call> before a think block is removed in the non-last segment.
text = 'pre <tool_call>{"name":"x"}</tool_call> <think>p</think> tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "<tool_call>" not in out
assert "<think>p</think>" in out
assert "pre" in out and "tail" in out
def test_all_route_cleanup_sites_use_protected_display_helper():
# Every route cleanup site must use _strip_tool_xml_for_display (think-preserving,
# balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only
# legitimate raw sub lives inside the helper itself.
raw_sub_lines = [
(i, line)
for i, line in enumerate(_src.splitlines(), 1)
if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#")
]
assert len(raw_sub_lines) == 1, (
"raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; "
f"found extra call sites: {raw_sub_lines!r}"
)
def test_route_display_strip_removes_mistral_tool_calls_with_nested_json():
# _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral
# balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON).
text = 'ok [TOOL_CALLS]web_search{"filters":{"date":"2024"},"query":"cats"} tail'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "web_search" not in out, out
assert out == "ok tail"
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
def test_strips_function_attribute_form():
# Attribute form ``<function name="...">`` (MiniCPM-5 / MiniMax-M2) must strip from the route too
# (it previously leaked into the UI); a dotted/hyphenated name also strips.
text = (
'Sure.\n<function name="get_weather">\n'
"<parameter=city>\nSydney\n</parameter>\n</function>\nDone."
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "<function name=" not in cleaned
assert "</function>" not in cleaned
assert "Sure." in cleaned and "Done." in cleaned
dotted = 'A <function name="srv.list-issues">x</function> B'
assert _TOOL_XML_RE.sub("", dotted) == "A B"
# Auto-Heal-disabled display contract still preserves literal markup.
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
assert "<function name=" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
# ── 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).
def test_strips_gemma_native_orphan_closing_tag():
cleaned = _TOOL_XML_RE.sub("", "Tool call drained.<tool_call|>Visible tail.")
assert "<tool_call|>" not in cleaned
assert "Tool call drained." in cleaned
assert "Visible tail." in cleaned
# ── 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_strips_complete_bracket_tag_keeps_trailing_prose():
# A complete Mistral call strips only its balanced JSON, leaving following prose intact.
cleaned = _TOOL_XML_RE.sub("", '[TOOL_CALLS]web_search{"q":"x"} and then prose')
assert "[TOOL_CALLS]" not in cleaned
assert "and then prose" in cleaned
def test_strips_unclosed_bracket_tail():
# Close brace lost to EOS: the truncated tail strips to the end instead of leaking.
cleaned = _TOOL_XML_RE.sub("", 'here [TOOL_CALLS]web_search{"query":"weather"')
assert "[TOOL_CALLS]" not in cleaned
assert cleaned.strip() == "here"
def test_strips_unclosed_rehearsal_tail():
cleaned = _TOOL_XML_RE.sub("", 'text python[ARGS]{"code":"print(1)"')
assert "[ARGS]" not in cleaned
assert cleaned.strip() == "text"
def test_strips_hyphenated_mcp_bracket_name():
cleaned = _TOOL_XML_RE.sub("", 'x [TOOL_CALLS]mcp__srv__list-issues{"q":"x"}')
assert "list-issues" not in cleaned
assert cleaned.strip() == "x"
def test_preserves_mid_string_parameter_in_code_sample():
# Tail-anchor on `</parameter>` 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
# ── Two-level-nested bracket JSON (balanced-scan strip) ──────────
def test_route_strip_two_level_nested_bracket_keeps_trailing_prose():
# Two-level-nested args must be removed whole so the trailing prose survives.
text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after'
cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert cleaned == "before after"
assert "[TOOL_CALLS]" not in cleaned
def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose():
text = 'note python[ARGS]{"a":{"b":{"c":1}}} done'
cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert cleaned == "note done"
assert "[ARGS]" not in cleaned
def test_route_strip_removes_call_with_literal_think_in_argument():
# A literal <think> inside a call argument strips with the call, not as reasoning.
text = (
'<tool_call>{"name":"write","arguments":'
'{"text":"compare <think> and </think> tags"}}</tool_call>'
)
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "<tool_call>" not in out and '"name"' not in out
def test_route_strip_removes_truncated_mistral_array():
# A canonical array truncated by EOS is stripped by the route fallback like other orphans.
text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ]
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "{" not in out
assert "before" in out
def test_route_strip_keeps_prose_mentioning_args_marker():
# ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line.
text = "Please pass foo[ARGS] to the template and continue reading."
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert out == text
def test_route_strip_handles_mistral_v11_call_id_args_shape():
# v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole.
text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out
assert "before" in out and "after" in out
# ── Mistral [/TOOL_CALLS] closer + literal <think> inside a call ───────────────
from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup
def test_core_strip_removes_orphan_tool_calls_closer_array_form():
# The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content.
text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]'
assert _strip_tool_call_markup(text, final = True) == ""
def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail():
text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail'
assert _strip_tool_call_markup(text, final = True) == "tail"
def test_core_strip_removes_call_with_literal_think_in_argument():
# An unclosed literal <think> inside call arguments strips with the call (argument data).
text = 'before <tool_call>{"name":"write","arguments":{"text":"literal <think> marker"}}</tool_call> after'
assert _strip_tool_call_markup(text, final = True) == "before after"
def test_route_display_strip_removes_orphan_tool_calls_closer_array_form():
text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert out.strip() == ""
def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail():
text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail'
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "[/TOOL_CALLS]" not in out
assert out.strip() == "tail"
def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped():
# An incomplete <tool_call> holding a literal <think> strips to EOS, not as a reasoning
# block (the unclosed tail _tool_call_markup_spans previously missed).
from core.tool_healing import parse_tool_calls_from_text as _parse
from core.tool_healing import strip_tool_call_markup as _strip
text = 'before <tool_call>{"name":"write","arguments":{"text":"literal <think> marker"}} after'
assert [c["function"]["name"] for c in _parse(text)] == ["write"]
assert _strip(text, final = True) == "before"
# A real reasoning block with no tool call is still preserved verbatim.
assert (
_strip("answer <think>real</think> done", final = True) == "answer <think>real</think> done"
)
# A complete call followed by a real reasoning block: call stripped, block kept.
mixed = '<tool_call>{"name":"a","arguments":{}}</tool_call> mid <think>r</think> end'
assert _strip(mixed, final = True) == "mid <think>r</think> end"
# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ──
def test_display_tool_name_gate_returns_active_names_or_none():
# Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior).
assert _display_tool_name_gate([]) is None
assert _display_tool_name_gate(None) is None
# OpenAI-shaped tool dicts -> set of function names, malformed entries dropped.
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "run_python"}},
{"type": "function"}, # no name
{"nope": 1}, # no function
]
assert _display_tool_name_gate(tools) == {"web_search", "run_python"}
def test_route_display_strip_keeps_inactive_rehearsal_when_gated():
# P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact.
gate = {"web_search"}
text = 'foo[ARGS]{"x":1} is just syntax.'
assert (
_strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate)
== text
)
# A bare marker with no JSON body is likewise prose when inactive.
assert (
_strip_tool_xml_for_display(
"use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate
)
== "use foo[ARGS] here"
)
def test_route_display_strip_removes_active_rehearsal_when_gated():
# Mirror case: an active tool name is a real rehearsal and still strips.
gate = {"web_search"}
out = _strip_tool_xml_for_display(
'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate
)
assert "web_search[ARGS]" not in out
assert out.strip() == "done"
def test_route_display_strip_ungated_strips_all_rehearsal_unchanged():
# Backwards-compat: with no gate (None) the bare rehearsal strips as before.
text = 'foo[ARGS]{"x":1} is just syntax.'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax."
assert (
_strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = None
).strip()
== "is just syntax."
)
def test_route_display_strip_control_token_stripped_regardless_of_gate():
# [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate.
gate = {"web_search"}
out = _strip_tool_xml_for_display(
'[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate
)
assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out
assert out.strip() == "keep"
def test_core_strip_gates_bare_rehearsal_on_enabled_tools():
# P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose
# and preserved, active names strip, ``None`` keeps legacy strip-all.
from core.tool_healing import strip_tool_call_markup as _strip
text = 'foo[ARGS]{"x":1} is just syntax.'
assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text
assert (
_strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"})
== "done"
)
assert _strip(text, final = True).strip() == "is just syntax."
assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax."
def test_route_display_strip_gate_preserves_inactive_history_rehearsal():
# The GGUF history sanitiser passes the gate, so a documented inactive shape survives in
# the replayed prompt context.
gate = _display_tool_name_gate([{"function": {"name": "web_search"}}])
text = 'To call it write foo[ARGS]{"x":1} in your reply.'
assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = gate
)
# An ACTIVE name is still stripped as a real rehearsed call.
assert "web_search[ARGS]" not in _strip_tool_xml_for_display(
'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate
)
# No gate (legacy) strips every NAME[ARGS]{...}.
assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate():
# Wiring guard: the GGUF history strip must forward the display gate like the live strip.
block = _re.search(
r"Strip stale tool-call XML from conversation history.*?\.strip\(\)",
_src,
_re.DOTALL,
)
assert block, "could not locate GGUF history sanitizer block"
assert "enabled_tool_names" in block.group(
0
), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display"
def test_route_history_and_passthrough_forward_the_display_gate():
# The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough
# must forward the gate so inactive examples survive in replayed prompt / final text.
blocks = {
"safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)",
"anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)",
"anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)",
}
for label, pat in blocks.items():
m = _re.search(pat, _src, _re.DOTALL)
assert m, f"could not locate {label} strip block"
assert "enabled_tool_names" in m.group(
0
), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display"
# ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ──
def test_strips_deepseek_space_opener_variant():
# The space-separated opener is parsed by the parser, so the display strip
# must remove it too (the shared opener alternation is reused here).
text = (
"pre <｜tool calls begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "tool" not in cleaned.replace("post", "").replace("pre", "")
assert cleaned == "pre post"
def test_strips_deepseek_escaped_underscore_opener_variant():
text = (
"pre <｜tool\\_calls\\_begin｜><｜tool▁call▁begin｜>get_y<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "pre post"
def test_strips_bare_kimi_call_without_section_wrapper():
# Kimi can emit a bare <|tool_call_begin|>...<|tool_call_end|> with no
# section wrapper; the parser accepts it, so the strip must cover it.
text = (
"pre <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
'{"a":1}<|tool_call_end|> post'
)
cleaned = _TOOL_XML_RE.sub("", text)
assert "tool_call_begin" not in cleaned
assert cleaned == "pre post"
@pytest.mark.parametrize(
"text",
[
# Prose that merely names a Kimi/DeepSeek marker (no real call follows) must
# survive: the call-shaped lookahead fires only on a real call or a bare EOF
# fragment, so an answer discussing the protocol is never truncated.
"See <|tool_call_begin|> in the docs. More prose after it.",
"The <|tool_calls_section_begin|> marker opens a batch. Read on.",
"DeepSeek uses <｜tool▁calls▁begin｜> to start a call block, then continues.",
],
)
def test_deepseek_kimi_false_alarm_prose_is_kept(text):
# Regression for the route arm truncating a prose answer that references a marker
# without a following call (parser _TOOL_ALL_PATS already had this lookahead).
assert _TOOL_XML_RE.sub("", text) == text
def test_deepseek_kimi_real_calls_still_strip_after_false_alarm_fix():
# The lookahead must not weaken real-call stripping: closed, truncated, and bare
# EOF-fragment forms all still get removed.
closed = (
"answer <|tool_call_begin|>functions.get_w:0<|tool_call_argument_begin|>"
'{"a":1}<|tool_call_end|> tail'
)
assert _TOOL_XML_RE.sub("", closed) == "answer tail"
eof_fragment = "prefix <|tool_call_begin|>"
assert _TOOL_XML_RE.sub("", eof_fragment) == "prefix "
deepseek = (
"reply <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_x<｜tool▁sep｜>"
'{"a":1}<｜tool▁call▁end｜><｜tool▁calls▁end｜>'
)
assert _TOOL_XML_RE.sub("", deepseek) == "reply "
# ── Llama-3 <|python_tag|> arm bounds on REAL sentinels only ──────
# Llama-3 <|python_tag|> arm bounds on REAL sentinels only
def test_python_tag_strip_consumes_literal_sentinel_in_arg():
# A <|python_tag|> tool call whose JSON argument carries a literal <|...|>
# token (here <|cite|>) must be stripped whole. The old `<(?!\|)` arm stopped
# at any `<|`, leaking the call tail (e.g. `<|cite|> here"}}`) into display.
text = '<|python_tag|>{"name": "send", "parameters": {"text": "use <|cite|> here"}}'
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "", f"python_tag call leaked at literal sentinel: {cleaned!r}"
@pytest.mark.parametrize(
"sentinel",
[
"<|eot_id|>",
"<|eom_id|>",
"<|start_header_id|>",
"<|end_header_id|>",
],
)
def test_python_tag_strip_stops_at_real_sentinel(sentinel):
# A genuine Llama control sentinel still bounds the strip so following
# assistant text is preserved (the arm must not swallow past it).
text = f'<|python_tag|>{{"name": "x", "parameters": {{}}}}{sentinel}visible answer'
cleaned = _TOOL_XML_RE.sub("", text)
assert (
cleaned == f"{sentinel}visible answer"
), f"strip did not stop at real sentinel {sentinel!r}: {cleaned!r}"
def test_python_tag_strip_restarts_on_second_python_tag():
# A second <|python_tag|> opens a new tool-call region, so the whole pair is
# stripped (the arm bounds the first, then the next match consumes the rest).
text = '<|python_tag|>{"name": "a"}<|python_tag|>{"name": "b"}'
cleaned = _TOOL_XML_RE.sub("", text)
assert cleaned == "", f"second python_tag region leaked: {cleaned!r}"
def test_glm_call_with_literal_close_tag_in_arg_value_is_stripped_whole():
# GLM 4.x emits <tool_call>NAME<arg_key>k</arg_key><arg_value>v</arg_value> ...</tool_call>.
text = (
"<tool_call>web_search\n<arg_key>query</arg_key>\n"
"<arg_value>find </tool_call> here</arg_value>\n</tool_call> done"
)
out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
assert "</arg_value>" not in out
assert "<arg_key>" not in out
assert out.strip() == "done"
def test_glm_normal_and_qwen_calls_still_stripped_by_route():
# Regression: a normal GLM call (no literal close tag) and a Qwen
# <tool_call>{json}</tool_call> are still stripped; trailing prose is kept.
glm = "<tool_call>get_time\n<arg_key>tz</arg_key>\n<arg_value>UTC</arg_value>\n</tool_call> ok"
assert _strip_tool_xml_for_display(glm, auto_heal_tool_calls = True).strip() == "ok"
qwen = '<tool_call>{"name":"web_search","arguments":{"q":"x"}}</tool_call> after'
assert _strip_tool_xml_for_display(qwen, auto_heal_tool_calls = True).strip() == "after"
def test_route_strip_removes_param_alias_close_tag():
# The parser accepts the <param name="...">...</param> attribute-form alias of
# <parameter=...>; the route tail cleanup must strip an orphan </param> close too.
assert _strip_tool_xml_for_display("answer </param>", auto_heal_tool_calls = True) == "answer "
assert (
_strip_tool_xml_for_display("answer </parameter>", auto_heal_tool_calls = True) == "answer "
)
def test_route_strip_uses_guarded_function_scan_for_literal_nested_markup():
# A literal <function=...></function> in a value must not truncate the strip: the route runs the
# parser's guarded function-XML scan before the regex, matching the core strip.
text = "<function=python><parameter=code><function=evil></function></parameter></function> tail"
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "tail"
def test_route_strip_gates_wrapperless_gemma_by_enabled_tools():
# The route strip must gate the markerless Gemma call:NAME{...} form on the enabled tool names,
# like the parser/loop, so a disabled/example name in prose is preserved in ...
prose = "To document syntax you write call:foo{query:example}. That shows the format."
assert "call:foo{query:example}" in _strip_tool_xml(prose, {"web_search"})
# An enabled name is still a real call and stripped.
assert "call:web_search" not in _strip_tool_xml(
"Answer. call:web_search{query:x}", {"web_search"}
)
# No gate (legacy) strips every closed call.
assert "call:foo" not in _strip_tool_xml(prose)
def test_gemma_strip_gate_empty_tools_preserves_prose():
# With NO tools enabled the gate must return an EMPTY set (strip nothing), not None: None falls
# back to strip-all and deletes an answer that documents the call:NAME{...} syntax.
assert _gemma_strip_gate([]) == set()
assert _gemma_strip_gate(None) == set()
assert _gemma_strip_gate([{"function": {"name": "web_search"}}]) == {"web_search"}
prose = "To document syntax you write call:foo{query:example}. That shows the format."
assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate([]))
assert "call:foo{query:example}" in _strip_tool_xml(prose, _gemma_strip_gate(None))
# An enabled tool's real call is still stripped.
assert "call:web_search" not in _strip_tool_xml(
"Answer. call:web_search{query:x}",
_gemma_strip_gate([{"function": {"name": "web_search"}}]),
)
def test_strip_keeps_prose_after_closed_function_call_with_literal_close():
# The call ends at its first non-data close: prose after it survives the
# strip even when it mentions a literal </function>.
from core.inference.tool_call_parser import strip_tool_markup
text = (
"<function=web_search><parameter=query>cats</parameter></function>"
" Done. The tag </function> closes a call."
)
assert strip_tool_markup(text, final = True) == "Done. The tag </function> closes a call."
def test_final_strip_keeps_prose_mentioning_bare_markers():
# A false-alarm marker in a normal answer must not lose everything after
# it; only text that looks like that family's call start drops.
from core.inference.tool_call_parser import strip_tool_markup
for text in (
"See [TOOL_CALLS] docs for details. More prose after.",
"<|python_tag|> is the Llama marker. Explanation continues.",
"The <|tool_call> opener wraps Gemma calls.",
):
assert strip_tool_markup(text, final = True) == text
# A bare marker at end-of-text is a fragment and still drops.
assert strip_tool_markup("Answer text [TOOL_CALLS]", final = True) == "Answer text"
def test_final_strip_still_drops_truncated_marker_calls():
from core.inference.tool_call_parser import strip_tool_markup
for text in (
'[TOOL_CALLS][{"name":"web_search","argu',
'[TOOL_CALLS]web_search[ARGS]{"q":"x',
'<|python_tag|>{"name":"web_search","par',
'<|python_tag|>foo.call(items=["a',
"<|tool_call>call:web_search{query:tru",
):
assert strip_tool_markup(text, final = True) == ""
def test_chained_bare_json_strip_consumes_all_calls():
# The loops keep this text as next-turn history: a leftover executed call
# would be replayed alongside the structured tool_calls.
from core.inference.tool_call_parser import strip_leading_bare_json_call
enabled = {"web_search", "python"}
chained = (
'{"name":"web_search","parameters":{"q":"first"}};'
'{"name":"python","parameters":{"code":"x"}}'
)
assert strip_leading_bare_json_call(chained, enabled_tool_names = enabled) == ""
assert (
strip_leading_bare_json_call(chained + " trailing prose", enabled_tool_names = enabled)
== "trailing prose"
)
# The chain stops at a non-call answer object, which stays visible.
call_then_answer = (
'{"name":"web_search","parameters":{"q":"x"}};{"name":"web_search","result":"data"}'
)
assert (
strip_leading_bare_json_call(call_then_answer, enabled_tool_names = enabled)
== '{"name":"web_search","result":"data"}'
)
Reference in a new issue View git blame Copy permalink