ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704)

Browse source 
* 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>
This commit is contained in:
Daniel Han 2026-07-06 18:52:13 -07:00 • committed by GitHub
commit f109e7f0e6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 3115 additions and 270 deletions
Download patch file Download diff file Expand all files Collapse all files

208
studio/backend/core/inference/llama_cpp.py
View file

@ -44,7 +44,8 @@ from core.inference.llama_server_args import (
from core.inference.tool_call_parser import (
_GEMMA_BARE_TC_PREFIX_RE,
_GEMMA_BARE_TC_RE,
_TOOL_ALL_PATS,
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
_balanced_brace_end,
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
@ -58,6 +59,16 @@ from core.inference.tool_call_parser import (
strip_llama3_leading_sentinels,
strip_tool_markup as _shared_strip_tool_markup,
)
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
# pattern lists, so the GGUF streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_strip_bracket_tag_calls,
apply_tool_strip_patterns,
strip_outside_think,
strip_tool_call_markup,
)
from utils.native_path_leases import child_env_without_native_path_secret
from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
from utils.subprocess_compat import (
@ -256,6 +267,72 @@ _FINAL_ANSWER_SIGNAL = re.compile(
)
def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in (active_tools or [])
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the
# Mistral [CALL_ID]...[ARGS] shape.
_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]")
def _gguf_rehearsal_signal_pos(text: str, active_tools: list[dict]) -> int:
"""Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A
bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors
``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode)."""
active = set(_gguf_active_tool_names(active_tools))
if not active:
return -1
for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text):
if m.group(1) in active:
return m.start()
return -1
def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool:
"""True when ``text`` holds a genuine tool-call boundary for one of ``signals``.
Unambiguous markers (``<tool_call>``, ``[TOOL_CALLS]``, ``<function=``) count on a
plain substring hit; an ``[ARGS]`` hit is genuine only when an active tool name
precedes it, so inactive-name prose is neither drained nor parsed."""
for sig in signals:
if sig == "[ARGS]":
if _gguf_rehearsal_signal_pos(text, active_tools) >= 0:
return True
continue
if sig in text:
return True
return False
def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool:
"""True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an
active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``.
Mirrors the safetensors loop so the split rehearsal call is not streamed."""
if not stripped or any(ch.isspace() for ch in stripped):
return False
for name in _gguf_active_tool_names(active_tools):
if stripped == name or f"{name}[ARGS]".startswith(stripped):
return True
return False
def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int:
"""Length of a trailing bare tool-name token that may be a split rehearsal call
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors."""
i = len(text)
while i > 0 and not text[i - 1].isspace():
i -= 1
tail = text[i:]
return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0
def _is_short_intent_without_action(text: str) -> bool:
stripped = text.strip()
return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None
@ -8418,6 +8495,13 @@ class LlamaCppBackend:
_reasoning_started_at: Optional[float] = None
_reasoning_summary_emitted = False
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate.
_enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
_detect_tools = list(tools or [])
def _reasoning_summary_event(started_at: float) -> dict:
return {
"type": "reasoning_summary",
@ -8436,25 +8520,42 @@ class LlamaCppBackend:
) -> str:
if not (auto_heal_tool_calls or force):
return text
# Delegate to the shared parser-side strip so the GGUF cleanup covers every family the
# parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function
# XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only
# forms (python_tag / Mistral name) and would leak them into display.
return _shared_strip_tool_markup(
text, final = final, enabled_tool_names = _enabled_tool_names
text, final = final, enabled_tool_names = _enabled_names_gate
)
def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
if not (auto_heal_tool_calls or force):
return text
# Shared parser patterns (not the legacy tool_healing set) so textual
# Mistral/python_tag calls entering DRAINING never leak. Balanced strips
# first (nested JSON removed whole); no final trim so length compares hold.
text = _strip_mistral_closed_calls(text)
text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names)
# Parser-accurate scans close at each call's REAL terminator before
# the regex arms: literal markup inside a value is data.
text = _strip_function_xml_calls(text, final = True)
text = _strip_glm_calls(text, final = True)
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
def _seg(segment: str, is_last: bool) -> str:
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
# strips first (nested JSON removed whole; literal markup inside a value is that
# call's data), then the guarded function-XML / GLM scans, then the regex arms
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal + markerless
# strips are name-gated on the ORIGINAL list (strip/detect aligned).
seg = _strip_mistral_closed_calls(segment)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate)
seg = _strip_function_xml_calls(seg, final = is_last)
seg = _strip_glm_calls(seg, final = is_last)
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if is_last:
seg = apply_tool_strip_patterns(
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate
)
return seg
# Preserve think blocks verbatim (a rehearsed call inside one must not be deleted).
return strip_outside_think(text, _seg)
def _build_metadata_event(usage, timings, finish_reason):
"""Final usage+timings metadata event for the given pass, merging its
@ -8814,12 +8915,18 @@ class LlamaCppBackend:
in_thinking = False
cumulative_display += token
cleaned = _strip_tool_markup_streaming(cumulative_display)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
# Hold a trailing bare active-tool-name (split rehearsal)
# until [ARGS] arrives; released by later prose or stream end.
_hold = _held_rehearsal_tail_len(cleaned, _detect_tools)
_emit = (
cleaned[: len(cleaned) - _hold] if _hold else cleaned
)
if len(_emit) > len(_last_emitted):
_last_emitted = _emit
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
"text": _emit,
}
elif detect_state == _S_BUFFERING:
@ -8828,7 +8935,8 @@ class LlamaCppBackend:
if not stripped_buf:
continue
# Check tool signal prefixes.
# Bracket tags arrive mid-buffer, so substring-check too;
# ``[ARGS]`` counts only as a regex-matched NAME[ARGS].
is_prefix = False
is_match = False
for sig in _tool_xml_signals:
@ -8838,6 +8946,31 @@ class LlamaCppBackend:
if sig.startswith(stripped_buf):
is_prefix = True
break
if sig == "[ARGS]":
# Active NAME[ARGS] only; inactive-name prose
# is gated out, not drained/parsed.
if (
_gguf_rehearsal_signal_pos(
stripped_buf, _detect_tools
)
>= 0
):
is_match = True
break
elif sig.startswith("[") and sig in stripped_buf:
is_match = True
break
# Split rehearsal: hold the bare name until
# its [ARGS] arrives and matches above.
is_rehearsal_prefix = False
if (
not is_match
and not is_prefix
and _is_rehearsal_prefix(stripped_buf, _detect_tools)
):
is_prefix = True
is_rehearsal_prefix = True
# Signal-less call shapes (mirror the safetensors
# loop): Llama-3.2 bare {"name":..} and Gemma
@ -8884,9 +9017,14 @@ class LlamaCppBackend:
# Tool signal -- flush any visible
# prefix before DRAINING so the
# route sends it before tool_start.
# Use the final strip (all families incl. Llama
# <|python_tag|> / Mistral name): the buffer holds
# the whole call, so a streaming closed-only strip
# would leak its open-ended markup as display text.
_flush_reasoning_and_buffer()
cleaned = _strip_tool_markup_streaming(
cleaned = _strip_tool_markup(
cumulative_display,
final = True,
force = True,
)
if len(cleaned) > len(_last_emitted):
@ -8898,8 +9036,14 @@ class LlamaCppBackend:
}
detect_state = _S_DRAINING
elif _hold_buffer or (
is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
is_prefix
and (
is_rehearsal_prefix
or len(stripped_buf) < _MAX_BUFFER_CHARS
)
):
# A rehearsal prefix is self-bounded; the buffer
# cap must not cut long MCP names short.
pass # keep buffering
else:
# Not a tool -- flush buffer
@ -8910,12 +9054,20 @@ class LlamaCppBackend:
cleaned = _strip_tool_markup(
cumulative_display,
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
# Same trailing-name hold as STREAMING for this
# first flush out of BUFFERING.
_hold = _held_rehearsal_tail_len(cleaned, _detect_tools)
_emit = (
cleaned[: len(cleaned) - _hold]
if _hold
else cleaned
)
if len(_emit) > len(_last_emitted):
_last_emitted = _emit
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
"text": _emit,
}
except json.JSONDecodeError:
@ -8933,7 +9085,9 @@ class LlamaCppBackend:
_is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json(
_bare_eos, _enabled_tool_names
)
if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
if stripped_buf and _gguf_has_genuine_tool_signal(
stripped_buf, _tool_xml_signals, _detect_tools
):
detect_state = _S_DRAINING
elif _is_bare_tc:
detect_state = _S_DRAINING
@ -9066,6 +9220,12 @@ class LlamaCppBackend:
"type": "content",
"text": forced_visible_text,
}
elif not _suppress_visible_output:
# Turn ended as a plain answer (no [ARGS] followed): the held
# rehearsal tail is real prose, release it.
_final_clean = _strip_tool_markup_streaming(cumulative_display)
if len(_final_clean) > len(_last_emitted):
yield {"type": "content", "text": _final_clean}
# Content was already streamed. Yield metadata.
yield {"type": "status", "text": ""}

64
studio/backend/core/inference/passthrough_healing.py
View file

@ -32,16 +32,15 @@ from typing import Any, Optional
from core.inference.tool_loop_controller import coerce_tool_arguments
from core.tool_healing import parse_tool_calls_from_text
# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing)
# actually promotes. The parser module's broader signal list also covers Llama
# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose
# full parser handles them; buffering those here would hold a streamed
# client-tool call until finalization and then flush it as prose (this healer
# cannot promote them), so the passthrough keeps its own aligned list.
# Only the formats this healer's parser can promote -- narrower than the loops'
# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare
# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a
# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in.
_HEAL_SIGNALS = (
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
)
@ -356,9 +355,10 @@ class StreamToolCallHealer:
events.append(("text", emit))
self._buffer = self._buffer[len(self._buffer) - keep :]
return events
# HOLD: handle the FIRST complete block per pass so events keep
# document order (a later declared call must not overtake an
# earlier undeclared one flushing as text).
# HOLD: drain the first contiguous run per pass so events keep document
# order (a later declared call must not overtake an earlier undeclared one
# flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS]
# array of contiguous spans, so later calls in it are not stranded as text.
parsed, spans = parse_tool_calls_from_text(
self._buffer,
id_offset = self._id_offset,
@ -379,26 +379,32 @@ class StreamToolCallHealer:
self._holding = False
continue
return events
start, end = spans[0]
promoted = _promote(
[parsed[0]],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
if start:
events.append(("text", self._buffer[:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
# Drop exactly the promoted markup span; everything else
# (leading text, later blocks) stays and is rescanned.
self._buffer = self._buffer[end:]
else:
# Undeclared or unusable name: its markup is DATA, flush it
# (and anything before it) verbatim, then rescan the rest.
events.append(("text", self._buffer[:end]))
self._buffer = self._buffer[end:]
pos = 0
run_end = spans[0][1]
for order, (call, (start, end)) in enumerate(zip(parsed, spans)):
# Stop at the first gap or incomplete trailing block: leave it for the
# next pass to re-hold and stream incrementally, not flush as text early.
if order and start != run_end:
break
promoted = _promote(
[call],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
# Flush any leading text, then drop the promoted markup span.
if self._buffer[pos:start]:
events.append(("text", self._buffer[pos:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
else:
# Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim.
events.append(("text", self._buffer[pos:end]))
pos = end
run_end = end
# Everything past the drained run (later blocks) stays and is rescanned.
self._buffer = self._buffer[run_end:]
self._holding = False
def finalize(self) -> list:

394
studio/backend/core/inference/safetensors_agentic.py
View file

@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via
``core.inference.tools``.
"""
import bisect
import re
import threading
from typing import Callable, Generator, Optional
@ -23,7 +24,8 @@ from loggers import get_logger
from core.inference.tool_call_parser import (
_GEMMA_BARE_TC_PREFIX_RE,
_GEMMA_BARE_TC_RE,
_TOOL_ALL_PATS,
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
_balanced_brace_end,
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
@ -39,6 +41,16 @@ from core.inference.tool_call_parser import (
strip_llama3_leading_sentinels,
strip_tool_markup,
)
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_strip_bracket_tag_calls,
_think_spans_outside_tool_markup,
apply_tool_strip_patterns,
strip_outside_think,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
coerce_tool_arguments,
@ -93,6 +105,147 @@ def _active_tool_names(active_tools: list[dict]) -> list[str]:
return [name for name in names if name]
def _active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal;
# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held.
_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?")
def _is_rehearsal_prefix(
stripped: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]``
rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space
means prose. Unrestricted mode accepts any identifier; else NAME must be active."""
if not stripped or any(ch.isspace() for ch in stripped):
return False
if unrestricted:
return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None
for name in _active_tool_names(active_tools):
if stripped == name or f"{name}[ARGS]".startswith(stripped):
return True
return False
def _held_rehearsal_tail_len(
text: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Length of a trailing bare tool-name token that may be a split rehearsal call
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
instead of leaking the name. Returns 0 for ordinary prose."""
i = len(text)
while i > 0 and not text[i - 1].isspace():
i -= 1
tail = text[i:]
return (
len(tail)
if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted)
else 0
)
def _rehearsal_name_start(
candidate: str,
signal_pos: int,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding
bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the
signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode)."""
if not candidate.startswith("[ARGS]", signal_pos):
return signal_pos
j = signal_pos
while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"):
j -= 1
if j < signal_pos and (
unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools)
):
return j
return signal_pos
def _earliest_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Index where the turn's first genuine tool-call boundary begins, or -1.
Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal
only when an active tool name (any name in unrestricted mode) precedes it, so a
literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a
real ``NAME[ARGS]`` the boundary is pulled back to NAME."""
best = -1
for sig in signals:
if sig != "[ARGS]":
p = candidate.find(sig)
if p >= 0 and (best < 0 or p < best):
best = p
continue
from_idx = 0
while True:
p = candidate.find("[ARGS]", from_idx)
if p < 0:
break
name_start = _rehearsal_name_start(
candidate, p, active_tools, unrestricted = unrestricted
)
if name_start < p:
# Genuine ``NAME[ARGS]``: the boundary is the start of NAME.
if best < 0 or name_start < best:
best = name_start
break
# Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found.
from_idx = p + len("[ARGS]")
return best
def _has_genuine_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``.
Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only
when an active tool name (any in unrestricted mode) precedes it. Mirrors the
``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not
drain inactive-name prose."""
for sig in signals:
if sig == "[ARGS]":
if (
_earliest_tool_signal(
candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted
)
>= 0
):
return True
continue
if sig in candidate:
return True
return False
def strip_tool_markup_streaming(
text: str,
*,
@ -101,25 +254,46 @@ def strip_tool_markup_streaming(
enabled_tool_names: Optional[set] = None,
) -> str:
"""Strip open-ended tool XML from display text without trimming whitespace.
``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a
disabled/example name in prose is kept (mirrors the parser gate)."""
Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so
streaming and final display agree: balanced strips first (nested JSON removed whole),
then the guarded function-XML / GLM scans that close at each call's REAL terminator so
literal markup inside argument values is data and trailing prose survives. Reasoning
``<think>`` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must
not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names``
keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose,
not a call), matching the parse / detection active-tool gate."""
if not (auto_heal_tool_calls or tool_protocol_active):
return text
# Mirror the final strip's scan order so streaming and final display agree:
# balanced strips first (nested JSON removed whole), then the guarded
# function-XML/GLM scans that close at each call's REAL terminator, so literal
# markup inside argument values is data and trailing prose survives. No final
# trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK]
# is dropped (bracket form, not the reasoning channel's <think>); an unclosed
# [THINK] holds until [/THINK] so the cleaned text stays monotonic.
# Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the
# ``<think>`` channel) so raw reasoning does not leak into streamed display; an unclosed
# leading block is held (dropped to EOF) until its closer streams in.
text = _strip_mistral_reasoning(text)
text = _strip_mistral_closed_calls(text)
text = _strip_gemma_wrapperless_calls(text, enabled_tool_names)
text = _strip_function_xml_calls(text, final = True)
text = _strip_glm_calls(text, final = True)
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
def _seg(segment: str, is_last: bool) -> str:
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
# strips first, then the guarded function-XML / GLM scans, then the regex arms
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal strips are name-gated.
seg = _strip_mistral_closed_calls(segment)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
seg = _strip_function_xml_calls(seg, final = is_last)
seg = _strip_glm_calls(seg, final = is_last)
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if is_last:
seg = apply_tool_strip_patterns(
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names
)
return seg
# Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then
# regrows the cumulative text, corrupting append-by-length consumers.
return strip_outside_think(text, _seg)
def _strip_tool_markup_final(
@ -149,23 +323,66 @@ def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set])
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional
# render-html card fires for bracket-tag serializations too.
_MISTRAL_RENDER_NAME_RE = re.compile(
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
)
_REHEARSAL_RENDER_NAME_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
def _detect_render_html_tool_start(content: str) -> bool:
"""Return True when the first drained tool call is clearly render_html."""
function_match = _FUNCTION_SIGNAL_RE.search(content)
tool_call_index = content.find("<tool_call>")
if not function_match and tool_call_index < 0:
"""Return True when the FIRST tool call in ``content`` is clearly render_html.
Covers every serialization the loop executes (XML ``<function=>`` / ``<tool_call>``,
Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a
render_html marker inside another call's argument is treated as data. Markers inside
a ``<think>`` / ``[THINK]`` block are dropped since the parser skips them."""
think_spans = _think_spans_outside_tool_markup(content)
_think_starts = [s for s, _e in think_spans]
def _in_think(pos: int) -> bool:
if not think_spans:
return False
i = bisect.bisect_right(_think_starts, pos) - 1
return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1]
def _first_outside(start: int, finder) -> int:
# First occurrence at/after ``start`` that is not inside a think span.
pos = finder(start)
while pos >= 0 and _in_think(pos):
pos = finder(pos + 1)
return pos
candidates: list[tuple[int, str]] = []
for fm in _FUNCTION_SIGNAL_RE.finditer(content):
if not _in_think(fm.start()):
candidates.append((fm.start(), fm.group(1)))
break
tc = _first_outside(0, lambda i: content.find("<tool_call>", i))
if tc >= 0:
nm = _TOOL_CALL_NAME_RE.search(content[tc:])
candidates.append((tc, nm.group(1) if nm else ""))
mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i))
if mt >= 0:
mm = _MISTRAL_RENDER_NAME_RE.match(content, mt)
if mm:
candidates.append((mt, mm.group(1)))
else:
# Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the
# first call through the parser (it reads top-level names).
arr_calls = parse_tool_calls_from_text(content[mt:])
if arr_calls:
candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or ""))
for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content):
if not _in_think(rm.start(1)):
candidates.append((rm.start(1), rm.group(1)))
break
if not candidates:
return False
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
return function_match.group(1) == "render_html"
if tool_call_index >= 0:
name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:])
return bool(name_match and name_match.group(1) == "render_html")
return False
_pos, name = min(candidates, key = lambda c: c[0])
return name == "render_html"
def _coerce_arguments_with_provenance(
@ -256,6 +473,12 @@ def run_safetensors_tool_loop(
conversation.extend(_auto["messages"])
unrestricted_tools = not tools
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted.
_enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools))
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
_detect_tools = [] if unrestricted_tools else list(tools or [])
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
@ -381,18 +604,18 @@ def run_safetensors_tool_loop(
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in tool_xml_signals:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
# Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is
# pulled back to NAME so the name is not flushed.
signal_pos = _earliest_tool_signal(
candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools
)
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup_streaming(
before_tool,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
@ -423,11 +646,20 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives;
# released by later prose or the end-of-stream flush.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
continue
# BUFFERING: hold until we know it is not a tool call.
@ -445,6 +677,34 @@ def run_safetensors_tool_loop(
if sig.startswith(stripped):
is_prefix = True
break
# Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS]
# counts only with an active NAME so prose is not drained into a no-op.
if sig == "[ARGS]":
if (
_earliest_tool_signal(
stripped,
("[ARGS]",),
_detect_tools,
unrestricted = unrestricted_tools,
)
>= 0
):
is_match = True
break
elif sig.startswith("[") and sig in stripped:
is_match = True
break
# Split rehearsal: hold the bare name until its [ARGS] arrives and matches above.
is_rehearsal_prefix = False
if (
not is_match
and not is_prefix
and tool_protocol_active
and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools)
):
is_prefix = True
is_rehearsal_prefix = True
# Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML
# signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses
@ -512,7 +772,7 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
@ -536,7 +796,8 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
else:
detect_state = _state_streaming
@ -545,24 +806,38 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Same trailing-name hold as STREAMING for this first flush out of BUFFERING.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
# Stream finished -- resolve what we collected.
if cancel_event is not None and cancel_event.is_set():
return
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content?
# Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal
# ``foo[ARGS]{...}`` is not parsed.
stripped = content_buffer.lstrip()
_bare_eos = strip_llama3_leading_sentinels(stripped)
if (
stripped
and tool_protocol_active
and any(sig in stripped for sig in tool_xml_signals)
and _has_genuine_tool_signal(
stripped,
tool_xml_signals,
_detect_tools,
unrestricted = unrestricted_tools,
)
):
detect_state = _state_draining
elif tool_protocol_active and _looks_like_enabled_bare_json(
@ -629,6 +904,17 @@ def run_safetensors_tool_loop(
# in full; route-level cleanup still applies the Auto-Heal policy.
if content_accum and any(sig in content_accum for sig in tool_xml_signals):
yield {"type": "content", "text": content_accum}
else:
# Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real
# prose, release it.
final_clean = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(final_clean) > len(last_emitted):
yield {"type": "content", "text": final_clean}
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
@ -636,19 +922,23 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
)
else:
# DRAINING: parse tool calls out of full content.
# DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the
# ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to
# drain here: a spent one-shot (render_html) is off the active list but its re-emitted
# ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of
# being dropped into a blank continuation.
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if not tool_calls:
# Parser found nothing. Auto-Heal-enabled display cleanup
@ -682,7 +972,7 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_tool_names,
enabled_tool_names = _enabled_names_gate,
)
if tool_calls:

116
studio/backend/core/inference/tool_call_parser.py
View file

@ -41,6 +41,9 @@ TOOL_XML_SIGNALS = (
"<|python_tag|>",
"[TOOL_CALLS]",
"<|tool_call>",
# Bare reasoning-rehearsal marker (``name[ARGS]{...}``, no leading [TOOL_CALLS]);
# keeps a rehearsed call held in the stream so it is promoted, not leaked as prose.
"[ARGS]",
# DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps.
"<｜tool▁calls▁begin｜>",
"<｜tool▁call▁begin｜>",
@ -90,7 +93,7 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
# follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops.
re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL),
re.compile(
r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$",
r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*(?:[\[{]|\s*$))|\s*$).*$",
re.DOTALL,
),
re.compile(
@ -572,26 +575,51 @@ def strip_tool_markup(
enabled_tool_names: Optional[set] = None,
) -> str:
"""Strip tool-call markup. ``final=False`` keeps in-progress markup buffered;
``final=True`` also drops trailing unclosed runs and trims. ``enabled_tool_names``
gates the markerless Gemma ``call:NAME{...}`` strip so a disabled/example name in
prose is kept (mirrors the parser gate); ``None`` strips every closed call."""
``final=True`` also drops trailing unclosed runs and trims.
``enabled_tool_names`` gates the name-conditioned forms so a disabled/example name in
prose is kept (mirrors the parser gate): the bare reasoning-rehearsal ``name[ARGS]{...}``
and the markerless Gemma ``call:NAME{...}`` strip. ``None`` strips every closed call.
"""
if final:
# Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket
# form is not the ``<think>`` the reasoning channel renders.
text = _strip_mistral_reasoning(text)
text = _strip_mistral_closed_calls(text)
if final:
text = _strip_gemma_wrapperless_calls(text, enabled_tool_names)
# Scan-strip the function-XML form (a literal ``<function=...>`` inside a value is
# data). The regex arms below cover the other formats but no-op on function calls here.
text = _strip_function_xml_calls(text, final = final)
# GLM 4.x: scan to the call's real </tool_call> so a literal one inside a value is data,
# not a leak. Qwen <tool_call>{json} is left to the regex arms.
text = _strip_glm_calls(text, final = final)
pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
for pat in pats:
text = pat.sub("", text)
return text.strip() if final else text
def _strip_segment(segment: str, is_last: bool) -> str:
seg_final = final and is_last
seg = _strip_mistral_closed_calls(segment)
# Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through
# the shared balanced scan, so strip them the same way (any nesting depth removed whole).
# The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept.
seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if seg_final:
# Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only.
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
# Scan-strip the function-XML form (parser-accurate: a literal ``<function=...>`` in a
# value is data, not a call); the regex arms below cover the other formats.
seg = _strip_function_xml_calls(seg, final = seg_final)
# GLM 4.x: scan to the call's real </tool_call> so a literal one inside a value is data,
# not a leak. Qwen <tool_call>{json} is left to the regex arms.
seg = _strip_glm_calls(seg, final = seg_final)
pats = _TOOL_ALL_PATS if seg_final else _TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if seg_final:
# Drop a trailing partial bare rehearsal (``name[ARGS]`` with a truncated or absent
# body) the balanced scan cannot close; gated so prose ``foo[ARGS] ...`` survives.
seg = _tool_healing.apply_tool_strip_patterns(
seg,
[_tool_healing._REHEARSAL_TAIL_STRIP_RE],
enabled_tool_names = enabled_tool_names,
)
return seg
# ``<think>`` / ``[THINK]`` reasoning is preserved verbatim (a rehearsed call inside it is
# not executed, so it must not be stripped from display either); a literal think marker
# inside a real call's arguments is that call's data and is stripped with the call.
result = _tool_healing.strip_outside_think(text, _strip_segment)
return result.strip() if final else result
def has_tool_signal(text: str) -> bool:
@ -711,6 +739,41 @@ def _xml_signal_inside_leading_mistral(content: str) -> bool:
return _mistral_region_end(content, trig) is not None
def _parse_bare_rehearsals(
content: str,
*,
id_offset: int = 0,
enabled_tool_names: Optional[set] = None,
) -> list[dict]:
"""Promote bare reasoning-rehearsal ``name[ARGS]{json}`` calls that a leading [TOOL_CALLS]
owns-the-turn parse would miss. Only the ``rehearsal`` kind is taken (a Mistral
``[TOOL_CALLS]name[ARGS]{..}`` yields ``name`` and is not double-counted), and a rehearsal
inside a ``<think>`` / ``[THINK]`` block is reasoning, so it is skipped."""
out: list[dict] = []
think_spans = _tool_healing._think_spans_outside_tool_markup(content)
for start, end, kind, m in _tool_healing._iter_bracket_spans(
content, enabled_tool_names = enabled_tool_names
):
if kind != "rehearsal":
continue
if any(s <= start < e for s, e in think_spans):
continue
try:
payload = json.loads(content[m.end() : end])
except (json.JSONDecodeError, ValueError):
continue
if not isinstance(payload, dict):
continue
out.append(
{
"id": f"call_{id_offset + len(out)}",
"type": "function",
"function": {"name": m.group(1), "arguments": json.dumps(payload)},
}
)
return out
_ATTR_FUNC_OPEN_RE = re.compile(r'<function\s+name="')
@ -919,6 +982,16 @@ def parse_tool_calls_from_text(
content, id_offset = id_offset, allow_incomplete = allow_incomplete
)
if calls:
# A bare rehearsal ``name[ARGS]{..}`` after the Mistral call is a peer tool call,
# not foreign XML the owns-the-turn guard protects against: promote it too so a
# Mistral call and a rehearsal in one message both parse.
calls.extend(
_parse_bare_rehearsals(
content,
id_offset = id_offset + len(calls),
enabled_tool_names = enabled_tool_names,
)
)
return calls
# DeepSeek/Kimi markers are unique, so try them first -- unless an outer envelope
@ -984,13 +1057,16 @@ def parse_tool_calls_from_text(
if calls:
return calls
# Qwen/Hermes, Qwen3.5 XML, and Gemma 4 go through the shared tool_healing
# parser (strict/Auto-Heal contract + nested-marker, trailing-prose, and
# ``<|"|>`` quoted-string handling the GGUF path relies on).
# Qwen/Hermes, Qwen3.5 XML, Gemma 4, plus Mistral [TOOL_CALLS] / bare rehearsal
# ``name[ARGS]{json}`` use the shared tool_healing parser (strict/Auto-Heal contract +
# nested-marker, trailing-prose, and ``<|"|>`` quoted-string handling the GGUF path
# relies on). ``enabled_tool_names`` gates the ambiguous bare-rehearsal form so an
# inactive ``foo[ARGS]{..}`` stays prose.
calls = _tool_healing.parse_tool_calls_from_text(
content,
id_offset = id_offset,
allow_incomplete = allow_incomplete,
enabled_tool_names = enabled_tool_names,
)
if calls:
return calls

521
studio/backend/core/tool_healing.py
View file

@ -1,38 +1,91 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge
# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026
# Antoine Zambelli, used under the MIT License.
"""Lightweight tool-call XML parsing and stripping helpers.
"""Lightweight tool-call parsing and stripping helpers.
External inference servers import this module without pulling in the inference
orchestrator, structlog, httpx, or the rest of the studio backend.
orchestrator, structlog, httpx, or the rest of the studio backend. Kept in
lockstep with ``core/inference/tool_call_parser.py`` so those servers
(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any
change here must also land there.
Handles these serializations (see ``parse_tool_calls_from_text``):
* ``<tool_call>{json}</tool_call>``
* ``<|tool_call>call:name{...}<tool_call|>`` (Gemma)
* ``<function=name><parameter=k>v</parameter></function>``
* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback)
* ``name[ARGS]{json}`` (reasoning-model rehearsal)
"""
# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9).
from __future__ import annotations
import bisect
import json
import re
# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs
# strip first so a closed call goes as a unit before any to-EOF sweep reaches
# nested markup; only the final list adds the .*$ EOF sweeps.
# One nesting level in the strip regexes; deeper may leak markup (still parsed).
_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"
# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed =
# complete body, tail = truncated; ``(?<!\[CALL_ID\])`` keeps the v11 call-id from reading as a name.
_REHEARSAL_CLOSED_STRIP_RE = re.compile(
r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*" + _BRACKETED_JSON_ONE_LEVEL, re.DOTALL
)
_REHEARSAL_TAIL_STRIP_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?:\{.*)?$", re.DOTALL)
# Tool-XML strip patterns; hyphen in the name class covers dashed MCP names.
# Closed-pair patterns are named so _PAT_REQUIRED_TOKEN can skip a doomed lazy rescan when
# the close token is absent: an unguarded ``<tag>.*?</tag>`` rescans to EOF from every opener
# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass.
_TC_JSON_CLOSED_PAT = re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL)
_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?<tool_call\|>", re.DOTALL)
_TC_FUNC_CLOSED_PAT = re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL)
_TC_GEMMA_END_PAT = re.compile(r"<tool_call\|>")
_TOOL_CLOSED_PATS = [
_TC_JSON_CLOSED_PAT,
_TC_GEMMA_CLOSED_PAT,
re.compile(r"<tool_call\|>"),
_TC_FUNC_CLOSED_PAT,
_TC_GEMMA_END_PAT,
# Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata.
re.compile(
r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*"
+ _BRACKETED_JSON_ONE_LEVEL,
re.DOTALL,
),
_REHEARSAL_CLOSED_STRIP_RE,
# Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind.
re.compile(r"\[/TOOL_CALLS\]"),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<\|tool_call>.*$", re.DOTALL),
# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF
# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by
# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data).
_TOOL_OPEN_XML_TAIL_PATS = [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<\|tool_call>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in
# their argument data cannot make the helper truncate the block and its tail.
_TOOL_ALL_PATS = (
_TOOL_CLOSED_PATS
+ _TOOL_OPEN_XML_TAIL_PATS
+ [
re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL),
_REHEARSAL_TAIL_STRIP_RE,
]
)
# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None.
_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE})
# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument
# data cannot make the helper truncate the block and its tail.
_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT]
# A lazy closed-pair pattern whose close token is absent rescans to EOF from
# every opener (quadratic, re-run per streamed token); skip that doomed pass.
# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every
# opener; skip that doomed (quadratic) pass. Shared by both strip helpers.
_PAT_REQUIRED_TOKEN = {
_TC_JSON_CLOSED_PAT: "</tool_call>",
_TC_GEMMA_CLOSED_PAT: "<tool_call|>",
@ -50,26 +103,102 @@ def strip_tool_patterns(text: str, patterns) -> str:
return text
def apply_tool_strip_patterns(
text: str,
patterns,
enabled_tool_names = None,
) -> str:
"""Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern
strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is
``None``); every other pattern is removed unconditionally. A closed-pair pattern whose
close token is absent is skipped so an unclosed-marker stream stays linear."""
for pat in patterns:
token = _PAT_REQUIRED_TOKEN.get(pat)
if token is not None and token not in text:
continue
if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS:
text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text)
else:
text = pat.sub("", text)
return text
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around
# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``.
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline).
# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>[^\S\n]*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_GEMMA_QUOTE = '<|"|>'
_PARAM_CLOSE_TAG = "</parameter>"
_FUNC_CLOSE_TAG = "</function>"
# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next
# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`,
# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value.
# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next
# `key:` pair. A comma NOT followed by a key token is part of the value (e.g.
# `location:New York, NY`), so it must not terminate the value. The key token
# must be identifier-shaped (start with a letter or underscore); a comma
# followed by digits-then-colon is value text such as a timestamp or ratio
# (`meet at 10:00, 11:00 tomorrow`), not a new key.
_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:")
# A candidate starting inside a think block is a rehearsal (block kept so literal tags in
# real args survive); ``$`` accepts an unclosed block mid-stream.
_THINK_TAG_RE = re.compile(r"<think>.*?(?:</think>|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL)
# Bare open/close markers for prefilled-reasoning turns (template opens <think> in the prompt).
_THINK_OPEN_RE = re.compile(r"<think>|\[THINK\]")
_THINK_CLOSE_RE = re.compile(r"</think>|\[/THINK\]")
# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects.
_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)")
# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name,
# and hyphens keep dashed MCP names whole.
_MISTRAL_BRACKET_RE = re.compile(
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
)
# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id
# from being taken as the function name.
_REHEARSAL_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
# Above this size skip the balanced scan; the linear regex catch-all bounds pathological output.
_MAX_BRACKET_SCAN_CHARS = 1_000_000
def _balanced_json_span(text: str, start: int) -> int | None:
"""Return the end index of a balanced JSON object opening at ``start``,
or ``None`` if the braces don't balance. Honors escapes and strings.
"""
if start >= len(text) or text[start] != "{":
return None
depth = 0
in_string = False
escape = False
for j in range(start, len(text)):
ch = text[j]
if escape:
escape = False
continue
if ch == "\\":
escape = True
continue
if in_string:
if ch == '"':
in_string = False
continue
if ch == '"':
in_string = True
elif ch == "{":
depth += 1
elif ch == "}":
depth -= 1
if depth == 0:
return j
return None
def _balanced_brace_end(
content: str,
@ -134,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int:
return -1
def _decode_array_items(text: str, body_start: int, body_end: int):
"""Return ``(objs, ends)`` for each top-level element of the JSON array between
``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded
object and its absolute exclusive end offset.
Decoding element-by-element with ``raw_decode`` tolerates the comma-less object
separators the repo's own Mistral/Ollama multi-call templates emit
(``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the
whole body rejects that form and would drop every call. The ends also tile the
region across the calls' spans so a with_spans consumer strips each exactly once."""
decoder = json.JSONDecoder()
objs: list = []
ends: list[int] = []
i = text.find("[", body_start)
if i < 0:
return objs, ends
i += 1
while i < body_end:
while i < body_end and text[i] in " \t\r\n,":
i += 1
if i >= body_end or text[i] == "]":
break
try:
obj, rel = decoder.raw_decode(text[i:body_end])
except (json.JSONDecodeError, ValueError):
break
i += rel
objs.append(obj)
ends.append(i)
return objs, ends
def _iter_bracket_spans(
text: str,
start: int = 0,
enabled_tool_names = None,
):
"""Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag
call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is
``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11
[CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}).
``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous
bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a
prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit
[TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric.
Balance-only (no JSON validation) so strip and parse share one scan. The cursor
jumps past each consumed span, so a marker inside consumed JSON is never
re-matched and each regex re-searches only once its match falls behind: linear."""
n = len(text)
specs = (
("array", _MISTRAL_ARRAY_RE),
("name", _MISTRAL_BRACKET_RE),
("rehearsal", _REHEARSAL_RE),
)
nexts = {kind: rx.search(text, start) for kind, rx in specs}
cursor = start
while cursor < n:
for kind, rx in specs:
m = nexts[kind]
if m is not None and m.start() < cursor:
nexts[kind] = rx.search(text, cursor)
live = [(kind, m) for kind, m in nexts.items() if m is not None]
if not live:
return
kind, m = min(live, key = lambda km: km[1].start())
if kind == "array":
end = _balanced_bracket_end(text, m.end())
end = None if end < 0 else end
else:
end = _balanced_json_span(text, m.end())
if end is None:
# Truncated body: skip and keep scanning; the caller's catch-all strips the tail.
cursor = m.end()
continue
if (
kind == "rehearsal"
and enabled_tool_names is not None
and m.group(1) not in enabled_tool_names
):
# Inactive-name rehearsal is prose: advance past its body without yielding.
cursor = end + 1
continue
yield (m.start(), end + 1, kind, m)
cursor = end + 1
def _split_top_level_commas(src: str) -> list:
"""Split on commas that are not inside a nested ``[]``/``{}`` or a string."""
parts: list[str] = []
@ -164,8 +381,14 @@ def _split_top_level_commas(src: str) -> list:
def _quote_gemma_array_elements(body: str) -> str:
"""Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds:
quote bare strings, recurse into objects/arrays, keep quoted/JSON literals."""
"""Normalise the elements of a Gemma array value so json.loads succeeds.
Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of
objects (``items:[{path:a}]``) whose keys/values also lack quotes; left
as-is json.loads fails and the whole call is dropped. Bare string elements
are quoted, object and nested-array elements are normalised recursively, and
quoted strings (already normalised from ``<|"|>``), numbers, and JSON
literals are preserved."""
out: list[str] = []
for element in _split_top_level_commas(body):
stripped = element.strip()
@ -173,9 +396,11 @@ def _quote_gemma_array_elements(body: str) -> str:
out.append(element)
continue
if stripped[0] == "{":
# Object element: quote its keys/bare values like a top-level object.
out.append(_quote_gemma_object_keys(stripped))
continue
if stripped[0] == "[":
# Nested array: normalise its elements too.
inner_end = _balanced_bracket_end(stripped, 0)
if inner_end == len(stripped) - 1:
out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]")
@ -240,8 +465,6 @@ def _quote_gemma_object_keys(src: str) -> str:
while i < len(src) and src[i].isspace():
i += 1
key_name_start = i
# Dots match the parser's key/name charset: Gemma emits dotted argument keys
# (user.name:...) for namespaced schemas.
while i < len(src) and (src[i].isalnum() or src[i] in "_-."):
i += 1
key_name = src[key_name_start:i]
@ -254,12 +477,15 @@ def _quote_gemma_object_keys(src: str) -> str:
parts.append(src[i:colon_pos])
parts.append(":")
i = colon_pos + 1
# Quote bare string values ({unit:celsius}); JSON stays as-is.
# Gemma may emit bare string values ({unit:celsius}); quote them so
# json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is.
ws = i
while i < len(src) and src[i].isspace():
i += 1
parts.append(src[ws:i])
if i < len(src) and src[i] == "[":
# Array value: quote bare string elements (e.g. labels:[bug,ui])
# so json.loads succeeds instead of dropping the call.
arr_end = _balanced_bracket_end(src, i)
if arr_end < 0:
parts.append(src[i:])
@ -269,7 +495,9 @@ def _quote_gemma_object_keys(src: str) -> str:
i = arr_end + 1
elif i < len(src) and src[i] not in '"{':
v_start = i
# Bare value: up to `}` or a comma that starts the next key:pair.
# Consume the bare value up to `}` or a comma that starts the
# next key:value pair; a comma inside the value (e.g.
# `New York, NY`) does not terminate it.
while i < len(src):
if src[i] == "}":
break
@ -329,7 +557,9 @@ def _func_close_index(content: str, body_start: int, body: str) -> int:
def _trim_param_value(val: str) -> str:
"""Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives."""
"""Trim the single wrapping newline the chat template adds around an XML
parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed
code/diff argument indentation)."""
if val.startswith("\n"):
val = val[1:]
if val.endswith("\n"):
@ -404,6 +634,7 @@ def parse_tool_calls_from_text(
*,
id_offset: int = 0,
allow_incomplete: bool = True,
enabled_tool_names = None,
with_spans: bool = False,
):
"""Parse OpenAI-format tool calls from model text.
@ -412,22 +643,36 @@ def parse_tool_calls_from_text(
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<|tool_call>call:web_search{query:"..."}<tool_call|>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
[TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback)
web_search[ARGS]{"query":"..."} (reasoning-model rehearsal)
A call rehearsed inside a ``<think>`` / ``[THINK]`` block is skipped, not
executed; the block is kept so a literal tag in a real argument is preserved.
With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]``
is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup
in ``content`` (including its close tag when present), so a caller can
remove exactly the parsed markup and keep every other byte intact.
"""
# Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a
# think marker opening inside a call is argument data (excluded from spans).
_think_spans = _think_spans_outside_tool_markup(content)
_think_starts = [s for s, _e in _think_spans]
def _in_think(pos: int) -> bool:
# Spans are ordered and non-overlapping; bisect gives O(log M) per candidate.
i = bisect.bisect_right(_think_starts, pos) - 1
return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1]
tool_calls: list[dict] = []
call_spans: list[tuple] = []
# Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside
# another call's coverage, or an open <parameter=> value, is data not executed.
markers = _build_markers(content)
coverage = _marker_coverage(content, markers)
# Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside
# another call's coverage (even one that failed to parse) is data, not executed. A
# marker opening inside a think block is a rehearsal and is skipped.
parsed_items = [] # (start, span_end, name, arguments) in document order
markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])]
coverage = _marker_coverage(content, markers)
for idx, (start, brace_end, kind, m) in enumerate(markers):
# A marker starting inside another's coverage is that call's data. The
# end is exclusive so a marker at a close's end is an adjacent sibling.
if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx):
continue
if brace_end < 0:
@ -441,7 +686,7 @@ def parse_tool_calls_from_text(
if kind == "json":
obj = json.loads(content[m.end() - 1 : brace_end + 1])
name = obj.get("name", "")
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes <tool_call>).
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes).
arguments = obj.get("arguments")
if arguments is None:
arguments = obj.get("parameters", {})
@ -452,7 +697,6 @@ def parse_tool_calls_from_text(
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))
except (json.JSONDecodeError, ValueError):
continue
# Span reaches through the close tag when present, else just the braces.
span_end = brace_end + 1
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
ws = len(content[span_end:]) - len(content[span_end:].lstrip())
@ -461,14 +705,11 @@ def parse_tool_calls_from_text(
span_end = close_m.end()
parsed_items.append((start, span_end, name, arguments))
# Function-XML calls promote in document order alongside marker calls (the
# #6801 contract). A <function=> inside any marker's coverage is excluded --
# even if that marker failed to parse -- so nested XML cannot escape; one
# after a balanced close-less marker is a sibling, not swallowed to EOF.
func_starts = [
fm
for fm in _TC_FUNC_START_RE.finditer(content)
if not _inside_open_parameter(content, fm.start())
and not _in_think(fm.start())
and not any(s <= fm.start() < e for s, e in coverage)
]
for idx, fm in enumerate(func_starts):
@ -545,11 +786,170 @@ def parse_tool_calls_from_text(
)
call_spans.append((start, span_end))
# Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document
# order, so a Mistral call and a rehearsal in one message both parse.
if not tool_calls:
for start, end, kind, m in _iter_bracket_spans(
content, enabled_tool_names = enabled_tool_names
):
if _in_think(start):
continue
# Extend the region over an immediately-following v11 closer so with_spans consumers strip it too.
closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:])
region_end = end + closer.end() if closer else end
if kind == "array":
# Decode elements individually (comma-tolerant): one json.loads of the whole
# body rejects the comma-less multi-call arrays Mistral/Ollama templates emit.
payload, item_ends = _decode_array_items(content, m.end(), end)
if not payload:
continue
# Tile the region so every byte belongs to exactly one span; a with_spans consumer
# keeps skipped bytes visible and strips promoted markup exactly once.
tile_start = start
last_span_idx = -1
for item_idx, item in enumerate(payload):
if not isinstance(item, dict) or "name" not in item:
continue
args = item.get("arguments", {})
if isinstance(args, str):
# ``arguments`` may itself be a JSON string (OpenAI spec).
try:
args = json.loads(args)
except (json.JSONDecodeError, ValueError):
pass
if not isinstance(args, (dict, str)):
# ``"arguments": null`` (or any non-object scalar) becomes {} like the
# <tool_call> path, not the string "null" auto-heal would mangle to
# a bogus {"query":"null"}.
args = {}
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": item.get("name", ""),
# A bare scalar string stays raw (like the <tool_call> path);
# json.dumps would double-encode it so the arg healer wraps
# "weather" with its literal quotes.
"arguments": args if isinstance(args, str) else json.dumps(args),
},
}
)
item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end
last_span_idx = len(call_spans)
call_spans.append((tile_start, item_end))
tile_start = item_end
if last_span_idx >= 0:
tile_start, _tile_end = call_spans[last_span_idx]
call_spans[last_span_idx] = (tile_start, region_end)
else:
try:
payload = json.loads(content[m.end() : end])
except (json.JSONDecodeError, ValueError):
continue
if not isinstance(payload, dict):
continue
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": m.group(1),
"arguments": json.dumps(payload),
},
}
)
call_spans.append((start, region_end))
if with_spans:
return tool_calls, call_spans
return tool_calls
def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str:
"""Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one
balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex
left two-level args behind). Truncated tails go to the caller's catch-all. Linear.
``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None
strips every span)."""
if len(text) > _MAX_BRACKET_SCAN_CHARS:
return text
out: list[str] = []
cursor = 0
for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names):
out.append(text[cursor:start])
cursor = end
out.append(text[cursor:])
return "".join(out)
def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]:
"""Spans of tool-call markup, so a literal <think>/[THINK] inside a call's args is
stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket
calls and an unclosed XML call (run via allow_incomplete); without the open-ended
span the unclosed call's markup would leak after execution."""
# Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan
# to EOF from every opener (quadratic on a stream of unclosed openers).
spans = [
m.span()
for pat in _TOOL_CLOSED_PATS
if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text)
for m in pat.finditer(text)
]
spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text))
# An unclosed opener is a real incomplete call only outside closed/bracket spans.
for pat in _TOOL_OPEN_XML_TAIL_PATS:
for m in pat.finditer(text):
if not any(s <= m.start() < e for s, e in spans):
spans.append(m.span())
return spans
def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]:
"""<think>/[THINK] block spans, minus any whose opening marker sits INSIDE a
tool-call span (that tag is argument data, not reasoning). Keeping it would drop a
real call after it as rehearsed and leak the call's markup. START tested only, so
a greedy unclosed <think> past the call is still that call's argument data."""
think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)]
call_spans = _tool_call_markup_spans(text)
# Prefilled reasoning: the template opens <think> in the prompt, so add a leading span
# (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe.
close = _THINK_CLOSE_RE.search(text)
if close is not None:
opener = _THINK_OPEN_RE.search(text)
if (
(opener is None or close.start() < opener.start())
and not any(cs <= close.start() < ce for cs, ce in call_spans)
and any(cs >= close.end() for cs, ce in call_spans)
):
think_spans = [(0, close.end())] + think_spans
if not think_spans:
return think_spans
if not call_spans:
return think_spans
return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)]
def strip_outside_think(text: str, strip_segment) -> str:
"""Apply ``strip_segment(segment, is_last)`` to visible text around <think>/[THINK]
blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal).
``is_last`` is True only after the final block, so trailing-tail patterns apply
only there. Shared by every strip path so they stay consistent."""
# A think marker opening inside a complete call is argument text; excluding it lets the
# stripper see the whole call. START-tested, so an unclosed match stays argument data.
think_spans = _think_spans_outside_tool_markup(text)
if not think_spans:
return strip_segment(text, True)
pieces: list[str] = []
prev = 0
for s, e in think_spans:
pieces.append(strip_segment(text[prev:s], False))
pieces.append(text[s:e])
prev = e
pieces.append(strip_segment(text[prev:], True))
return "".join(pieces)
def _strip_gemma_native_spans(text: str, *, final: bool) -> str:
"""Remove complete Gemma-native spans, brace/quote-balanced so a literal
``<tool_call|>`` in a quoted argument cannot truncate the span. An incomplete
@ -635,26 +1035,43 @@ def _strip_closed_blocks_outside_gemma(text: str) -> str:
return text
def strip_tool_markup_final(text: str) -> str:
"""Final display strip, shared with the streaming wrappers so all paths order
the passes identically: Gemma-aware closed JSON/function blocks first, then
well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed
spans and drop any unclosed remainder to EOF. Whitespace is kept."""
def _strip_markup_segment(
text: str,
*,
final: bool,
enabled_tool_names = None,
) -> str:
# Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth,
# rehearsal name-gated); then the quote-aware Gemma-native passes so a literal
# <tool_call|> in an argument cannot truncate a block; finally the regex XML/tail sweeps.
text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names)
text = _strip_closed_blocks_outside_gemma(text)
text = _strip_gemma_native_spans(text, final = True)
return strip_tool_patterns(text, _TOOL_ALL_PATS)
text = _strip_gemma_native_spans(text, final = final)
patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS
return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names)
def strip_tool_call_markup(text: str, *, final: bool = False) -> str:
def strip_tool_call_markup(
text: str,
*,
final: bool = False,
enabled_tool_names = None,
) -> str:
"""Strip tool-call XML markup from text.
When ``final`` is False, only fully closed tool-call blocks are removed.
When ``final`` is True, trailing incomplete tool-call blocks are removed
too, and the result is stripped of surrounding whitespace.
``<think>`` / ``[THINK]`` reasoning is preserved verbatim (see
``strip_outside_think``); the trailing-tail patterns apply only after the
last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}``
example visible (it is prose, not a call) so display cleanup matches detection.
"""
if final:
return strip_tool_markup_final(text).strip()
# Non-final: same ordering as the final path, but incomplete blocks are kept.
text = _strip_closed_blocks_outside_gemma(text)
text = _strip_gemma_native_spans(text, final = False)
return strip_tool_patterns(text, _TOOL_CLOSED_PATS)
result = strip_outside_think(
text,
lambda seg, is_last: _strip_markup_segment(
seg, final = final and is_last, enabled_tool_names = enabled_tool_names
),
)
return result.strip() if final else result

277
studio/backend/routes/inference.py
View file

@ -1366,17 +1366,60 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict:
return flags
def _generation_prompt_opens_think(template: Optional[str]) -> bool:
"""True when rendering the template's generation prompt ends INSIDE an unclosed ``<think>``.
Distinguishes templates that PREFILL an open ``<think>`` in the assistant generation
prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing
``</think>`` and the extractor must start in reasoning mode -- from templates that merely
render PAST assistant ``<think>...</think>`` history while leaving the generation prompt
open with no ``<think>`` (e.g. Kimi-K2-Thinking), where the model self-emits its own block
and the extractor must start in normal mode. Renders a single-user-message probe with the
same sandbox transformers uses; on any failure returns True, preserving the historical
always-on prefill for templates that cannot be rendered here.
"""
if not template:
return False
try:
from jinja2.sandbox import ImmutableSandboxedEnvironment
def _raise_exception(message: str):
raise RuntimeError(message)
env = ImmutableSandboxedEnvironment(
trim_blocks = True,
lstrip_blocks = True,
extensions = ["jinja2.ext.loopcontrols"],
)
env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False)
env.globals["raise_exception"] = _raise_exception
rendered = env.from_string(template).render(
messages = [{"role": "user", "content": "hi"}],
add_generation_prompt = True,
bos_token = "",
eos_token = "",
)
except Exception:
return True
# ``<think>`` is not a substring of ``</think>`` (the ``/`` breaks it), so the last open
# tag sitting after the last close tag means the prompt ends inside an open block.
return rendered.rfind("<think>") > rendered.rfind("</think>")
def _sf_reasoning_prefill_mode(
features: dict,
enable_thinking: Optional[bool],
template: Optional[str] = None,
reasoning_effort: Optional[str] = None,
) -> bool:
"""Whether this request begins INSIDE an unclosed ``<think>`` (Qwen3/Qwen3.5/GLM prefill it).
"""Whether a safetensors/MLX generation begins INSIDE an unclosed ``<think>``.
Gated on the STANDARD ``<think>``/``</think>`` markers: a bespoke reasoning channel (e.g. gemma)
never emits ``</think>``, so prefilled mode would swallow the whole answer -- excluded, as are
gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill.
``enable_thinking`` templates (Qwen3/GLM) prefill an open ``<think>`` so the model
emits only the closing ``</think>``, and the extractor must start in reasoning mode.
Gated on the STANDARD ``<think>``/``</think>`` markers: bespoke channels (gemma's
``<|think|>``) never emit ``</think>`` and would swallow the answer, so they and
gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None
defaults thinking ON, so a plain request still prefills.
"""
if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"):
return False
@ -1384,16 +1427,21 @@ def _sf_reasoning_prefill_mode(
if "</think>" not in tpl and "<think>" not in tpl:
return False
if features.get("reasoning_always_on"):
return True
# enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps
# thinking on, so always-on wins over reasoning_effort and we prefill.
if features.get("reasoning_style") == "enable_thinking_effort":
return True
# ``reasoning_always_on`` fires on paired ``<think>...</think>`` anywhere in the
# template, including markup that only renders PAST assistant history (Kimi-K2-Thinking)
# while the generation prompt opens none. Prefill only when the generation prompt opens
# one, else the extractor captures a normal answer as reasoning_content and returns blank.
return _generation_prompt_opens_think(tpl)
if not features.get("supports_reasoning"):
return False
if enable_thinking is False:
return False
# A reasoning_effort="none" request disables thinking for enable_thinking_effort
# (GLM-5.2) models the same way enable_thinking=False does (see
# ``_request_reasoning_kwargs``). Without this, the model emits no ``</think>`` and
# a plain answer is swallowed whole into reasoning_content, leaving the visible
# response empty.
# Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it
# so we don't prefill and capture the answer. Plain enable_thinking models ignore effort.
if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none":
return False
return True
@ -1669,11 +1717,17 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str:
return nudge + " " + _RAG_GROUNDING_NUDGE
# Strip leaked tool-call markup: every shared-parser format plus the four leak
# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN
# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a
# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation
# is the parser's own, so a signal we parse is never left un-stripped.
# Strip leaked tool-call markup: every shared-parser format plus the leak shapes
# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary:
# 1. well-formed `<tool_call>...</tool_call>` / `<function=...>...</function>`
# 2. orphan opening to EOF (close was DRAINED)
# 3. bare orphan close (open was DRAINED)
# 4. tail-only `</parameter>` (outer close truncated by EOS); anchored to
# `\Z` so mid-text `<parameter>` in user code samples survives.
# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced
# scan removes the whole call (a non-greedy regex would truncate nested JSON).
# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal
# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own.
from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC
_TOOL_XML_RE = _re.compile(
@ -1694,6 +1748,17 @@ _TOOL_XML_RE = _re.compile(
r"|</(?:tool_call|function)>"
r"|<tool_call\|>"
r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*"
r"|\[/TOOL_CALLS\]"
# Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove
# it, so strip its tail here.
r"|\[TOOL_CALLS\]\s*\[.*\Z"
# Named / v11 forms and bare rehearsal; arms aligned with the parser regexes.
r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)"
# Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]``
# survives); NAME captured as ``reh`` for the inactive-name display gate.
r"|(?<!\[CALL_ID\])\b(?P<reh>[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)"
# DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls;
# each arm carries a call-shaped lookahead so prose merely mentioning a marker survives.
r"|"
+ _DS_OPEN_SRC
+ r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)"
@ -1705,6 +1770,17 @@ _TOOL_XML_RE = _re.compile(
_re.DOTALL,
)
# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms
# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``.
_TOOL_XML_CLOSED_RE = _re.compile(
r"<(?:tool_call|function=[\w-]+)>.*?</(?:tool_call|function)>"
r"|<\|tool_call>.*?<tool_call\|>"
r"|</(?:tool_call|function)>"
r"|<tool_call\|>"
r"|\[/TOOL_CALLS\]",
_re.DOTALL,
)
def _gemma_strip_gate(tools) -> set:
"""Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the
@ -1720,18 +1796,18 @@ def _gemma_strip_gate(tools) -> set:
return names
def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str:
"""Combine the parser's scan-based strips (Mistral balanced-brace, gated
Gemma wrapper-less, GLM real-close, guarded function-XML) with
``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so
literal markup inside argument values is data, not a leaked tail.
``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call."""
cleaned = _strip_glm_calls(
_strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names),
final = True,
)
cleaned = _strip_function_xml_calls(cleaned, final = True)
return _TOOL_XML_RE.sub("", cleaned)
def _display_tool_name_gate(active_tools):
"""Active tool NAMES for gating the rehearsal display strip, or None when no tools
are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate:
a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every
identifier stays ambiguous, so strip."""
names = {
(t.get("function") or {}).get("name")
for t in (active_tools or [])
if isinstance(t, dict) and isinstance(t.get("function"), dict)
}
names.discard(None)
return names or None
def _strip_tool_xml_for_display(
@ -1740,12 +1816,56 @@ def _strip_tool_xml_for_display(
auto_heal_tool_calls: bool,
enabled_tool_names: Optional[set] = None,
) -> str:
"""Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml``
so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no
``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip."""
"""Apply route-level XML leak cleanup only when Auto-Heal is enabled.
Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma
wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator
so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the
DeepSeek / Kimi / orphan forms. ``<think>`` blocks are preserved verbatim and the
``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a
block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal
``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool
list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip
unconditionally regardless of NAME."""
if not auto_heal_tool_calls:
return text
return _strip_tool_xml(text, enabled_tool_names)
from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think
def _keep_inactive_rehearsal(m) -> str:
# Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive
# NAME[ARGS]{...} is prose -- keep it.
if enabled_tool_names is not None:
name = m.groupdict().get("reh")
if name is not None and name not in enabled_tool_names:
return m.group(0)
return ""
def _strip_segment(seg: str, is_last: bool) -> str:
# Scan strips close at each call's REAL terminator (a literal ``</function>`` or a
# nested marker quoted inside a value cannot truncate the strip); the regex arms below
# cover the attribute form and the DeepSeek / Kimi / orphan families.
seg = _strip_mistral_closed_calls(seg)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
seg = _strip_glm_calls(seg, final = is_last)
seg = _strip_function_xml_calls(seg, final = is_last)
if is_last:
return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg)
return _TOOL_XML_CLOSED_RE.sub("", seg)
return strip_outside_think(text, _strip_segment)
def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str:
# Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then
# the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside
# _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names``
# gates the Gemma wrapper-less strip; ``None`` strips every closed call.
text = _strip_mistral_closed_calls(text)
return _strip_tool_xml_for_display(
text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names
)
logger = get_logger(__name__)
@ -6010,14 +6130,18 @@ async def openai_chat_completions(
_gguf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# Active tool names gating the bare-rehearsal strip, matching the loop gate.
_gguf_display_tool_names = _display_tool_name_gate(tools_to_use)
# ── Strip stale tool-call XML from conversation history ─
for _msg in gguf_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
# Gate on enabled tool names, like the live strip, so a documented inactive
# ``foo[ARGS]{...}`` survives in the replayed prompt context.
_msg["content"] = _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(tools_to_use),
enabled_tool_names = _gguf_display_tool_names,
).strip()
def gguf_generate_with_tools():
@ -6151,7 +6275,7 @@ async def openai_chat_completions(
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(tools_to_use),
enabled_tool_names = _gguf_display_tool_names,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
@ -6258,7 +6382,7 @@ async def openai_chat_completions(
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(tools_to_use),
enabled_tool_names = _gguf_display_tool_names,
)
return full_text, usage, finish
finally:
@ -6631,14 +6755,17 @@ async def openai_chat_completions(
_sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template")
_sf_features = _detect_safetensors_features(backend, _sf_tpl)
# Split prefilled-``<think>`` output into reasoning_content deltas (GGUF parity) so the UI
# renders the thinking block for safetensors AND MLX.
# GGUF parity: enable_thinking templates prefill an unclosed <think>; split into
# reasoning_content deltas so the UI renders the block for safetensors and MLX.
_sf_parse_think = bool(
_sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on")
)
# Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded.
# Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode.
_sf_reasoning_prefilled = _sf_reasoning_prefill_mode(
_sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort
_sf_features,
payload.enable_thinking,
_sf_tpl,
reasoning_effort = payload.reasoning_effort,
)
def _new_sf_reasoning_extractor():
@ -6722,6 +6849,8 @@ async def openai_chat_completions(
_sf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# Active tool names gating the bare-rehearsal strip, matching the loop gate.
_sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use)
# Strip stale tool-call XML from prior assistant turns.
_sf_chat_messages = []
@ -6733,7 +6862,7 @@ async def openai_chat_completions(
"content": _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
enabled_tool_names = _sf_display_tool_names,
).strip(),
}
)
@ -6792,7 +6921,7 @@ async def openai_chat_completions(
reasoning_extractor = _new_sf_reasoning_extractor()
def _sf_flush_reasoning():
# Drain the extractor at a turn boundary / stream end (GGUF parity); only visible text reaches the monitor.
# Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor.
fr, fv = reasoning_extractor.finish()
out = []
if fr:
@ -6818,7 +6947,7 @@ async def openai_chat_completions(
if event["type"] == "status":
if not event["text"]:
# Iteration boundary: flush reasoning, then start a fresh extractor for the next turn.
# Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn.
for _c in _sf_flush_reasoning():
yield _c
prev_text = ""
@ -6834,7 +6963,7 @@ async def openai_chat_completions(
if event["type"] in ("tool_start", "tool_end"):
if event["type"] == "tool_start":
# Flush reasoning before the tool_start line so the thinking block closes ahead of the tool card.
# Flush reasoning before tool_start so the thinking block closes ahead of the card.
for _c in _sf_flush_reasoning():
yield _c
prev_text = ""
@ -6847,7 +6976,7 @@ async def openai_chat_completions(
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
enabled_tool_names = _sf_display_tool_names,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
@ -6939,12 +7068,12 @@ async def openai_chat_completions(
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use),
enabled_tool_names = _sf_display_tool_names,
)
return full_text
content_text = await asyncio.to_thread(_drain_to_text)
# Split prefilled <think> reasoning out of the visible answer (GGUF parity); monitor gets visible text only.
# Split prefilled <think> out of the visible answer (GGUF parity); the monitor gets visible text only.
_reasoning_text, _visible_text = _extract_responses_reasoning(
content_text,
parse_think_markers = _sf_parse_think,
@ -7043,7 +7172,7 @@ async def openai_chat_completions(
yield _chat_role_chunk(completion_id, created, model_name)
prev_text = ""
# Split prefilled <think> into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves MLX.
# Split prefilled <think> into reasoning_content deltas (GGUF parity); single turn, serves MLX.
reasoning_extractor = _new_sf_reasoning_extractor()
# Run the sync generator in a thread pool to avoid blocking the
# event loop. Critical for compare mode: two SSE requests arrive
@ -7150,7 +7279,7 @@ async def openai_chat_completions(
for token in generate():
full_text = token
# Split prefilled <think> reasoning from the visible answer (GGUF parity); also covers MLX.
# Split prefilled <think> reasoning (GGUF parity); also covers MLX via the shared generate().
_reasoning_text, _visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _sf_parse_think,
@ -8000,8 +8129,8 @@ class _ResponsesReasoningExtractor:
reasoning_prefilled: bool = False,
) -> None:
self._buffer = ""
# ``reasoning_prefilled``: output begins INSIDE an unclosed ``<think>`` (Qwen3/GLM prefill),
# so start in reasoning to capture leading text until the first ``</think>``. Callers default False.
# reasoning_prefilled: the template inserts an unclosed <think>, so output begins inside
# the block; start in reasoning until the first close tag. Existing callers pass False.
self._in_reasoning = reasoning_prefilled
# Splitting requires marker parsing; a prefilled open implies it.
self._parse_think_markers = parse_think_markers or reasoning_prefilled
@ -8033,8 +8162,8 @@ class _ResponsesReasoningExtractor:
self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :]
self._in_reasoning = False
continue
# Hold back a trailing partial of EITHER marker: the close (clean chunk-boundary split)
# and a stray open (so a re-emitted ``<think>`` isn't leaked into the reasoning drawer).
# Hold back a trailing partial of either marker: the close (clean split across chunks)
# and a stray open (a re-emitted <think> is suppressed, not leaked).
keep = _responses_marker_holdback(
self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN)
)
@ -9919,11 +10048,15 @@ async def anthropic_messages(
else:
openai_messages.insert(0, {"role": "system", "content": _nudge})
# Strip stale tool-call XML from conversation
# Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS]
# prose survive), gated on enabled tool names so documented inactive examples are kept.
_anthropic_history_gate = _display_tool_name_gate(openai_tools)
for _msg in openai_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_msg["content"] = _strip_tool_xml(
_msg["content"], _gemma_strip_gate(openai_tools)
_msg["content"] = _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = True,
enabled_tool_names = _anthropic_history_gate,
).strip()
def _run_tool_gen():
@ -10023,6 +10156,10 @@ async def _anthropic_tool_stream(
"""Streaming response for the tool-calling path."""
_sentinel = object()
# Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final
# answer is prose and must survive in the delivered text.
_display_names = _display_tool_name_gate(openai_tools)
# Prompt-token count for message_start.usage.input_tokens. count_chat_tokens
# makes blocking HTTP calls to llama-server, so run it off the event loop.
# Pass the tools so tool-schema tokens are counted (the generator renders
@ -10074,9 +10211,15 @@ async def _anthropic_tool_stream(
captured_finish_reason = _fr
# Strip leaked tool-call XML from content events first, so a
# content event that was purely tool XML doesn't count as text.
# Protected helper preserves <think> rehearsal and balanced
# [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both).
if etype == "content":
event = dict(event)
event["text"] = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools))
event["text"] = _strip_tool_xml_for_display(
event["text"],
auto_heal_tool_calls = True,
enabled_tool_names = _display_names,
)
# disable_parallel_tool_use: keep only the first tool_use block,
# dropping every later tool_start and its paired tool_end (robust
# to empty tool-call ids — tracked by state, not id matching).
@ -10250,6 +10393,9 @@ async def _anthropic_tool_non_streaming(
usage = {}
prev_text = ""
captured_finish_reason = None
# Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final
# answer is prose and must survive in the delivered text.
_display_names = _display_tool_name_gate(openai_tools)
# Pending client tool_use; cleared by tool_end (server execution) or
# trailing text. See the stop_reason mapping below.
ends_on_tool_use = False
@ -10259,8 +10405,10 @@ async def _anthropic_tool_non_streaming(
for event in events:
etype = event.get("type", "")
if etype == "content":
# Strip leaked tool-call XML
clean = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools))
# Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose).
clean = _strip_tool_xml_for_display(
event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names
)
new = clean[len(prev_text) :]
prev_text = clean
if new:
@ -10730,13 +10878,16 @@ async def _anthropic_passthrough_non_streaming(
text = message.get("content") or ""
if text:
# Keep unpromoted bytes when healing is active; legacy stripping is
# only for opted-out or no-client-tool requests. Use the full
# _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded
# function-XML leaks are cleaned too, not just _TOOL_XML_RE forms,
# with the Gemma display gate so a disabled/example call:NAME{...}
# in prose survives.
# only for opted-out or no-client-tool requests. Protected helper (not
# raw _TOOL_XML_RE.sub): preserves <think> rehearsal and balanced
# [TOOL_CALLS] trailing prose, gated on the declared tools so an
# inactive NAME[ARGS]{...} example in the final text is kept.
if not healing_active:
text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools))
text = _strip_tool_xml_for_display(
text,
auto_heal_tool_calls = True,
enabled_tool_names = _display_tool_name_gate(openai_tools),
)
text = text.strip()
if text:
content_blocks.append(AnthropicResponseTextBlock(text = text))

18
studio/backend/tests/test_anthropic_messages.py
View file

@ -889,6 +889,24 @@ class TestAnthropicToolNonStreaming:
assert tool_blocks[0]["name"] == "render_html"
assert tool_blocks[0]["input"] == {"code": "<!doctype html><html></html>"}
def test_display_strip_gates_on_declared_tools(self):
# A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared
# ``foo`` markup is prose and survives, the declared web_search rehearsal strips.
def _run_gen():
yield {
"type": "content",
"text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.',
}
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
response = asyncio.run(
_anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools)
)
body = json.loads(response.body)
text = "".join(b["text"] for b in body["content"] if b["type"] == "text")
assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose
assert "web_search[ARGS]" not in text # active name stripped from display
# =====================================================================
# Pass-through emitter tests (client-side tool execution path)

14
studio/backend/tests/test_gemma_tool_parse_edge_cases.py
View file

@ -160,6 +160,20 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call():
assert [c["function"]["name"] for c in calls] == ["python"], calls
def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call():
# A literal <think> inside a completed call's arguments is argument data; both calls must parse.
text = '[TOOL_CALLS]a{"x":"literal <think> marker"} b[ARGS]{"y":2}'
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["a", "b"], calls
def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal():
# A genuine reasoning block still hides its rehearsal while a real call after it parses.
text = '<think>web_search[ARGS]{"q":"draft"}</think>real[ARGS]{"q":"go"}'
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["real"], calls
def test_wrapperless_nested_object_argument_is_parsed():
# skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare.
calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}")

306
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -2130,6 +2130,312 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch):
assert "prompt_tokens_details" not in metadata[-1]["usage"]
def test_gguf_rehearsal_name_split_before_args_is_not_leaked(monkeypatch):
"""Finding 6: a rehearsal call whose name (``web_search``) and ``[ARGS]{...}``
arrive in separate content deltas must hold the bare name in the buffer until
``[ARGS]`` flips it to a drain. Without _is_rehearsal_prefix the GGUF path
streams the tool name as visible content before the call executes."""
first_stream = [
_sse({"content": "web_search"}),
_sse({"content": '[ARGS]{"query":"cats"}'}),
_done(),
]
final_stream = [_sse({"content": "Found cats."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search cats"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [("web_search", {"query": "cats"})], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert all("web_search" not in t for t in content_texts), content_texts
assert all("[ARGS]" not in t for t in content_texts), content_texts
def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch):
"""The first flush out of BUFFERING (prose plus a trailing active-tool-name in
the first delta, ``[ARGS]{...}`` in the next) must apply the same trailing-name
hold the STREAMING branch uses. The first delta has spaces so it is not a
rehearsal prefix and falls to the initial flush, which previously emitted the
bare name before the call drained."""
first_stream = [
_sse({"content": "I will use web_search"}),
_sse({"content": '[ARGS]{"query":"cats"}'}),
_done(),
]
final_stream = [_sse({"content": "Found cats."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search cats"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [("web_search", {"query": "cats"})], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert all("web_search" not in t for t in content_texts), content_texts
assert all("[ARGS]" not in t for t in content_texts), content_texts
def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch):
"""Finding 9: the BUFFERING guard only covers a rehearsal at the turn start.
When prose has already streamed (STREAMING state) and the model then emits the
tool name and ``[ARGS]{...}`` in later deltas, the bare name must still be held,
not flushed as visible content before the call drains."""
first_stream = [
_sse({"content": "Let me think. "}),
_sse({"content": "I will search "}),
_sse({"content": "web_search"}),
_sse({"content": '[ARGS]{"query":"cats"}'}),
_done(),
]
final_stream = [_sse({"content": "Found cats."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search cats"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [("web_search", {"query": "cats"})], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert all("web_search" not in t for t in content_texts), content_texts
def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch):
"""End-of-stream flush: a plain answer that ENDS on a tool-name word with no
``[ARGS]`` following is real prose and must not be dropped by the streaming
rehearsal hold."""
first_stream = [
_sse({"content": "I think "}),
_sse({"content": "you should "}),
_sse({"content": "web_search"}),
_done(),
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream], payloads)
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "advise"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert any(t.rstrip().endswith("web_search") for t in content_texts), content_texts
def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypatch):
"""Finding 11: a realistic MCP name longer than the 32-char buffer cap split as
NAME then [ARGS]{...} must still be held (a rehearsal prefix is self-bounding),
so the name does not leak and the call executes."""
name = "mcp__github__create_pull_request"
assert len(name) >= 32, len(name)
first_stream = [
_sse({"content": name}),
_sse({"content": '[ARGS]{"x":1}'}),
_done(),
]
final_stream = [_sse({"content": "done"}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda n, a, **_k: (calls.append((n, a)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "go"}],
tools = [{"type": "function", "function": {"name": name}}],
max_tool_iterations = 1,
)
)
assert calls == [(name, {"x": 1})], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert not any(name in t for t in content_texts), content_texts
def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch):
"""F4: the GGUF streaming strip must run its open-ended ``[ARGS]`` tail cleanup
only on the LAST segment. A bare ``foo[ARGS]`` (no JSON body, ``foo`` not a tool)
before a <think> block is prose, not a truncated call, so the final visible text
must keep it verbatim instead of dropping ``foo[ARGS]`` and corrupting the
sentence."""
first_stream = [
_sse({"content": "Please pass foo[ARGS] "}),
_sse({"content": "<think>pause</think> "}),
_sse({"content": "to the template."}),
_done(),
]
backend = _make_backend(monkeypatch, [first_stream], [])
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "x"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
assert content_texts, events
assert content_texts[-1] == "Please pass foo[ARGS] <think>pause</think> to the template."
def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch):
"""BUG A: an inactive-name ``foo[ARGS]{...}`` in a prose answer must not be treated
as a tool call. The BUFFERING and end-of-stream safety-net ``[ARGS]`` checks gate on
active tool names (like the safetensors loop and the mid-stream path), so ``foo``
(``web_search`` is the only enabled tool) is neither drained/parsed into a disabled
no-op nor forced into another generation turn."""
first_stream = [
_sse({"content": 'foo[ARGS]{"x":1} is just syntax.'}),
_done(),
]
backend = _make_backend(monkeypatch, [first_stream], [])
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "x"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 2,
)
)
# No tool executed for the inactive name; a spurious no-op re-prompt would exhaust the
# single supplied stream and error.
assert calls == [], calls
assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
# The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence.
assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts
def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch):
"""BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}``
in one delta must NOT swallow the real call; web_search executes while the inactive
rehearsal stays visible as prose."""
first_stream = [
_sse({"content": 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'}),
_done(),
]
final_stream = [_sse({"content": "Found cats."}), _done()]
backend = _make_backend(monkeypatch, [first_stream, final_stream], [])
calls: list[tuple[str, dict]] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"),
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search cats"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
# The real call runs; ``foo`` is not executed as a phantom disabled call.
assert calls == [("web_search", {"query": "cats"})], calls
content_texts = [e.get("text", "") for e in events if e.get("type") == "content"]
# The inactive rehearsal is preserved as prose; the active one is stripped.
assert any('foo[ARGS]{"a":1}' in t for t in content_texts), content_texts
assert all("web_search[ARGS]" not in t for t in content_texts), content_texts
def test_gguf_rehearsal_detection_recognises_spent_one_shot_with_original_tools():
# Rehearsal detection is fed the ORIGINAL tool list, so a spent one-shot's re-emitted
# repeat is still detected (matching the strip gate) instead of blanking the turn.
from core.inference.llama_cpp import _gguf_has_genuine_tool_signal
from core.inference.tool_call_parser import TOOL_XML_SIGNALS
repeat = 'render_html[ARGS]{"code":"<html>x</html>"}'
active_only = [{"type": "function", "function": {"name": "web_search"}}]
original = active_only + [{"type": "function", "function": {"name": "render_html"}}]
assert not _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, active_only)
assert _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, original)
def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot():
# The BUFFERING prefix check and STREAMING/flush tail-holds use the ORIGINAL tool list,
# so a spent one-shot's split repeat is held rather than leaked as visible text.
from core.inference.llama_cpp import _held_rehearsal_tail_len, _is_rehearsal_prefix
active_only = [{"type": "function", "function": {"name": "web_search"}}]
original = active_only + [{"type": "function", "function": {"name": "render_html"}}]
assert not _is_rehearsal_prefix("render_html", active_only)
assert _is_rehearsal_prefix("render_html", original)
assert _held_rehearsal_tail_len("answer render_html", active_only) == 0
assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html")
def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch):
"""An oversized bare-JSON call drains rather than streams, and still executes via the safety net."""

89
studio/backend/tests/test_passthrough_healing.py
View file

@ -280,6 +280,56 @@ class TestStreamHealer:
assert [c["id"] for c in calls] == ["call_0", "call_1"]
assert _events_text(events).strip() == "then"
def test_mistral_array_multiple_calls_all_promoted_in_stream(self):
# A canonical Mistral [TOOL_CALLS] array carries several calls under a
# SINGLE signal. Draining only the first call would leave the residue
# starting at ",{...}]" (no signal), so later calls in the same array
# must be promoted in the same pass, not flushed as raw text.
healer = StreamToolCallHealer({"get_weather", "get_time"})
array = (
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
)
events = healer.feed(array) + healer.finalize()
calls = _events_calls(events)
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
assert [c["id"] for c in calls] == ["call_0", "call_1"]
assert _events_text(events) == ""
def test_mistral_array_multiple_calls_promoted_char_by_char(self):
healer = StreamToolCallHealer({"get_weather", "get_time"})
array = (
'[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},'
'{"name":"get_time","arguments":{"tz":"UTC"}}]'
)
events = []
for ch in array:
events += healer.feed(ch)
events += healer.finalize()
calls = _events_calls(events)
assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"]
assert _events_text(events) == ""
def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self):
# A mid-array element for a tool that is not declared must survive as
# text while the declared neighbours on either side still promote in
# document order.
healer = StreamToolCallHealer({"a", "c"})
array = (
'[TOOL_CALLS][{"name":"a","arguments":{}},'
'{"name":"b","arguments":{}},{"name":"c","arguments":{}}]'
)
events = healer.feed(array) + healer.finalize()
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"]
assert '"b"' in _events_text(events)
def test_mistral_array_then_trailing_prose(self):
healer = StreamToolCallHealer({"a", "b"})
array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]'
events = healer.feed(f"{array} all done") + healer.finalize()
assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"]
assert "all done" in _events_text(events)
def test_incomplete_call_healed_at_finalize(self):
healer = StreamToolCallHealer({"Bash"})
events = healer.feed('<tool_call>{"name":"Bash","arguments":{"cmd":"ls"}}')
@ -1356,3 +1406,42 @@ class TestOpenaiStreamingRoute:
assert chunks[0] == line + "\n\n" # byte-for-byte relay
asyncio.run(_run())
class TestHealerSignalAlignment:
"""The passthrough healer buffers only formats its parser can promote.
The loops' bare [ARGS] rehearsal signal is gated on active tool names
there; ungated in the healer it would stall legitimate prose until
finalization without ever producing a promotable call."""
def test_heal_signals_are_promotable_formats_only(self):
from core.inference.passthrough_healing import _HEAL_SIGNALS
assert set(_HEAL_SIGNALS) == {
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
}
def test_prose_with_bare_args_marker_streams_through(self):
healer = StreamToolCallHealer({"Bash"})
chunks = [
"Use the pattern foo",
"[ARGS] in templates when calling tools, ",
"and remember to close it.",
]
streamed = ""
for chunk in chunks:
streamed += _events_text(healer.feed(chunk))
# Incremental relay: nothing withheld for finalize.
assert streamed == "".join(chunks)
final = healer.finalize()
assert not _events_calls(final)
assert not healer.healed
def test_bracket_tool_calls_still_promote_in_stream(self):
healer = StreamToolCallHealer({"web_search"})
events = healer.feed('[TOOL_CALLS]web_search{"query": "unsloth docs"}') + healer.finalize()
(call,) = _events_calls(events)
assert call["function"]["name"] == "web_search"
assert healer.healed

11
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -1990,8 +1990,8 @@ class TestTranslatedMessagesValidate:
ChatMessage(**m.model_dump(exclude_none = True))
# reasoning_prefilled mode: Qwen3/GLM enable_thinking templates prefill an unclosed <think>, so
# generation begins inside the think block and emits only the closing </think>; the extractor starts in reasoning.
# reasoning_prefilled: enable_thinking templates prefill an unclosed <think>, so
# generation begins inside the block; the extractor must start in reasoning.
class TestReasoningPrefilledExtractor:
def test_prefilled_single_feed_splits_lone_close(self):
# T1: reasoning...</think>answer with a prefilled (unseen) open tag.
@ -2077,9 +2077,7 @@ class TestReasoningPrefilledExtractor:
assert visible == "hi"
def test_not_prefilled_lone_close_preserves_current_behavior(self):
# T9: GGUF-parity guard -- WITHOUT prefilled, a lone </think> keeps the
# pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and
# every existing caller are byte-identical.
# T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard).
reasoning, visible = _extract_responses_reasoning(
"reasoning</think>ans",
parse_think_markers = True,
@ -2099,8 +2097,7 @@ class TestReasoningPrefilledExtractor:
assert visible == "v"
def test_prefilled_ignored_when_markers_not_parsed(self):
# T11: a non-reasoning model (parse_think_markers False) still passes text
# straight through even if reasoning_prefilled were mistakenly set False.
# T11: a non-reasoning model passes text through even with reasoning_prefilled False.
reasoning, visible = _extract_responses_reasoning(
"just an answer",
parse_think_markers = False,

37
studio/backend/tests/test_safetensors_capability_advertise.py
View file

@ -183,7 +183,9 @@ def test_detect_safetensors_features_llama3_template_keeps_tools_on():
def test_detect_safetensors_features_mistral_template_keeps_tools_on():
"""Mistral emits [TOOL_CALLS]; parser now supports it."""
"""Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses
(the shared bracket-tag parser). The gate must no longer suppress it, or the
PR's Mistral tool support is unreachable through normal capability detection."""
from routes.inference import _detect_safetensors_features
backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3")
@ -706,13 +708,28 @@ def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json():
assert flags["supports_tools"] is True
# _sf_reasoning_prefill_mode gates the prefilled-<think> extractor so safetensors/MLX reach
# GGUF reasoning-block parity for enable_thinking models.
# _sf_reasoning_prefill_mode gates the prefilled-<think> extractor (GGUF reasoning parity).
class TestSafetensorsReasoningPrefillGate:
# A minimal Qwen3-style template with the standard <think>/</think> markers.
_QWEN_TPL = "{% if enable_thinking %}<think>{% endif %}...</think>..."
# gemma-style bespoke reasoning channel -- no standard markers.
_GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought<channel|>"
# always-on template whose GENERATION PROMPT opens an unclosed <think> (DeepSeek-R1 / QwQ /
# Qwen3-Thinking shape): the model emits only the closing </think>, so prefill.
_ALWAYS_ON_OPEN_TPL = (
"{% for m in messages %}{{ m['content'] }}{% endfor %}"
"{% if add_generation_prompt %}<|assistant|><think>\n{% endif %}"
)
# always-on template that renders PAST assistant <think>...</think> history but leaves the
# generation prompt open with no <think> (Kimi-K2-Thinking shape): the model self-emits its
# own block, so prefill mode would blank a normal answer.
_ALWAYS_ON_HISTORY_TPL = (
"{% for m in messages %}"
"{% if m['role'] == 'assistant' %}<think>{{ m.get('reasoning_content', '') }}</think>"
"{{ m['content'] }}{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}"
)
def _features(self, **over):
base = {
@ -756,11 +773,19 @@ class TestSafetensorsReasoningPrefillGate:
feats = self._features(supports_reasoning = False, reasoning_style = None)
assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False
def test_g7_reasoning_always_on(self):
# G7: hardcoded-<think> template -> prefilled regardless of the flag.
def test_g7_reasoning_always_on_prompt_opens_think(self):
# G7: always-on template whose generation prompt opens <think> -> prefilled regardless of the flag.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_always_on = True)
assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True
assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True
def test_g7b_reasoning_always_on_history_only_not_prefilled(self):
# G7b (#5704): always-on classification from rendered assistant HISTORY <think></think>
# (Kimi-K2-Thinking) whose generation prompt opens no <think>. Prefill mode would capture a
# normal answer entirely as reasoning_content and blank the visible answer, so it must be off.
from routes.inference import _sf_reasoning_prefill_mode
feats = self._features(reasoning_always_on = True)
assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False
def test_g8_gemma_bespoke_channel_excluded(self):
# G8: gemma's <|think|>/<|channel> format has no </think> -> NOT prefilled

45
studio/backend/tests/test_safetensors_reasoning_stream.py
View file

@ -3,9 +3,11 @@
"""Safetensors/MLX reasoning-block parity with GGUF.
enable_thinking templates prefill an unclosed ``<think>``, so the stream must split the leading
text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy
of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events.
enable_thinking templates (Qwen3/GLM) prefill an unclosed ``<think>`` so the model
emits only the closing ``</think>`` then the answer; the safetensors stream must
split the leading text into ``reasoning_content`` deltas (plain stream and tool
loop), resetting per turn and appending only visible text to the monitor. Replays a
copy of ``sf_tool_stream``'s reasoning loop against synthetic events.
"""
from __future__ import annotations
@ -24,8 +26,40 @@ from routes.inference import (
)
_THINK_TPL = "...<think>...</think>..."
_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True}
_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True}
def test_prefill_mode_on_for_enable_thinking_default():
assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True
def test_prefill_mode_off_when_thinking_disabled():
assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False
def test_prefill_mode_off_for_reasoning_effort_none():
# enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode
# would capture the whole answer as reasoning_content.
assert (
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none")
is False
)
assert (
_sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high")
is True
)
def test_prefill_mode_off_without_think_markers():
assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False
def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict:
"""Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end."""
"""Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content``
snapshot, feed the delta through the extractor, and reset (flushing first) on
``tool_start`` / empty ``status`` so each turn splits independently."""
prev_text = ""
extractor = _ResponsesReasoningExtractor(
parse_think_markers = True, reasoning_prefilled = prefilled
@ -151,9 +185,6 @@ def test_s5_thinking_off_no_reasoning_deltas():
assert out["monitor"] == "Just the plain answer, no thinking."
_THINK_TPL = "...{% if enable_thinking %}<think>{% endif %}...</think>..."
def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort():
# GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and
# enable_thinking omitted) disables thinking exactly like enable_thinking=False, so

818
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -115,6 +115,22 @@ class TestParser:
assert result[0]["function"]["name"] == "python"
assert "print('hi')" in result[0]["function"]["arguments"]
def test_xml_param_preserves_leading_indentation(self):
import json
# Only the wrapping newline is trimmed; code-argument indentation survives.
text = (
"<function=python><parameter=code>\n"
" indented = 1\n"
" more\n"
"</parameter></function>"
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert json.loads(result[0]["function"]["arguments"]) == {
"code": " indented = 1\n more"
}
def test_xml_unclosed(self):
# Closing tags omitted; parser must still extract the value.
text = "<function=terminal><parameter=command>ls -la"
@ -183,6 +199,8 @@ class TestParser:
assert has_tool_signal("blah <tool_call> x")
assert has_tool_signal("blah <|tool_call>call:terminal")
assert has_tool_signal("hi <function=foo>...")
assert has_tool_signal("ok [TOOL_CALLS]web_search{...")
assert has_tool_signal("fine python[ARGS]{...")
assert not has_tool_signal("hello world")
def test_render_html_start_detector_uses_first_tool(self):
@ -197,6 +215,44 @@ class TestParser:
'<tool_call>{"name":"python","arguments":{"code":"<function=render_html>"}}'
)
def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self):
# The provisional render-html card must fire for bracket-tag forms too, not only XML.
assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":"<html>"}')
assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}')
assert _detect_render_html_tool_start(
'[TOOL_CALLS] [{"name":"render_html","arguments":{}}]'
)
assert _detect_render_html_tool_start('render_html[ARGS]{"code":"<html>"}')
# A different first tool (or a prose mention with no JSON body) must not fire.
assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}')
assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}')
assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}')
assert not _detect_render_html_tool_start("use render_html[ARGS] to render")
def test_render_html_start_detector_skips_think_block_rehearsal(self):
# A render_html rehearsed inside think must not fire the card; the outside-think call decides.
assert not _detect_render_html_tool_start(
'<think>draft render_html[ARGS]{"code":"x"}</think>python[ARGS]{"code":"print(1)"}'
)
assert not _detect_render_html_tool_start(
'[THINK]render_html[ARGS]{"code":"x"}[/THINK]web_search[ARGS]{"q":"y"}'
)
# A real render_html AFTER a rehearsed non-render_html inside think still fires.
assert _detect_render_html_tool_start(
'<think>web_search[ARGS]{"q":"x"}</think>render_html[ARGS]{"code":"<html>"}'
)
# A render_html rehearsed inside think with no real call after does not fire.
assert not _detect_render_html_tool_start('<think>render_html[ARGS]{"code":"x"}</think>')
def test_render_html_start_detector_reads_top_level_array_name(self):
# Array form: the name is the object's top-level ``"name"``, not an argument key.
assert not _detect_render_html_tool_start(
'[TOOL_CALLS] [{"arguments":{"name":"render_html"},"name":"python"}]'
)
assert _detect_render_html_tool_start(
'[TOOL_CALLS] [{"arguments":{"name":"python"},"name":"render_html"}]'
)
def test_strip_markup_closed(self):
text = "before <tool_call>{}</tool_call> after"
assert strip_tool_markup(text) == "before after"
@ -237,6 +293,376 @@ class TestParser:
== "before "
)
# Mistral [TOOL_CALLS] bracket-tag.
def test_mistral_bracket_basic(self):
# Devstral / Mistral-Small fallback when bypassing native FC.
text = '[TOOL_CALLS]web_search{"query":"weather"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
assert isinstance(result[0]["function"]["arguments"], str)
assert "weather" in result[0]["function"]["arguments"]
def test_rehearsal_inside_unclosed_think_is_ignored(self):
"""Rehearsal-shaped markup inside an unclosed <think> block must
not be executed as a real tool call. Mid-stream the </think>
tag has not arrived yet, so the strip regex has to accept
end-of-string as a terminator. Regression for the Gemini
high-severity flag on this PR."""
text = (
"<think>I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.'
)
result = parse_tool_calls_from_text(text)
# Inside an unclosed think block no calls are yielded.
assert result == []
def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self):
text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.'
result = parse_tool_calls_from_text(text)
assert result == []
def test_rehearsal_after_closed_think_still_parsed(self):
text = "<think>planning</think>" 'python[ARGS]{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_rehearsal_inside_prefilled_think_is_ignored(self):
"""Reasoning models (Qwen3.5 enable_thinking) open <think> in the PROMPT,
so generated content starts inside the thought and carries only a closing
</think>. A call rehearsed in that leading thought must be skipped, while a
real call after the close still fires."""
text = 'planning web_search[ARGS]{"query":"draft"}</think>python[ARGS]{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_literal_close_think_in_leading_argument_not_prefill(self):
"""A </think> literal inside a real leading call's arguments must not be
read as a prefilled-reasoning close (which would skip the call)."""
text = 'web_search[ARGS]{"query":"what is </think>"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_stray_close_after_real_call_not_treated_as_prefill(self):
"""A real leading call followed by a stray </think> and no further call is
a normal answer, not prefilled reasoning; the call must still fire (the
virtual span only applies when a real call follows the close)."""
text = 'Now web_search[ARGS]{"query":"x"}</think> answer'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_mistral_bracket_with_whitespace(self):
# Optional whitespace (incl. newlines) between the name and the opening brace.
text = '[TOOL_CALLS]python \n {"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
assert "print(1)" in result[0]["function"]["arguments"]
def test_mistral_bracket_nested_json(self):
# Brace-balance scan handles nested objects and braces inside string literals.
text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
import json as _json
args = _json.loads(result[0]["function"]["arguments"])
assert args["query"] == "a {nested} brace"
assert args["opts"] == {"limit": 5}
def test_mistral_bracket_with_prose(self):
# Bracket-tag surrounded by prose is still recognised.
text = (
"Sure, I will look that up.\n"
'[TOOL_CALLS]web_search{"query":"weather"}\n'
"Calling now."
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_mistral_bracket_bad_json_dropped(self):
text = "[TOOL_CALLS]web_search{not valid}"
result = parse_tool_calls_from_text(text)
# No usable tool call; callers fall back to text.
assert result == []
def test_mistral_bracket_object_with_array_value(self):
# Args must be a JSON object; a dict wrapping an array value is accepted.
text = '[TOOL_CALLS]web_search{"opts":[1,2,3]}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
# Rehearsal syntax name[ARGS]{json}.
def test_rehearsal_basic(self):
text = 'python[ARGS]{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
assert "print(1)" in result[0]["function"]["arguments"]
def test_rehearsal_with_prose(self):
text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_rehearsal_bad_json_dropped(self):
text = "python[ARGS]{not valid json}"
result = parse_tool_calls_from_text(text)
assert result == []
def test_mistral_bracket_hyphenated_mcp_name(self):
# Dashed MCP names must be captured whole, not truncated at the first dash.
text = '[TOOL_CALLS]mcp__srv__list-issues{"q":"x"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "mcp__srv__list-issues"
def test_rehearsal_hyphenated_mcp_name(self):
text = 'mcp__srv__list-issues[ARGS]{"q":"x"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "mcp__srv__list-issues"
def test_streaming_strip_removes_partial_bracket_marker(self):
# A bracket tag streamed before its opening brace must strip on the final pass, not leak.
assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer"
assert strip_tool_markup("text python[ARGS]", final = True) == "text"
# Non-final must keep the in-progress tag buffered (not yet stripped).
partial = "answer [TOOL_CALLS]web_search"
assert strip_tool_markup(partial, final = False) == partial
def test_strip_removes_two_level_nested_bracket_call_keeps_prose(self):
# Two-level-nested args must be removed whole; the balanced scan handles any depth.
text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after'
assert strip_tool_markup(text, final = False) == "before after"
assert strip_tool_markup(text, final = True) == "before after"
def test_strip_removes_call_with_literal_think_in_argument(self):
# A literal think block inside arguments strips with the call, not as a reasoning block.
text = (
'<tool_call>{"name":"write","arguments":'
'{"text":"compare <think> and </think> tags"}}</tool_call>'
)
assert strip_tool_markup(text, final = True) == ""
def test_strip_preserves_real_think_but_strips_call_with_literal_think(self):
text = (
"<think>planning</think> ok "
'<tool_call>{"name":"w","arguments":{"t":"<think>x</think>"}}</tool_call> done'
)
out = strip_tool_markup(text, final = True)
assert "<think>planning</think>" in out
assert "<tool_call>" not in out and '"name"' not in out
assert "ok" in out and "done" in out
def test_prose_mentioning_args_marker_is_not_truncated(self):
# ``foo[ARGS] to the template`` is prose; the catch-all must not delete the sentence.
text = "Please pass foo[ARGS] to the template and continue reading."
assert strip_tool_markup(text, final = True) == text
def test_streaming_strip_handles_mistral_v11_call_id_args(self):
# The streaming strip uses the regex patterns directly, so they must cover the v11
# [CALL_ID]/[ARGS] metadata (aligned with the parser).
raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after'
out = strip_tool_markup_streaming(raw)
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
# <think> pre-strip.
def test_think_block_stripped_before_xml(self):
# The think block is stripped before matching so the post-thinking call is recognised.
text = (
"<think>I will use web_search to find the weather.</think>"
'<tool_call>{"name":"web_search","arguments":{"query":"sf"}}</tool_call>'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_think_block_stripped_before_bracket_tag(self):
text = (
"<think>Let me search for that.</think>\n" '[TOOL_CALLS]web_search{"query":"weather"}'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "web_search"
def test_uppercase_think_tag_stripped(self):
# Some templates use [THINK]...[/THINK] instead of <think>.
text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_think_block_hides_inner_tool_call(self):
# A call mentioned inside think is a rehearsal; the wrapper strip removes the inner markup.
text = (
"<think>I might call "
'<tool_call>{"name":"web_search","arguments":{}}</tool_call> '
"but I am not sure</think>\n"
"Let me just answer directly."
)
result = parse_tool_calls_from_text(text)
assert result == []
def test_think_literal_inside_real_tool_argument_is_preserved(self):
# A real call whose argument contains a literal think tag must not be corrupted.
text = (
'<tool_call>{"name":"write","arguments":'
'{"text":"compare <think> and </think> tags"}}</tool_call>'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert json.loads(result[0]["function"]["arguments"])["text"] == (
"compare <think> and </think> tags"
)
def test_bracket_tag_argument_with_think_literal_is_preserved(self):
text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks"
def test_real_call_after_think_with_rehearsal_inside(self):
# A rehearsal inside <think> is skipped, but the real call after the close tag parses.
text = '<think>plan: search[ARGS]{"q":"x"}</think>search[ARGS]{"q":"real"}'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert json.loads(result[0]["function"]["arguments"])["q"] == "real"
# XML takes precedence over bracket-tag.
def test_xml_wins_over_bracket(self):
# When a model emits both forms in one message, the XML form is canonical and wins.
text = (
'<tool_call>{"name":"primary","arguments":{}}</tool_call>'
'[TOOL_CALLS]secondary{"k":"v"}'
)
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert result[0]["function"]["name"] == "primary"
# Strip patterns include bracket-tag and rehearsal.
def test_strip_bracket_tag_closed(self):
text = 'before [TOOL_CALLS]web_search{"q":"hi"} after'
assert "[TOOL_CALLS]" not in strip_tool_markup(text)
assert "before" in strip_tool_markup(text)
assert "after" in strip_tool_markup(text)
def test_strip_rehearsal_closed(self):
text = 'prose python[ARGS]{"code":"x"} more prose'
cleaned = strip_tool_markup(text)
assert "[ARGS]" not in cleaned
assert "prose" in cleaned
assert "more prose" in cleaned
def test_strip_bracket_tag_unclosed_final(self):
text = 'before [TOOL_CALLS]web_search{"q":"part'
# Final-mode strip drops the trailing unclosed run.
cleaned = strip_tool_markup(text, final = True)
assert "TOOL_CALLS" not in cleaned
assert cleaned == "before"
# Canonical Mistral array, v11 [CALL_ID], unified multi-call (PR review fixes).
def test_mistral_canonical_array_is_parsed(self):
# Canonical multi-call array: every call must parse (was dropped then deleted to EOS).
text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}},{"name":"b","arguments":{"y":2}}]'
result = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in result] == ["a", "b"]
assert json.loads(result[0]["function"]["arguments"]) == {"x": 1}
assert json.loads(result[1]["function"]["arguments"]) == {"y": 2}
def test_mistral_array_string_arguments_are_decoded(self):
# OpenAI-spec arguments arrive as a JSON string; decode to an object.
text = '[TOOL_CALLS] [{"name":"a","arguments":"{\\"x\\":1}"}]'
result = parse_tool_calls_from_text(text)
assert len(result) == 1
assert json.loads(result[0]["function"]["arguments"]) == {"x": 1}
def test_mistral_array_scalar_string_argument_not_double_encoded(self):
# A bare scalar string argument in the Mistral array form must be kept
# raw, exactly like the <tool_call> path, so the downstream argument
# healer wraps ``weather`` into the single-string tool's key -- not
# ``"weather"`` with literal quotes from a redundant json.dumps.
array = parse_tool_calls_from_text(
'[TOOL_CALLS][{"name":"web_search","arguments":"weather"}]'
)
xml = parse_tool_calls_from_text(
'<tool_call>{"name":"web_search","arguments":"weather"}</tool_call>'
)
assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather"
healed = _coerce_arguments(
array[0]["function"]["arguments"], heal = True, tool_name = "web_search"
)
assert healed == {"query": "weather"}
def test_mistral_array_strip_keeps_trailing_prose(self):
# The array form must be removed whole, not deleted to end-of-string.
text = 'answer [TOOL_CALLS] [{"name":"a","arguments":{}}] tail'
assert strip_tool_markup(text, final = True) == "answer tail"
def test_mistral_and_rehearsal_in_one_message_both_parse(self):
# A Mistral call and a rehearsal call together: both must parse.
text = '[TOOL_CALLS]a{"x":1} then b[ARGS]{"y":2}'
result = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in result] == ["a", "b"]
def test_mistral_v11_call_id_is_not_the_function_name(self):
# v11 shape: the function name is ``name``, never the opaque call-id token.
result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}')
assert len(result) == 1
assert result[0]["function"]["name"] == "get_weather"
assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"}
# v11 without a call-id parses the same name.
r2 = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[ARGS]{"q":"y"}')
assert r2[0]["function"]["name"] == "get_weather"
def test_strip_preserves_rehearsal_inside_think(self):
# A rehearsal inside <think> is reasoning; strip keeps it verbatim.
text = '<think>plan: search[ARGS]{"q":"x"}</think> A'
out = strip_tool_markup(text, final = True)
assert out == text
assert "search[ARGS]" in out
def test_streaming_strip_preserves_rehearsal_inside_think(self):
# The streaming strip must also preserve a think rehearsal: a mid-stream strip shrinks
# then regrows the cumulative text (corrupts append-by-length consumers). Matches GGUF.
text = '<think>plan: search[ARGS]{"q":"x"}</think> A'
assert strip_tool_markup_streaming(text) == text
assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text
# An unclosed block during streaming is preserved too (the parser keeps it).
partial = '<think>plan: search[ARGS]{"q":"x"}'
assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial
def test_streaming_strip_still_removes_real_call_outside_think(self):
# The think guard must not stop the streaming strip removing a call outside the block.
text = '<think>reason</think> web_search[ARGS]{"q":"x"}'
out = strip_tool_markup_streaming(text, tool_protocol_active = True)
assert "web_search[ARGS]" not in out
assert "<think>reason</think>" in out
def test_strip_bracket_calls_is_linear(self):
# Many complete bracket calls must strip in ~linear time (was O(n^2) per match).
import time
text = '[TOOL_CALLS]f{"a":1}' * 4000 # ~80KB, 4000 complete calls
t0 = time.perf_counter()
out = strip_tool_markup(text, final = True)
elapsed = time.perf_counter() - t0
assert "[TOOL_CALLS]" not in out
assert elapsed < 1.0, f"strip took {elapsed * 1000:.0f}ms on 4000 bracket calls"
def test_streaming_strip_handles_nested_mistral_json(self):
# The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the
# balanced helper must remove the whole call so no trailing brace leaks to the streaming ...
@ -1390,6 +1816,254 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success():
assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation():
# A spent one-shot (render_html) stays in the ORIGINAL tool list; detection is gated on
# that list (matching the strip gate) so a re-emitted repeat is drained and routed to the
# repeat no-op instead of stripped into a blank continuation.
exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
turns = iter(
[
[
'<tool_call>{"name":"render_html","arguments":{"code":"<html>one</html>"}}</tool_call>'
],
['render_html[ARGS]{"code":"<html>two</html>"}'], # spent one-shot rehearsal
["The chart is above."],
]
)
def gen(_messages, *, active_tools = None):
try:
chunks = next(turns)
except StopIteration:
return
acc = ""
for c in chunks:
acc += c
yield acc
events = _collect_events(
run_safetensors_tool_loop(
single_turn = gen,
messages = [{"role": "user", "content": "make a chart"}],
tools = [
{"type": "function", "function": {"name": "render_html"}},
{"type": "function", "function": {"name": "web_search"}},
],
execute_tool = exec_fn,
max_tool_iterations = 5,
)
)
contents = [e["text"] for e in events if e["type"] == "content"]
# render_html ran exactly once; the repeat was a no-op, not a second execution.
assert exec_fn.calls == [("render_html", {"code": "<html>one</html>"})], exec_fn.calls
# The loop continued past the repeat to the real answer (not a blank continuation).
assert any("The chart is above." in t for t in contents), contents
# The raw rehearsal markup never leaked as visible content.
assert not any("render_html[ARGS]" in t for t in contents), contents
def test_rehearsal_call_name_is_not_streamed_before_args():
# A rehearsal whose name and [ARGS] arrive together must drain, not stream the bare name.
loop, exec_fn = _make_loop(
turns = [['web_search[ARGS]{"query":"cats"}'], ["Found."]],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search" in t for t in contents), contents
def test_rehearsal_call_name_split_before_args_is_not_streamed():
# Finding 5: name and [ARGS] in separate chunks -- the bare name is held until [ARGS] arrives.
loop, exec_fn = _make_loop(
turns = [["web_search", '[ARGS]{"query":"cats"}'], ["Found."]],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search" in t for t in contents), contents
def test_plain_word_matching_no_tool_still_streams():
# The prefix guard must not swallow prose: a non-tool bare word streams.
loop, _exec = _make_loop(
turns = [["weather", " is nice today."]],
max_tool_iterations = 1,
)
events = _collect_events(loop)
contents = "".join(e["text"] for e in events if e["type"] == "content")
assert "weather is nice today." in contents, contents
def test_rehearsal_name_after_prose_in_streaming_is_not_streamed():
# After prose has streamed (STREAMING state), a split rehearsal name must still be held.
loop, exec_fn = _make_loop(
turns = [
# _make_loop accumulates these deltas into cumulative snapshots.
["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'],
["Found."],
],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search" in t for t in contents), contents
def test_rehearsal_name_after_prose_same_chunk_in_streaming_is_not_streamed():
# Prose then ``web_search[ARGS]{...}`` in one chunk: the boundary is pulled back over the name.
loop, exec_fn = _make_loop(
turns = [
["Sure. ", 'now web_search[ARGS]{"query":"cats"}'],
["Found."],
],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search" in t for t in contents), contents
def test_initial_buffer_flush_holds_split_rehearsal_name():
# First flush out of BUFFERING applies the same trailing-name hold as STREAMING.
loop, exec_fn = _make_loop(
turns = [["I will use python", '[ARGS]{"code":"print(1)"}'], ["done"]],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("python", {"code": "print(1)"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("python" in t for t in contents), contents
def test_think_rehearsal_streams_monotonically_and_keeps_reasoning():
# A think rehearsal streams the same text the final strip keeps: cumulative content is
# monotonically non-decreasing and ends with the markup intact.
loop, exec_fn = _make_loop(
turns = [["<think>plan ", 'search[ARGS]{"q":"x"}', "</think> visible"]],
max_tool_iterations = 1,
)
events = _collect_events(loop)
contents = [e["text"] for e in events if e["type"] == "content"]
assert exec_fn.calls == [], exec_fn.calls
assert all(len(b) >= len(a) for a, b in zip(contents, contents[1:])), contents
final = contents[-1] if contents else ""
assert 'search[ARGS]{"q":"x"}' in final, contents
assert "visible" in final, contents
def test_plain_answer_ending_with_tool_name_word_is_preserved():
# End-of-stream flush: a plain answer ending on a tool-name word is prose, not dropped.
loop, exec_fn = _make_loop(
turns = [["I think ", "you should ", "web_search"]],
max_tool_iterations = 1,
)
events = _collect_events(loop)
assert exec_fn.calls == [], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert any(t.rstrip().endswith("web_search") for t in contents), contents
def test_long_tool_name_split_rehearsal_is_not_capped_and_executes():
# Finding 10/11: an MCP name longer than the buffer cap, split before [ARGS], is still
# held (self-bounding prefix); no leak and the call executes.
from core.inference.safetensors_agentic import _MAX_BUFFER_CHARS
name = "mcp__github__create_pull_request"
assert len(name) >= _MAX_BUFFER_CHARS, len(name)
exec_fn = FakeExecuteTool(["RESULT"])
_turns = iter([[name, name + '[ARGS]{"x":1}'], ["done"]])
def st(_messages, active_tools = None):
yield from next(_turns)
events = _collect_events(
run_safetensors_tool_loop(
single_turn = st,
messages = [{"role": "user", "content": "go"}],
tools = [{"type": "function", "function": {"name": name}}],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
)
assert exec_fn.calls == [(name, {"x": 1})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any(name in t for t in contents), contents
def test_unrestricted_mode_split_rehearsal_name_is_not_streamed():
# Finding 6: unrestricted mode treats any bare identifier as a possible rehearsal NAME.
exec_fn = FakeExecuteTool(["RESULT"])
_turns = iter([["web_search", 'web_search[ARGS]{"q":"x"}'], ["done"]])
def st(_messages, active_tools = None):
yield from next(_turns)
events = _collect_events(
run_safetensors_tool_loop(
single_turn = st,
messages = [{"role": "user", "content": "go"}],
tools = [], # unrestricted
execute_tool = exec_fn,
max_tool_iterations = 2,
)
)
assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search" in t for t in contents), contents
def test_unrestricted_mode_split_after_bracket_is_not_streamed():
# Unrestricted mode: a chunk split right after ``NAME[`` is still held (parity with the
# restricted-mode startswith hold).
exec_fn = FakeExecuteTool(["RESULT"])
_turns = iter([["web_search[", 'web_search[ARGS]{"q":"x"}'], ["done"]])
def st(_messages, active_tools = None):
yield from next(_turns)
events = _collect_events(
run_safetensors_tool_loop(
single_turn = st,
messages = [{"role": "user", "content": "go"}],
tools = [], # unrestricted
execute_tool = exec_fn,
max_tool_iterations = 2,
)
)
assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert not any("web_search[" in t for t in contents), contents
def test_unrestricted_mode_plain_prose_still_streams():
# The unrestricted hold releases a held identifier once the rest of the sentence follows.
def st(_messages, active_tools = None):
for snap in ("Hello", "Hello there friend."):
yield snap
events = _collect_events(
run_safetensors_tool_loop(
single_turn = st,
messages = [{"role": "user", "content": "hi"}],
tools = [],
execute_tool = FakeExecuteTool([]),
max_tool_iterations = 1,
)
)
contents = "".join(e["text"] for e in events if e["type"] == "content")
assert "Hello there friend." in contents, contents
def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call():
# A late call caught by the safety net: an unclosed ``<tool_call>`` heals only with Auto-Heal on;
# off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call.
@ -1977,6 +2651,42 @@ class TestLoopBasic:
assert tool_starts[0]["tool_name"] == "python"
assert exec_fn.calls == [("python", {"code": "print('<function=render_html>')"})]
def test_render_html_rehearsed_in_think_block_emits_no_provisional_start(self):
# BUG B: a render_html rehearsed inside think before a real python call must not emit a
# provisional render_html card; only the outside-think call fires.
exec_fn = FakeExecuteTool(["ok"])
turn_iter = iter(
[
[
'<think>draft render_html[ARGS]{"code":"x"}</think>',
'python[ARGS]{"code":"print(1)"}',
],
["Done."],
]
)
def _gen(_messages):
chunks = next(turn_iter)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "run code"}],
tools = [
{"type": "function", "function": {"name": "render_html"}},
{"type": "function", "function": {"name": "python"}},
],
execute_tool = exec_fn,
)
events = _collect_events(loop)
tool_starts = [e for e in events if e["type"] == "tool_start"]
assert [e["tool_name"] for e in tool_starts] == ["python"], tool_starts
assert exec_fn.calls == [("python", {"code": "print(1)"})]
def test_render_html_success_blocks_second_canvas_call(self):
exec_fn = FakeExecuteTool(["Rendered HTML canvas."])
turn_iter = iter(
@ -3653,6 +4363,114 @@ if __name__ == "__main__":
pytest.main([__file__, "-v"])
def test_streaming_strip_keeps_bare_args_before_think_block():
# F3: a bare ``foo[ARGS]`` before a think block is prose; EOS-anchored tail arms run only
# on the last segment.
text = "Please pass foo[ARGS] <think>pause</think> to the template."
out = strip_tool_markup_streaming(text, tool_protocol_active = True)
assert out == text
def test_streaming_strip_still_removes_complete_call_before_think_block():
# A complete bracket call before a think block still strips in the non-last segment.
text = 'go web_search[ARGS]{"q":"x"} <think>z</think> done'
out = strip_tool_markup_streaming(text, tool_protocol_active = True)
assert "web_search[ARGS]" not in out
assert "<think>z</think>" in out
assert "go" in out and "done" in out
def test_prose_args_marker_before_real_call_does_not_drain_the_prose():
# F5: an inactive ``foo[ARGS]`` in prose is not a call boundary; the prose streams in
# full and the later real call still executes.
loop, exec_fn = _make_loop(
turns = [
["Intro ", "foo[ARGS] syntax. ", 'web_search[ARGS]{"query":"cats"}'],
["Cats are great."],
],
exec_results = ["RESULT"],
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls
contents = [e["text"] for e in events if e["type"] == "content"]
# The prose between the bogus marker and the real call must survive.
assert any("foo[ARGS] syntax." in t for t in contents), contents
# The real call markup is never shown as content.
assert not any("web_search[ARGS]" in t for t in contents), contents
def test_inactive_name_args_with_body_is_not_parsed_into_disabled_noop():
# BUG A: a prose answer with an inactive ``foo[ARGS]{...}`` is not drained into a
# disabled no-op extra turn; the [ARGS] checks are name-gated.
turns = [['foo[ARGS]{"x":1} is just syntax.']]
turn_calls: list[int] = []
def _gen(_messages):
turn_calls.append(1)
chunks = turns[len(turn_calls) - 1] if len(turn_calls) <= len(turns) else []
acc = ""
for chunk in chunks:
acc += chunk
yield acc
exec_fn = FakeExecuteTool([])
loop = run_safetensors_tool_loop(
single_turn = _gen,
messages = [{"role": "user", "content": "explain"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 3,
)
events = _collect_events(loop)
assert exec_fn.calls == [], exec_fn.calls
assert not any(e["type"] in ("tool_start", "tool_end") for e in events), events
# Exactly one generation turn -- no disabled ``foo`` no-op re-prompt.
assert len(turn_calls) == 1, turn_calls
contents = [e["text"] for e in events if e["type"] == "content"]
assert any("is just syntax." in t for t in contents), contents
class TestEnabledToolNameGate:
"""The safetensors loop passes the active tool names into parse/strip so the
ambiguous bare-rehearsal ``NAME[ARGS]{json}`` is treated as a call only when NAME
is an active tool (#5704). Without the gate an inactive ``foo[ARGS]{...}`` in prose
was parsed into a disabled no-op call and stripped from the visible text."""
def _names(self, calls):
return [c["function"]["name"] for c in calls]
def test_parse_inactive_rehearsal_does_not_swallow_active_call(self):
text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'
calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"})
assert self._names(calls) == ["web_search"]
assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"}
def test_parse_inactive_rehearsal_alone_is_prose(self):
assert (
parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == []
)
def test_streaming_strip_keeps_inactive_rehearsal(self):
raw = 'answer foo[ARGS]{"x":1} tail'
assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw
def test_streaming_strip_removes_active_rehearsal(self):
raw = 'answer web_search[ARGS]{"q":1} tail'
out = strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"})
assert "web_search[ARGS]" not in out
assert out == "answer tail"
def test_final_strip_keeps_inactive_rehearsal(self):
text = 'foo[ARGS]{"x":1} is just syntax.'
assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text
def test_gate_none_preserves_legacy_strip_and_parse(self):
text = 'foo[ARGS]{"x":1} tail'
assert self._names(parse_tool_calls_from_text(text)) == ["foo"]
assert strip_tool_markup_streaming(text) == " tail"
def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled():
# F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did
# not parse must stay visible (disabled-Auto-Heal contract: malformed markup is

141
studio/backend/tests/test_tool_call_parser_strict.py
View file

@ -266,6 +266,133 @@ class TestHealingPathUnaffected:
assert healed[span[0] : span[1]] == "<function=web_search><parameter=query>dogs"
class TestEnabledToolNameGate:
"""``enabled_tool_names`` disambiguates the ambiguous bare-rehearsal
``NAME[ARGS]{json}`` form (#5704): NAME is a call only when it is an active tool,
otherwise it is prose. ``None`` (the default) keeps the legacy unrestricted parse
so existing callers are unaffected."""
def _names(self, calls):
return [c["function"]["name"] for c in calls]
def test_inactive_rehearsal_before_active_call_does_not_swallow_it(self):
# P1: an inactive ``foo[ARGS]{...}`` before a real call must not consume the real call.
text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'
calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"})
assert self._names(calls) == ["web_search"]
assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"}
def test_inactive_rehearsal_alone_is_not_a_call(self):
text = 'foo[ARGS]{"a":1}'
assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == []
def test_active_rehearsal_is_still_parsed(self):
text = 'web_search[ARGS]{"query":"cats"}'
calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"})
assert self._names(calls) == ["web_search"]
def test_unrestricted_gate_none_preserves_legacy_behavior(self):
# Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed.
text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'
assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"]
assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [
"foo",
"web_search",
]
class TestBracketCallSpans:
"""with_spans tiling for Mistral bracket calls: promoted markup strips
exactly once, filtered calls' bytes stay visible, closers strip too."""
def test_mixed_array_filtered_first_keeps_its_bytes_only(self):
from core.inference.passthrough_healing import heal_openai_message_events
tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}]
content = (
'[TOOL_CALLS][{"name":"bad","arguments":{"x":1}},'
'{"name":"lookup","arguments":{"q":"cats"}}]'
)
events = heal_openai_message_events(
{"role": "assistant", "content": content}, {"lookup"}, tools
)
kinds = [k for k, _v in events]
assert kinds == ["text", "tool_call"]
text = events[0][1]
assert '"bad"' in text
# The promoted call's markup must not survive in the text event.
assert '"lookup"' not in text
def test_mixed_array_filtered_second_stays_visible(self):
from core.inference.passthrough_healing import heal_openai_message_events
tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}]
content = (
'[TOOL_CALLS][{"name":"lookup","arguments":{"q":"cats"}},'
'{"name":"bad","arguments":{"x":1}}]'
)
events = heal_openai_message_events(
{"role": "assistant", "content": content}, {"lookup"}, tools
)
assert events[0][0] == "tool_call"
trailing = "".join(v for k, v in events if k == "text")
assert '"bad"' in trailing
def test_v11_closer_inside_span(self):
from core.tool_healing import parse_tool_calls_from_text as parse_with_spans
text = '[TOOL_CALLS]web_search[ARGS]{"query":"cats"}[/TOOL_CALLS] after'
calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True)
(call,) = calls
assert call["function"]["name"] == "web_search"
(span,) = spans
assert text[span[0] : span[1]].endswith("[/TOOL_CALLS]")
assert text[span[1] :] == " after"
def test_fully_promoted_array_strips_whole_region(self):
from core.inference.passthrough_healing import heal_openai_message_events
tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}]
content = (
'[TOOL_CALLS][{"name":"lookup","arguments":{"q":"a"}},'
'{"name":"lookup","arguments":{"q":"b"}}] after'
)
events = heal_openai_message_events(
{"role": "assistant", "content": content}, {"lookup"}, tools
)
assert [k for k, _v in events] == ["tool_call", "tool_call", "text"]
assert events[2][1] == " after"
class TestMistralArrayHealing:
"""Draining the whole [TOOL_CALLS] array for the shapes the repo's own
Mistral/Ollama templates emit."""
def test_comma_less_multi_call_array_parses_all_calls(self):
# ollama_template_mappers.py renders multi-call turns as [{...}{...}] with no
# comma separator; a single json.loads of the body rejects it and dropped every
# call. The element-by-element decode must recover all of them.
text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}}{"name":"b","arguments":{"y":2}}]'
calls = parse_tool_calls_from_text(text)
assert [c["function"]["name"] for c in calls] == ["a", "b"]
assert json.loads(calls[0]["function"]["arguments"]) == {"x": 1}
assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2}
def test_comma_separated_and_single_arrays_still_parse(self):
both = parse_tool_calls_from_text(
'[TOOL_CALLS] [{"name":"a","arguments":{}},{"name":"b","arguments":{}}]'
)
assert [c["function"]["name"] for c in both] == ["a", "b"]
one = parse_tool_calls_from_text('[TOOL_CALLS] [{"name":"a","arguments":{}}]')
assert [c["function"]["name"] for c in one] == ["a"]
def test_mistral_array_null_arguments_normalized_to_empty_object(self):
# ``"arguments": null`` is a no-arg call; it must become {} (as the <tool_call>
# path does), not the string "null" that auto-heal turns into {"query":"null"}.
calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]')
assert calls[0]["function"]["arguments"] == "{}"
class TestGlmStrict:
def test_closed_glm_call_is_accepted(self):
text = (
@ -740,22 +867,26 @@ class TestMistralOuterOverXmlLiteral:
class TestHealerSignalAlignment:
"""The healer buffers only promotable formats; Mistral/Llama text calls stream through."""
"""The healer buffers only formats its shared parser can promote. Mistral's
``[TOOL_CALLS]`` is promotable (rescued), so it is a heal signal; the loop-only
text-call markers (Llama ``<|python_tag|>``, bare ``[ARGS]``) are not, so they
stream through instead of stalling as prose that never yields a call."""
def test_heal_signals_subset_of_promotable_formats(self):
from core.inference.passthrough_healing import _HEAL_SIGNALS
assert set(_HEAL_SIGNALS) == {"<tool_call>", "<|tool_call>", "<function="}
assert set(_HEAL_SIGNALS) == {"<tool_call>", "<|tool_call>", "<function=", "[TOOL_CALLS]"}
def test_stream_healer_does_not_hold_mistral_text(self):
def test_stream_healer_does_not_hold_llama_python_tag_text(self):
from core.inference.passthrough_healing import StreamToolCallHealer
healer = StreamToolCallHealer(
{"web_search"},
[{"type": "function", "function": {"name": "web_search", "parameters": {}}}],
)
events = list(healer.feed('[TOOL_CALLS]web_search[ARGS]{"query":"cats"}'))
# Llama <|python_tag|> is not a healer-promotable format, so it streams through as text.
events = list(healer.feed('<|python_tag|>web_search.call(query="cats")'))
text_out = "".join(v for k, v in events if k == "text")
assert "[TOOL_CALLS]" in text_out # streamed through, not buffered
assert "<|python_tag|>" in text_out # streamed through, not buffered
assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize())

324
studio/backend/tests/test_tool_xml_strip.py
View file

@ -52,6 +52,11 @@ _ns = {
}
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.
@ -66,16 +71,19 @@ assert "_strip_mistral_closed_calls" in _xml_helper.group(
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 _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?: .+\n)+",
r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)",
_src,
_re.DOTALL,
)
assert _helper, "could not extract _strip_tool_xml_for_display source"
# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the
# extracted body actually reached that call rather than truncating early.
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)+",
@ -95,6 +103,56 @@ def test_route_display_strip_respects_disabled_auto_heal_contract():
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).
@ -234,6 +292,32 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws():
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 = (
@ -362,6 +446,238 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam():
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) ──