unsloth/studio/backend
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
..
assets Studio: require signed capability tokens for /p preview links (#6666) 2026-06-25 21:40:48 -07:00
auth Studio: opt-in OpenAI /v1 model auto-switch and idle keep-warm (#6392) 2026-07-01 06:42:23 -07:00
core Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) 2026-07-03 08:22:42 -07:00
hub studio: allow updating HF models through UI (#5388) 2026-07-01 01:54:57 +03:00
loggers Studio: trim serving-log noise and surface llama-server engine stats (#6377) 2026-06-17 05:37:57 -07:00
models Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) 2026-07-03 08:22:42 -07:00
plugins Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
requirements Whole-document context for RAG chat attachments (#6693) 2026-06-30 15:55:23 +02:00
routes Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) 2026-07-03 08:22:42 -07:00
state Studio: shareable per-checkpoint preview links (#6486) 2026-06-24 06:31:53 -07:00
storage Studio: customizable RAG embedding model with HF search, settings tab reorganization (#6800) 2026-07-02 05:26:33 -07:00
tests Tool-call healing (default on) and opt-in nudging for the client-tool passthrough (#6801) 2026-07-03 08:22:42 -07:00
utils [Studio] Add --with-llama-cpp-dir installer flag to reuse a local llama.cpp (#6472) 2026-07-02 22:11:20 +01:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
cloudflare_tunnel.py Reap Studio child processes when the parent dies abnormally (#6425) 2026-06-18 05:51:22 -07:00
colab.py Studio Colab: opt-in shareable Cloudflare tunnel link (#6684) 2026-06-26 00:56:23 -07:00
main.py feat: Implementation of the Portuguese (Brazil) language and VRAM/RAM monitor (#6509) 2026-07-02 18:06:39 +02:00
run.py studio: explicit Cloudflare tunnel notice and public-exposure warning at startup (#6515) 2026-06-30 17:47:48 +02:00
startup_banner.py Fix Windows Studio UTF-8 startup handling (#6614) 2026-07-01 13:47:33 +01:00