Commit graph

4 commits

Author SHA1 Message Date
Daniel Han
85b49eeb56
Studio: Inkling support fixes (#7153)
* Studio: Inkling support fixes (context sizing, tool-call healing, reasoning effort, audio icon)

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-15 11:22:38 -07:00
Apoze
6a9b77ee37
Studio: harden OpenAI-compatible GGUF streaming (#6950)
---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
2026-07-09 12:09:08 -03:00
Daniel Han
f109e7f0e6
Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes (#5704)
* Studio: parse Mistral [TOOL_CALLS] and rehearsal tool-call shapes

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* Studio: preserve XML parameter indentation in tool_healing

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

* Studio: tighten Mistral/rehearsal tool-call comments

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Round-2 review follow-ups.

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

Regression tests added for both.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

* studio: skip tool calls rehearsed in prefilled reasoning

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

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

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

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

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

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

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

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

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

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

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

* studio: keep tool_healing importable on Python 3.9

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

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

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

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

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

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

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

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

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

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

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

Add regression tests for both loops and the history strip.

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

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

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

Add regression tests.

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

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

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

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

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

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

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

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

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

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

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

* verify_import_hoist: exempt __future__ imports and same-diff relocations

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

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

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

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

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

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

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

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

* Gate safetensors reasoning prefill on the rendered generation prompt

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

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

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

* Tighten tool-call rescue and reasoning-prefill comments

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

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

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-07-06 18:52:13 -07:00
Daniel Han
308ea5a93c
Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801)
* inference: add passthrough tool-call healing core (heal_gate, heal_openai_message, StreamToolCallHealer, nudge helpers)

Small GGUF models often emit tool calls as text (<tool_call>{...}</tool_call>,
Gemma <|tool_call>, <function=> XML) instead of structured tool_calls. Studio's
enable-tools loop already heals these, but the client-tool passthrough
(unsloth run --disable-tools, unsloth start agents) relays them verbatim, so
the agent sees prose and the turn dies.

This module is the shared response-side repair layer the passthrough routes
will call: promote parsed text-form calls to structured calls, but only for
function names the client actually declared; coerce arguments through the same
canonical-key healing as the tool loop; never touch the upstream request body
(llama-server KV/slot reuse stays byte-identical). StreamToolCallHealer is the
streaming buffer-and-repair state machine: prose forwards immediately, only a
partial-signal tail or a suspected tool block is held, false alarms flush
verbatim, and a 64 KiB bound caps memory. nudge_should_retry/nudge_messages
support an opt-in single-retry nudge for non-streaming routes (wired later).

Kill-switch: UNSLOTH_DISABLE_TOOL_CALL_HEALING=1. Reuses
core/tool_healing.parse_tool_calls_from_text, strip_tool_call_markup, and
tool_loop_controller.coerce_tool_arguments unchanged.

* inference: heal text-form tool calls on the OpenAI and Responses passthrough

Wire the passthrough healing core into /v1/chat/completions and /v1/responses,
default ON whenever the request declares client tools:

Non-streaming: heal_openai_message runs inside the existing response-mutation
loop; a promoted call flips finish_reason to tool_calls and nulls the content,
and the verbatim-bytes fast path still applies when nothing was healed.
/v1/responses non-streaming inherits this through openai_chat_completions.

Streaming: a StreamToolCallHealer per stream. Ordinary prose relays
byte-for-byte (a fast path keeps upstream bytes when the healer passes a chunk
through whole); once a tool signal appears, content is held, and at the
finish/[DONE] boundary either synthetic delta.tool_calls chunks replace the
markup (finish_reason rewritten to tool_calls, including the synthetic-finish
path) or a false alarm flushes the held text verbatim. Structured upstream
deltas put the healer to sleep after flushing anything held, so grammar-mode
responses stay byte-identical. The Responses stream feeds healed calls through
the same per-call state machinery as structured deltas (indexes live in a
disjoint range so a healed call can never merge into a structured call's
state), and the visible/reasoning split runs first so reasoning text is never
promoted. parallel_tool_calls=false caps healed calls on every path.

The upstream request body is never touched and healing issues no extra
generation, so llama-server slot/KV-cache reuse is unchanged. Opt-out per
request with auto_heal_tool_calls=false (Responses reads it from the
extra-body); requests without tools relay verbatim.

* inference: heal text-form tool calls on the Anthropic /v1/messages passthrough

Streaming: AnthropicPassthroughEmitter.enable_healing(allowed_tools) routes
content deltas through the shared StreamToolCallHealer. A promoted call closes
any open text block (only the safe prose prefix ever streamed into it), opens a
synthetic tool_use block with a fresh toolu_* id, carries one input_json_delta,
and closes; finish() then forces stop_reason to tool_use unless a truncation
(max_tokens) wins. Structured upstream deltas flush anything held and put the
healer to sleep, so grammar-mode responses are untouched, as is every stream
where enable_healing is never called (Studio's own loop, no-tools requests).
disable_parallel_tool_use caps healed calls too.

Non-streaming: the OpenAI message dict is healed BEFORE block building, so the
existing tool_use promotion loop and stop_reason line treat promoted calls
exactly like native ones (finish_reason length still maps to max_tokens). The
legacy tool-XML strip still runs on remaining text, so opted-out requests keep
today's cleanup behavior byte-for-byte.

auto_heal_tool_calls is now a typed field on AnthropicMessagesRequest
(default True, mirroring Chat Completions) and threads into both passthrough
calls. Healing never touches the upstream request body.

* inference: opt-in single-retry tool-call nudge on the non-streaming passthrough

When the model clearly tried to call a tool (a tool signal in the text) but
healing produced nothing usable, re-ask once: the retry body is the original
body plus an assistant turn (the model's own failed text) and a short user
nudge naming the declared tools. The prompt prefix stays byte-identical, so
llama-server reuses the slot's KV cache and only the two-message suffix is
prefilled. The retry replaces the original response only when it actually
yields a promotable or structured call; on any error or still-garbage output
the original response is returned unchanged. Exactly one retry, non-streaming
OpenAI and Anthropic passthroughs only (a stream has already emitted bytes).

OPT-IN per user decision: nudge_tool_calls=true per request (typed on both
ChatCompletionRequest and AnthropicMessagesRequest, lifted from the Responses
extra-body), or UNSLOTH_TOOL_CALL_NUDGE=1 to flip the process default.
auto_heal_tool_calls=false disables healing AND the nudge.

Also align the non-streaming heal on allow_incomplete=True: the response is
final, so a trailing unclosed tool block is a model failure worth repairing,
matching the enable-tools loop's drain semantics.

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

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

* inference: never assume the upstream response shape in the nudge helpers

llama-server error bodies can carry message: null (or no choices at all), and
_last_assistant_text / response_has_promotable_calls / nudge_should_retry
called .get() on the message without a dict check, so a malformed upstream
response raised an AttributeError the surrounding except tuples did not catch,
failing the request instead of degrading to 'nothing to heal'. Route the shape
probing through one _first_choice_message helper that returns None for any
non-dict message, and add a parametrized test over the malformed shapes.

* inference: constrain healing by tool_choice, preserve length finish_reason, keep healed event order in Responses streams

Three review findings on the passthrough healer:

- heal_gate now honors the request's tool_choice: "none" disables healing
  outright and a forced function narrows the promotion allowlist to that
  one function, so healing can never contradict the request's tool-choice
  constraint. Wired through the OpenAI chat (stream and non-stream),
  Responses, and Anthropic (converted shape) passthroughs.
- The OpenAI non-streaming heal only upgrades finish_reason "stop" to
  "tool_calls"; a truncated generation keeps "length" (the healed call
  stays attached) matching the streaming and Anthropic paths.
- The Responses stream emits healer events in order instead of collapsing
  all text ahead of the healed calls, so text after a healed call no longer
  jumps ahead of the function_call item and output indexes are claimed in
  the order the model produced them.

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

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

* inference: all-or-nothing promotion when a response mixes declared and undeclared text-form calls

Promoting a subset used to strip ALL tool markup from the content, which
silently deleted the text of any call naming an undeclared tool. The heal
now declines entirely when any parsed call is unpromotable, so the whole
message relays verbatim (pre-PR behavior) and no bytes are ever lost. In
streaming, a declared call that completed before an undeclared one arrived
is already emitted; the late undeclared markup still flushes as raw text.
The nudge helpers mirror the same contract via a shared predicate.

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

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

* tests: wrap long lines in the Responses healing tests to the project style

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

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

* inference: span-exact healing, disjoint healed stream indexes, per-call Responses message items, allowlisted nudge acceptance

Four review findings on the passthrough healer:

- parse_tool_calls_from_text gains an optional with_spans return so healing
  removes EXACTLY the promoted calls' markup. This supersedes the previous
  all-or-nothing rule: declared calls promote and every unpromoted byte
  (undeclared calls, unparseable closed blocks, suppressed alternate
  formats such as a <function=...> block after a JSON call) relays as text.
  The stream healer also processes one block per pass, so text between two
  healed calls keeps its document position instead of trailing them.
- The OpenAI chat stream shifts native tool-call delta indexes past any
  already-emitted healed calls; clients merge deltas by index, so a healed
  call and a later native call can no longer merge into one.
- A healed call in the Responses stream closes the open message item and
  trailing text opens a fresh one with a later output index, matching the
  native stream shape; response.completed snapshots every message item
  with its own text.
- The nudge retry only replaces the original response when the retry's
  structured call names a DECLARED tool; a hallucinated undeclared call is
  not an improvement.

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

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

* Studio: stop the heal path folding trailing prose into a closed function call

parse_tool_calls_from_text(allow_incomplete=True) cut a <function=...> body only
at an end-anchored </function>, so a fully closed call followed by trailing prose
(<function=..>..</parameter></function> words) folded </parameter></function> and
the prose into the tool argument and deleted the prose from visible content. The
strict path (allow_incomplete=False) already cut at the real </function> via rfind.
Do the same in both modes: trim the body at the real </function> when present and
end the removal span there, falling back to the end-anchored strip and body_end
only when the call is genuinely truncated. Add a regression test.

* inference: one shared single-call budget for healed and native calls

Codex round 5: the parallel-call caps counted healed and native calls
separately, so a healed text-form call followed by a native structured
delta double-emitted on all three streaming surfaces when the client
disabled parallel calls.

- OpenAI SSE: once a healed call went out with parallel_tool_calls
  false, native tool_call deltas are dropped instead of index-shifted.
- Anthropic emitter: native deltas skip block allocation when the
  healed-plus-native count already filled the single slot, and healed
  emission counts open native states too.
- Responses stream: native deltas that survived the chunk-level cap are
  skipped once a healed call claimed the slot.

Also adds a span assertion for the closed-</function> trailing-prose
parse fixed in the previous commit.

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

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

* studio: relay undeclared text-form calls as text on Anthropic non-streaming

heal_openai_message promotes only declared text-form tool calls and
span-trims just their markup, deliberately leaving every unpromoted byte
(undeclared text-form calls included) in the content to relay as text.
The Anthropic non-streaming builder then ran a blanket _TOOL_XML_RE strip
over that content unconditionally, deleting the undeclared block before
building the text part, so Anthropic clients silently lost a call the
OpenAI non-streaming path preserves. The strip was harmless when healing
was all-or-nothing but became data loss once healing turned span-exact.

Gate the legacy strip on whether healing promoted a call, matching the
OpenAI passthrough and the intent already stated in the comment above.
Add a route-level regression test for the mixed declared+undeclared case.

* inference: require fully declared nudge retries; keep unpromoted Anthropic text

Codex round 6, two findings:

- response_has_promotable_calls accepted a nudge retry when any one
  structured call named a declared tool, so a mixed retry (hallucinated
  undeclared call plus a declared one) replaced the original and the
  caller forwarded the undeclared call, or with parallel_tool_calls
  false could keep only it. All structured retry calls must be declared.

- The Anthropic non-streaming builder still ran the legacy _TOOL_XML_RE
  strip after span-exact healing, deleting undeclared or malformed call
  text that healing deliberately preserved. The legacy strip now runs
  only when healing is off (no declared tools, or opted out), matching
  the OpenAI passthrough.

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

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

* inference: keep unpromoted Anthropic text whenever healing is active

The previous commit skipped the legacy strip only when a call was
actually promoted, so an undeclared-only (or malformed-only) response
was still silently emptied: exactly the dead-turn shape this path
exists to fix, and inconsistent with the OpenAI passthrough, which
relays those bytes verbatim. Gate the strip on healing being active
instead; opt-out and no-tools requests keep the legacy strip.

* Fix schema-aware tool healing for PR #6801

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

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

* Fix passthrough healing ordering for PR #6801

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

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

* Fix stream finish ordering for PR #6801

* [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>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com>
2026-07-03 08:22:42 -07:00