Studio: client-tool passthrough healing for safetensors and MLX - #6870
Conversation
PR 6801 made response-side tool-call healing default-on for the client-tool passthrough, but only on the GGUF path: the passthrough branch in /v1/chat/completions is gated on using_gguf, and the safetensors section never reads payload.tools, so a client-tools request against a safetensors or MLX model silently dropped the tool schemas and returned prose with no tool_calls. Add the missing leg. When a non-GGUF model is loaded, the request declares client tools (or carries tool-role history), server-side tools are off, and the template supports tools, the route now: - renders the tools into the chat template for a single turn via the existing backend.generate_chat_response(..., tools=...) seam (worker templating already accepts role=tool and assistant.tool_calls messages, normalized with _openai_messages_for_passthrough); - non-streaming: promotes text-form calls with heal_openai_message, honors the opt-in nudge single retry (nudge_should_retry / nudge_messages), caps healed calls when parallel_tool_calls=false (covers the nudge retry too), and sets finish_reason=tool_calls with content null on a pure tool-call turn; - streaming: derives deltas from the worker's cumulative snapshots and feeds StreamToolCallHealer, emitting healed tool-call deltas and the correct finish chunk, guarded against repeated or shrinking snapshots. heal_gate semantics are identical to the GGUF passthrough: default on, auto_heal_tool_calls=false or UNSLOTH_DISABLE_TOOL_CALL_HEALING=1 relays verbatim, tool_choice narrows promotion, undeclared names stay text. MLX rides the same orchestrator seam, so both local backends gain the behavior. CompletionMessage.content becomes Optional so a promoted pure tool-call turn matches the OpenAI contract (content null when only tool_calls return). Adds tests/test_sf_client_tools_passthrough.py (22 cases: healing, gating, opt-outs, streaming deltas, tool-role history, dict-arguments history, forced tool_choice, parallel cap, usage, nudge on/off/double-failure, generator error hygiene, disconnect reset, empty output, MLX path).
|
@codex review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements client-tool passthrough healing for the safetensors and MLX backends, aligning their behavior with the GGUF backend. When a request declares client tools but server-side tools are disabled, text-form tool calls are automatically promoted back into structured tool calls using a shared healer. This is supported in both streaming and non-streaming modes, and includes an optional nudge retry mechanism for unparseable markup. Additionally, the CompletionMessage model was updated to allow an optional content field to accommodate pure tool-call turns, and a comprehensive suite of unit tests was added to verify these changes. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5d87168478
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| and image is None | ||
| and not _sf_is_gptoss | ||
| and _sf_features.get("supports_tools", False) | ||
| and ((payload.tools and len(payload.tools) > 0) or _sf_has_tool_msgs) |
There was a problem hiding this comment.
When tool_choice is "none", this predicate still becomes true for any request that includes a tools array, and the branch below passes those tools into the safetensors/MLX chat template. These local generators do not receive tool_choice, while heal_gate() disables promotion for "none", so clients that keep their tools array but force a final-answer turn with tool_choice: "none" can still prompt the model to emit tool-call markup that is then relayed as ordinary content. Please skip client-tool advertising, or pass no tools, when tool_choice == "none".
Useful? React with 👍 / 👎.
| gen_kwargs["messages"] = _openai_messages_for_passthrough(payload) | ||
| gen_kwargs["system_prompt"] = "" |
There was a problem hiding this comment.
When a tools request contains an OpenAI developer message, this path replaces the already-normalized chat_messages/system_prompt with _openai_messages_for_passthrough() and then clears system_prompt. Unlike _extract_content_parts() and the GGUF passthrough's _set_or_prepend_system_message(), that leaves the developer role in the message list; local templates that do not accept that role can fail rendering, and the transformers fallback formatter drops developer turns. Fold system/developer messages into a leading system prompt before clearing system_prompt.
Useful? React with 👍 / 👎.
| for token in generate( | ||
| [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] | ||
| ): | ||
| retry_text = token |
There was a problem hiding this comment.
When nudge_tool_calls is enabled and the first safetensors/MLX response contains an unhealable tool signal, this retry loop runs inside the main completion try without its own error handling. If the retry generation fails or is cancelled after the original response was already generated, the outer handler turns the whole request into a 500 instead of returning the original text; the GGUF passthrough nudge path catches retry failures and keeps the first response. Wrap only the retry generation and fall back to _msg/full_text on retry errors.
Useful? React with 👍 / 👎.
| if heal_openai_message(_msg, _sf_heal, payload.tools): | ||
| _finish = "tool_calls" |
There was a problem hiding this comment.
When this non-streaming healer succeeds, it mutates _msg into the structured response the client receives, but full_text remains the original <tool_call>... markup and is later stored with api_monitor.set_reply(monitor_id, full_text). For every healed pure tool-call turn, the API monitor/UI therefore shows raw XML rather than the actual tool call; set the monitor reply from _msg['content'] or _monitor_tool_calls_text(_msg['tool_calls']) after promotion.
Useful? React with 👍 / 👎.
…monitor reply Four review follow-ups on the safetensors/MLX client-tool passthrough leg: - tool_choice="none" keeps the tool-history templating but no longer advertises the tools, so a forced final-answer turn is not prompted into emitting markup that the (correctly disabled) healer would relay as prose. Mirrors the GGUF passthrough where llama-server honors tool_choice itself. - OpenAI "developer" messages fold into a single leading system message via _set_or_prepend_system_message before templating; local templates reject the role and the fallback formatter drops it. - A nudge retry that fails or is cancelled after the original answer exists falls back to the first response instead of surfacing a 500, matching the GGUF nudge path. - The API monitor records the healed tool call summary instead of the raw markup on a promoted turn. Adds four regression tests.
|
@codex review |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f9be9e0868
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # prompted to emit tool markup that heal_gate (correctly off for | ||
| # "none") then relays as ordinary content. Mirrors the GGUF passthrough, | ||
| # where llama-server receives and honors tool_choice itself. | ||
| gen_kwargs["tools"] = None if payload.tool_choice == "none" else payload.tools |
There was a problem hiding this comment.
When a non-GGUF client-tools request uses a forced tool_choice (for example forcing search while declaring both lookup and search), this branch still renders every declared tool into the local chat template and never passes tool_choice to the worker/template. The healer later narrows promotion to only the forced name, so if the model follows the advertised but unforced tool schema it can emit a valid lookup call that is relayed as plain content with finish_reason="stop" instead of a tool call. Filter gen_kwargs["tools"] to the forced function, or otherwise pass/enforce tool_choice, before templating.
Useful? React with 👍 / 👎.
| gen_kwargs["messages"] = _set_or_prepend_system_message( | ||
| _openai_messages_for_passthrough(payload), system_prompt |
There was a problem hiding this comment.
When a tool-bearing request includes an image_url part with a remote URL, _extract_content_parts leaves image as None, so this new branch runs and _openai_messages_for_passthrough forwards the original multimodal content list to the safetensors/MLX tokenizer. The previous non-GGUF path flattened these messages to text; now text-only local templates can receive unsupported image dicts and raise during apply_chat_template instead of returning a completion. Gate on the original message content or flatten/sanitize before calling this helper.
Useful? React with 👍 / 👎.
| else: | ||
| for line in _sf_heal_events_to_sse( | ||
| healer.feed(new_text), |
There was a problem hiding this comment.
In streaming safetensors/MLX passthrough, when healing is active and the model emits <tool_call>..., the client receives structured delta.tool_calls from this branch, but the monitor entry has already been appended with the raw new_text before feeding the healer. That makes /api/monitor show XML/tool JSON that was not actually returned (unlike the non-streaming path, which rewrites the monitor reply to [tool_calls]...). Append based on healer events or replace the monitor reply after finalization.
Useful? React with 👍 / 👎.
…g, stream monitor parity - A forced tool_choice function is now the only schema rendered into the local template, so the advertised tools and the healer allowlist can no longer disagree (llama-server enforces tool_choice itself on the GGUF path). - Content-part lists are flattened to their text parts before templating. Remote image URLs are not decodable locally, so such requests reached this path with part lists that raise inside apply_chat_template on text-only templates; the plain non-GGUF path has always flattened them. - The streaming monitor entry is now fed from the healed events the client actually receives, recording promoted calls as the [tool_calls] summary the non-streaming path records.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 32070dc350
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) | ||
| _sf_client_tools = ( | ||
| not _effective_enable_tools(payload) | ||
| and not payload.mcp_enabled |
There was a problem hiding this comment.
When mcp_enabled=true is present but the MCP/server-side loop is not actually used (for example a CLI --disable-tools policy is active, or no enabled MCP tools are returned), this raw flag still makes _sf_client_tools false. The request then falls through to plain safetensors/MLX generation where payload.tools are not passed, so OpenAI-compatible clients that set mcp_enabled globally lose their declared tool schemas instead of getting the new passthrough/healing behavior. Gate this on whether _sf_use_tools actually survived, not on the raw request flag.
Useful? React with 👍 / 👎.
| yield _chat_content_chunk(completion_id, created, model_name, new_text) | ||
| else: | ||
| for line in _sf_heal_events_to_sse( | ||
| healer.feed(new_text), |
There was a problem hiding this comment.
If the backend emits a shorter cumulative snapshot and then resumes, prev_text is reset to the shorter value, so the next full snapshot feeds only the suffix into healer.feed(new_text). After a prior complete tool call has already been promoted, that suffix no longer contains the opening <tool_call> signal and is streamed as normal content, so clients can receive a structured tool_calls delta followed by leaked XML tail. Detect shrinking snapshots before feeding the healer, or keep the cursor monotonic/reset the healer consistently.
Useful? React with 👍 / 👎.
| gen_kwargs["messages"] = _set_or_prepend_system_message( | ||
| _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)), | ||
| system_prompt, |
There was a problem hiding this comment.
On follow-up client-tool turns from standard OpenAI clients, prior assistant tool_calls carry function.arguments as a JSON string (the same shape this endpoint returns). This forwards that OpenAI shape directly into safetensors/MLX chat templates; templates such as studio/backend/assets/chat_templates/gemma-4.jinja:252-264 explicitly raise unless arguments is a mapping, after which the fallback formatter drops tool roles/tool_calls. That makes the next turn lose the tool-result context or fail to render, so deserialize arguments for template input while keeping response output OpenAI-shaped.
Useful? React with 👍 / 👎.
…ize templated arguments - The client-tools gate now keys on _sf_use_tools (whether the server-side tool path actually claimed the request) instead of the raw mcp_enabled flag: with an empty MCP registry or a CLI --disable-tools policy, a client that sets mcp_enabled while declaring its own tools fell through to plain generation with the tools silently dropped. The GGUF passthrough gate has no mcp_enabled clause either. - New _structured_tool_history_for_local_template deserializes assistant tool_calls[].function.arguments JSON strings into mappings for the templated copy only: spec-compliant clients send strings, but local chat templates iterate arguments as a mapping or raise on strings, which crashed or misrendered multi-turn tool history. The HTTP response and the GGUF wire shape keep strings.
|
@codex review |
for more information, see https://pre-commit.ci
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
1 similar comment
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 923abf4fbe
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for line in _sf_heal_events_to_sse( | ||
| healer.feed(new_text), |
There was a problem hiding this comment.
When a non-GGUF tool-passthrough stream receives a cumulative snapshot that is shorter than the previous one, prev_text is overwritten with that shorter value and the next full snapshot's already-processed suffix is fed into the healer here. For example, after a healed <tool_call>... snapshot followed by a shorter/reset snapshot, the next full snapshot can leak the tail of the XML as normal content after the tool_calls delta, so clients see both a tool call and stray assistant prose. Please skip or reset shorter snapshots before feeding new_text to the healer.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The safetensors path here consumes a single cumulative generator whose text only grows, so a snapshot shorter than the previous one is not produced on this path and prev_text cannot regress on legitimate streaming.
| for token in generate( | ||
| [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] | ||
| ): |
There was a problem hiding this comment.
When nudge_tool_calls=true and the retry completes but still does not produce a healable tool call, this second generate(...) call overwrites the shared stats_holder; the code then keeps the original response text but records usage from the discarded retry below. In that double-failure scenario the API monitor attributes token counts/timings to a completion the client never received, so the first attempt's stats should be restored unless the retry response replaces the original.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct that the discarded-retry double-failure records the retry's usage; that is a low-priority monitor-accounting nuance (the client still gets the first response), tracking separately.
There was a problem hiding this comment.
Fixed in 6369a3a: the discarded nudge retry now restores the first attempt's stats on both the no-heal and exception paths, so the monitor reports the usage of the response actually delivered.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6369a3a7e6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| yield _chat_content_chunk(completion_id, created, model_name, new_text) | ||
| else: | ||
| for line in _sf_heal_events_to_sse( | ||
| healer.feed(new_text), |
There was a problem hiding this comment.
When a safetensors/MLX stream with client-tool healing emits cumulative snapshots that shrink or restart, prev_text is updated to the shorter snapshot and the next full snapshot turns new_text into a suffix such as _call>...; feeding that suffix to the healer relays it as ordinary content after the tool call was already promoted. Clients then receive raw tool markup alongside finish_reason="tool_calls", so the offset needs to stay monotonic or the healer should be reset/dropped when the new cumulative text no longer extends the previous one.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
new_text cannot become a mid-token suffix like _call>: that needs cumulative to reintroduce <tool_call> at a position prev_text already passed without prev_text containing it, which is impossible when cumulative is append-only ASCII and prev_text is set to the last cumulative each step. No raw tool markup is relayed alongside finish_reason=tool_calls on this single-generation passthrough path.
| yield _chat_final_chunk(completion_id, created, model_name, "stop") | ||
| if healer is not None: | ||
| for line in _sf_heal_events_to_sse( | ||
| healer.finalize(), |
There was a problem hiding this comment.
When a streaming request is cancelled via the cancel registry, the loop breaks on cancel_event just above but still calls healer.finalize(). finalize() accepts incomplete tool markup at EOF, so a cancelled stream that has buffered <function=...> or a balanced <tool_call>{...} without its closing tag can still emit a structured tool_calls delta and finish_reason="tool_calls", causing clients to execute a tool after the user cancelled; skip healing finalization on cancelled streams or flush the held text instead.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 57428a0: a cancelled stream now skips healer.finalize() and ends with finish_reason=stop, so buffered-but-unclosed tool markup is no longer promoted after the user hits Stop. The disconnect path already returned before finalize; the registry Stop path is now guarded on cancel_event too. Added a regression test driving a Stop mid-emission.
The streaming client-tool heal path breaks out of the token loop when cancel_event is set (the registry "Stop" path), but then still fell through to healer.finalize(), which heals incomplete tool markup at EOF (allow_incomplete) and emits a tool_calls delta plus finish_reason=tool_calls. Because the Stop request only sets the event and leaves the SSE socket open, the client received that promoted call and executed a tool the user had just cancelled. The disconnect path already returns before finalize; guard finalize and the finish_reason on cancel_event too, so a cancelled stream ends with finish_reason=stop and no tool call. Adds a regression test driving a Stop mid-emission with buffered markup.
|
@codex review |
for more information, see https://pre-commit.ci
|
Codex Review: Didn't find any major issues. Another round soon, please! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8e3b37cf45
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| yield _chat_content_chunk(completion_id, created, model_name, new_text) | ||
| else: | ||
| for line in _sf_heal_events_to_sse( | ||
| healer.feed(new_text), |
There was a problem hiding this comment.
When the local generator yields a shorter cumulative snapshot after a longer one, prev_text has already been reset to that shorter value, so the next full snapshot feeds only its suffix into the healer here. In the shrunk-snapshot case covered by the new test shape (full <tool_call>..., then '<tool', then full <tool_call>... again), the first full snapshot is promoted to tool_calls, but the final snapshot streams the raw _call>... </tool_call> tail as content. Guard by ignoring non-prefix snapshots or resetting the healer/diff state when the cumulative text shrinks.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Both streamers are monotonic per the streaming contract: safetensors appends tokens and only strips special tokens or the EOS tail, and MLX decodes a growing token list with skip_special_tokens, so neither can regress the cursor back across an already-promoted tool_call block. The only real non-monotonicity is a 1-2 char UTF-8/space tail correction, which the empty-delta guard (if not new_text: continue) already absorbs. This is the third restatement of the same theoretical concern at line 7030; the two priors were rejected and the cited shrinking test shape does not exist in the PR.
|
@codex review |
|
Codex Review: Didn't find any major issues. 🎉 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
for more information, see https://pre-commit.ci
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
…rompt A vision turn renders through render_prompt_with_boundary, which never took tools, so every safetensors loop turn carrying an image was prompted without the catalog the loop still parses calls against. A catalog now renders through apply_chat_template_for_generation, the choke point that already owns the TypeError ladder and the re-sweep an attempt without tools= needs (unslothai#7066), rather than a second thinner copy of it. The no-tools path is untouched, so the continuation splice from "Studio: continue a response that stopped early (unslothai#8064)" runs exactly as before. The client-tool passthrough gate from "Studio: client-tool passthrough healing for safetensors and MLX (unslothai#6870)" still withdrew tools for an attachment. That guard predates the loop being able to carry a picture; it is widened, not dropped, so only a model that cannot read images gives its tools up, matching the server-side gate. A fresh attachment plus tool-returned pictures in history never matched the history shape's marker count, so the render collapsed to one turn and dropped the exchange. The attachment is marked on the newest user turn and each shape carries the pixel order its markers imply. Capped the pictures a long tool loop accumulates. trim_image_turns drops the oldest markers together with their own pixels, or the processor is left counting image tokens it was given none for. The image turn states when it carries fewer pictures than the tool result's note reports, so the model is not left waiting for the rest.
1 participant
Summary
PR #6801 made response-side tool-call healing default-on for the client-tool passthrough, but only on the GGUF path: the passthrough branch in
/v1/chat/completionsis gated onusing_gguf, and the safetensors section never readspayload.tools. A client-tools request against a safetensors or MLX model silently dropped the tool schemas and returned prose with notool_calls.This adds the missing leg. When a non-GGUF model is loaded, the request declares client tools (or carries tool-role history), server-side tools are off, and the template supports tools, the route now:
backend.generate_chat_response(..., tools=...)seam (worker templating already acceptsrole=toolandassistant.tool_callsmessages, normalized with_openai_messages_for_passthrough);heal_openai_message, honors the opt-in nudge single retry, caps healed calls whenparallel_tool_calls=false(covering the nudge retry), and setsfinish_reason="tool_calls"withcontent: nullon a pure tool-call turn;StreamToolCallHealer, emitting healed tool-call deltas and the correct finish chunk, guarded against repeated or shrinking snapshots.heal_gatesemantics match the GGUF passthrough exactly: default on,auto_heal_tool_calls=falseorUNSLOTH_DISABLE_TOOL_CALL_HEALING=1relays verbatim,tool_choicenarrows promotion, undeclared names stay text. MLX rides the same orchestrator seam, so both local backends gain the behavior.passthrough_healing.pyis reused as-is.CompletionMessage.contentbecomesOptionalso a promoted pure tool-call turn matches the OpenAI contract.Tests
New
tests/test_sf_client_tools_passthrough.py(22 cases): healing, declared-only gating, opt-outs and env kill-switch, streaming healed deltas, tool-role history round-trip, dict-arguments history, forced tool_choice, parallel cap, usage propagation, nudge on/off/double-failure, generator error hygiene, disconnect reset, empty output, MLX path parity.Validated on the staging fork with all 46 checks green: Linux backend pytest across Python 3.10-3.13, MLX CI on a real Mac M1 (imports + end-to-end MLX smoke), the Mac GGUF tool-calling smoke, and cross-platform parity.
Follow-up to #6801.