vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Studio: client-tool passthrough healing for safetensors and MLX (#6870)

Browse source 
* Studio: client-tool passthrough healing for safetensors and MLX

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).

* Address review: tool_choice none, developer folding, retry fallback, 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.

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

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

* Address review: forced tool_choice templating, content-part flattening, 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.

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

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

* Address review: gate passthrough on the engaged server path, deserialize 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.

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

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

* Tighten comments and docstrings in the client-tools passthrough

* Report first-attempt usage when a nudge retry is discarded

When nudge_should_retry fires but the retry produces no healable tool call
(or raises), the first response is still delivered to the client. The retry's
generate() had already overwritten stats_holder, so _monitor_usage recorded
the unseen retry's token counts against the request instead of the first
attempt that was actually returned. Capture the first attempt's stats before
the retry and restore them on both the no-heal and exception paths so the
monitor reports the usage of the response the caller received.

* Do not promote buffered tool markup when a stream is cancelled

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.

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

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

* Trim comments in the client-tools passthrough

* Trim client-tools passthrough comments further

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-06 19:48:36 -07:00 • committed by GitHub
parent 8ba46b566a
commit 46ab683065
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 1102 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

3
studio/backend/models/inference.py
View file

@ -1146,7 +1146,8 @@ class CompletionMessage(BaseModel):
"""The assistant's complete response message."""
role: Literal["assistant"] = "assistant"
content: str
# ``None`` on a pure tool-call turn (OpenAI content=null); string otherwise.
content: Optional[str] = None
refusal: Optional[str] = None
reasoning_content: Optional[str] = None
tool_calls: Optional[list[dict]] = None

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

@ -625,6 +625,66 @@ def _chat_final_chunk(completion_id, created, model_name, finish_reason) -> str:
)
def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> str:
"""Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``)."""
return _chat_chunk_sse(
completion_id,
created,
model_name,
delta = ChoiceDelta(tool_calls = tool_calls),
finish_reason = None,
)
def _sf_heal_events_to_sse(
events,
completion_id,
created,
model_name,
state,
parallel_tool_calls,
monitor_id = None,
):
"""Serialize ``StreamToolCallHealer`` events into chat SSE lines.
``state["idx"]`` tracks the call index across ``feed``/``finalize``;
``parallel_tool_calls is False`` caps promotion to one call (GGUF parity).
The monitor is fed from the same events the client receives, never the
healed-away markup."""
lines = []
for kind, value in events:
if kind == "text":
if value:
lines.append(_chat_content_chunk(completion_id, created, model_name, value))
api_monitor.append_reply(monitor_id, value)
continue
if parallel_tool_calls is False and state["idx"] >= 1:
continue
lines.append(
_chat_tool_calls_chunk(
completion_id,
created,
model_name,
[
{
"index": state["idx"],
"id": value["id"],
"type": "function",
"function": value["function"],
}
],
)
)
_fn = value.get("function") or {}
api_monitor.append_reply(
monitor_id,
("[tool_calls] " if state["idx"] == 0 else "; ")
+ f"{_fn.get('name', '')}({_fn.get('arguments', '')})",
)
state["idx"] += 1
return lines
def _rewrite_cmpl_id(raw: bytes) -> bytes:
"""Rewrite llama-server's chat-style ``chatcmpl-`` ids to the ``cmpl-``
prefix OpenAI's legacy /v1/completions use. Anchored on the ``"id":`` key
@ -7190,25 +7250,87 @@ async def openai_chat_completions(
if payload.preserve_thinking is not None:
gen_kwargs["preserve_thinking"] = payload.preserve_thinking
# ── Client-tool passthrough (safetensors + MLX) ──────────────
# Client tools (or tool-result history) without server-side tools: render
# tools into the template, generate one turn, heal text-form calls (#6801).
# supports_tools=False falls through to plain relay (GGUF gate parity).
_sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages)
# Gate on _sf_use_tools (did the server-side path claim the request?), not
# raw mcp_enabled: an empty MCP registry must not silently drop client tools.
_sf_client_tools = (
not _effective_enable_tools(payload)
and not _sf_use_tools
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)
)
_sf_heal = (
heal_gate(payload.auto_heal_tool_calls, payload.tools, payload.tool_choice)
if _sf_client_tools
else None
)
if _sf_client_tools:
# Re-derive from payload.messages so tool_calls / role="tool" history
# survives templating; fold system/developer into one leading system
# message (templates reject "developer") and clear prompt to avoid a dup.
gen_kwargs["messages"] = _set_or_prepend_system_message(
_structured_tool_history_for_local_template(
_flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload))
),
system_prompt,
)
gen_kwargs["system_prompt"] = ""
# tool_choice="none": keep history templating but advertise no tools
# (heal_gate is off, markup would relay as prose). A forced function
# narrows templating to that one schema. Both mirror the GGUF path,
# where llama-server honors tool_choice itself.
_sf_tc = payload.tool_choice
_sf_forced = None
if isinstance(_sf_tc, dict) and isinstance(_sf_tc.get("function"), dict):
_sf_forced = _sf_tc["function"].get("name")
if _sf_tc == "none":
gen_kwargs["tools"] = None
elif isinstance(_sf_forced, str):
gen_kwargs["tools"] = [
t
for t in payload.tools or []
if isinstance(t, dict)
and isinstance(t.get("function"), dict)
and t["function"].get("name") == _sf_forced
] or None
else:
gen_kwargs["tools"] = payload.tools
# Request-scoped usage/timings receptacle (filled at gen_done).
stats_holder: dict = {}
if payload.use_adapter is not None:
def generate():
def generate(messages_override = None):
kw = (
gen_kwargs
if messages_override is None
else {**gen_kwargs, "messages": messages_override}
)
return backend.generate_with_adapter_control(
use_adapter = payload.use_adapter,
cancel_event = cancel_event,
stats_holder = stats_holder,
**gen_kwargs,
**kw,
)
else:
def generate():
def generate(messages_override = None):
kw = (
gen_kwargs
if messages_override is None
else {**gen_kwargs, "messages": messages_override}
)
return backend.generate_chat_response(
cancel_event = cancel_event,
stats_holder = stats_holder,
**gen_kwargs,
**kw,
)
# ── Streaming response ────────────────────────────────────────
@ -7224,6 +7346,11 @@ async def openai_chat_completions(
try:
yield _chat_role_chunk(completion_id, created, model_name)
# Client-tool passthrough: heal text-form calls on the fly
# (None => relay verbatim).
healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None
heal_state = {"idx": 0}
prev_text = ""
# Split prefilled <think> into reasoning_content deltas (GGUF parity); single turn, serves MLX.
reasoning_extractor = _new_sf_reasoning_extractor()
@ -7255,22 +7382,76 @@ async def openai_chat_completions(
prev_text = cumulative
if not new_text:
continue
# Split prefilled <think> reasoning first (GGUF/MLX parity),
# then route only the visible text through the client-tool
# healer so tool markup inside a reasoning block is not promoted.
reasoning_delta, visible_delta = reasoning_extractor.feed(new_text)
if reasoning_delta:
yield _chat_reasoning_chunk(
completion_id, created, model_name, reasoning_delta
)
if visible_delta:
api_monitor.append_reply(monitor_id, visible_delta)
yield _chat_content_chunk(completion_id, created, model_name, visible_delta)
if healer is None:
# Monitor mirrors the verbatim relay; with healing on,
# _sf_heal_events_to_sse records the healed events instead.
api_monitor.append_reply(monitor_id, visible_delta)
yield _chat_content_chunk(
completion_id, created, model_name, visible_delta
)
else:
for line in _sf_heal_events_to_sse(
healer.feed(visible_delta),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
final_reasoning, final_visible = reasoning_extractor.finish()
if final_reasoning:
yield _chat_reasoning_chunk(completion_id, created, model_name, final_reasoning)
if final_visible:
api_monitor.append_reply(monitor_id, final_visible)
yield _chat_content_chunk(completion_id, created, model_name, final_visible)
yield _chat_final_chunk(completion_id, created, model_name, "stop")
if healer is None:
api_monitor.append_reply(monitor_id, final_visible)
yield _chat_content_chunk(completion_id, created, model_name, final_visible)
else:
for line in _sf_heal_events_to_sse(
healer.feed(final_visible),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
# A cancelled stream must not promote buffered-but-incomplete
# markup: finalize()'s allow_incomplete heal would execute a tool
# the user just cancelled. Disconnect returns earlier; "Stop" only
# sets cancel_event, so guard on it here too.
_cancelled = cancel_event.is_set()
if healer is not None and not _cancelled:
for line in _sf_heal_events_to_sse(
healer.finalize(),
completion_id,
created,
model_name,
heal_state,
payload.parallel_tool_calls,
monitor_id,
):
yield line
_finish = (
"tool_calls"
if (healer is not None and not _cancelled and healer.healed)
else "stop"
)
yield _chat_final_chunk(completion_id, created, model_name, _finish)
# Usage chunk (choices=[], usage set), same shape as the
# GGUF path so the speed popover works for MLX too.
# Request-scoped holder, so concurrent streams cannot
@ -7332,27 +7513,96 @@ async def openai_chat_completions(
for token in generate():
full_text = token
# Split prefilled <think> reasoning (GGUF parity); also covers MLX via the shared generate().
# Split prefilled <think> reasoning (GGUF parity); also covers MLX via
# the shared generate(). Client-tool healing then runs on the visible
# text so tool markup inside a reasoning block is never promoted.
_reasoning_text, _visible_text = _extract_responses_reasoning(
full_text,
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
_plain_msg_kwargs = {"content": _visible_text}
# Client-tool passthrough: promote text-form calls; opt-in single
# nudge retry on unparseable tool markup.
_msg = {"role": "assistant", "content": _visible_text}
if _reasoning_text:
_plain_msg_kwargs["reasoning_content"] = _reasoning_text
_msg["reasoning_content"] = _reasoning_text
_finish = "stop"
if _sf_heal:
if heal_openai_message(_msg, _sf_heal, payload.tools):
_finish = "tool_calls"
elif nudge_enabled(payload.nudge_tool_calls):
_data = {
"choices": [{"message": {"role": "assistant", "content": _visible_text}}]
}
if nudge_should_retry(_data, _sf_heal, payload.tools):
# A failed retry must not 500 the request; keep the first
# response (GGUF nudge parity). The retry's generate()
# overwrites stats_holder, so save the first attempt's stats
# and restore them if the retry is discarded.
_first_stats = stats_holder.get("stats")
try:
retry_text = ""
for token in generate(
[*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)]
):
retry_text = token
# Re-split reasoning on the retry so its visible text is
# what heals into a call (and reaches the monitor).
_retry_reasoning, _retry_visible = _extract_responses_reasoning(
retry_text,
parse_think_markers = _sf_parse_think,
reasoning_prefilled = _sf_reasoning_prefilled,
)
retry_msg = {"role": "assistant", "content": _retry_visible}
if _retry_reasoning:
retry_msg["reasoning_content"] = _retry_reasoning
if heal_openai_message(retry_msg, _sf_heal, payload.tools):
_visible_text, _msg, _finish = (
_retry_visible,
retry_msg,
"tool_calls",
)
else:
# Retry produced no healable call -> first response wins.
stats_holder["stats"] = _first_stats
except Exception as retry_exc:
logger.debug(
"Nudge retry failed; keeping first response: %s", retry_exc
)
stats_holder["stats"] = _first_stats
# parallel_tool_calls=false: cap to one call (GGUF parity).
if payload.parallel_tool_calls is False:
_tcs = _msg.get("tool_calls")
if isinstance(_tcs, list) and len(_tcs) > 1:
_msg["tool_calls"] = _tcs[:1]
response = ChatCompletion(
id = completion_id,
created = created,
model = model_name,
choices = [
CompletionChoice(
message = CompletionMessage(**_plain_msg_kwargs),
finish_reason = "stop",
message = CompletionMessage(
content = _msg["content"],
reasoning_content = _msg.get("reasoning_content"),
tool_calls = _msg.get("tool_calls"),
),
finish_reason = _finish,
)
],
)
api_monitor.set_reply(monitor_id, _visible_text)
_monitor_reply = _msg.get("content") or ""
if _finish == "tool_calls":
_tcs = _msg.get("tool_calls") or []
_calls_text = "; ".join(
f"{(tc.get('function') or {}).get('name', '')}"
f"({(tc.get('function') or {}).get('arguments', '')})"
for tc in _tcs
)
_monitor_reply = (_msg.get("content") or "") + (
f"[tool_calls] {_calls_text}" if _calls_text else ""
)
api_monitor.set_reply(monitor_id, _monitor_reply)
_stats = stats_holder.get("stats")
if _stats:
_monitor_usage(monitor_id, _stats.get("usage"))
@ -11189,6 +11439,55 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
return messages
def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]:
"""Flatten OpenAI content-part lists to plain strings.
Local text templates take string content and raise on part lists (e.g. a
remote ``image_url`` that leaves ``image is None``): keep the text parts,
drop the rest, like the plain non-GGUF path. GGUF keeps the parts."""
out = []
for msg in messages:
content = msg.get("content")
if isinstance(content, list):
text_parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
msg = {**msg, "content": "\n".join(text_parts) if text_parts else ""}
out.append(msg)
return out
def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]:
"""Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to
mappings for local templating.
Clients send prior-turn arguments as JSON strings, but local templates take
mappings (some raise on strings). Only the internal messages copy is
rewritten; the HTTP response stays OpenAI-shaped and unparseable strings
are left untouched."""
out = []
for msg in messages:
tool_calls = msg.get("tool_calls")
if isinstance(tool_calls, list) and tool_calls:
new_calls = []
for tc in tool_calls:
fn = tc.get("function") if isinstance(tc, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except ValueError:
parsed = None
if isinstance(parsed, dict):
tc = {**tc, "function": {**fn, "arguments": parsed}}
new_calls.append(tc)
msg = {**msg, "tool_calls": new_calls}
out.append(msg)
return out
def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]:
"""Build llama-server messages for the standard GGUF chat path.

786
studio/backend/tests/test_sf_client_tools_passthrough.py Normal file
View file

@ -0,0 +1,786 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Client-tools passthrough healing for the safetensors/MLX backend.
Parity for #6801: when a NON-GGUF model is loaded and the request declares its
own ``tools`` with server-side tools OFF, text-form tool calls are promoted back
into structured ``tool_calls`` (declared tools only) via the shared healer. MLX
rides the same orchestrator path, so a single scripted backend covers both.
"""
import asyncio
import json
from types import SimpleNamespace
from models.inference import ChatCompletionRequest, ChatMessage
from routes.inference import openai_chat_completions
from core.inference.api_monitor import ApiMonitor
LOOKUP_TOOL = {
"type": "function",
"function": {
"name": "lookup",
"description": "Look something up",
"parameters": {
"type": "object",
"properties": {"q": {"type": "string"}},
"required": ["q"],
},
},
}
SEARCH_TOOL = {
"type": "function",
"function": {
"name": "search",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
}
_CALL_XML = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}</tool_call>'
_SEARCH_XML = '<tool_call>{"name": "search", "arguments": {"query": "dogs"}}</tool_call>'
class _Request:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/chat/completions")
method = "POST"
scope: dict = {}
async def is_disconnected(self):
return False
class _ScriptedBackend:
"""Non-GGUF backend: ``generate_chat_response`` replays scripted
CUMULATIVE snapshots. ``responder(messages, tools)`` returns the snapshot
list for one generation, so nudge tests can vary output across turns."""
active_model_name = "sf-model"
def __init__(
self,
responder,
*,
stats = None,
):
self.models = {
"sf-model": {
"chat_template_info": {"template": "<tool_call> chatml"},
"context_length": 2048,
}
}
self._responder = responder
self._stats = stats
self.calls: list = []
self.reset_count = 0
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
snapshots = self._responder(messages, tools)
if stats_holder is not None and self._stats is not None:
stats_holder["stats"] = self._stats
for snap in snapshots:
yield snap
def reset_generation_state(self):
self.reset_count += 1
def _fixed(*snapshots):
"""Responder that always replays the given cumulative snapshots."""
return lambda messages, tools: list(snapshots)
def _llama_stub():
return SimpleNamespace(
is_loaded = False,
supports_tools = False,
is_vision = False,
context_length = None,
)
def _install(
monkeypatch,
backend,
*,
supports_tools = True,
):
import routes.inference as inf
from state.tool_policy import reset_tool_policy
reset_tool_policy()
monitor = ApiMonitor(max_entries = 8)
monkeypatch.setattr(inf, "api_monitor", monitor)
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _llama_stub())
monkeypatch.setattr(inf, "get_inference_backend", lambda: backend)
monkeypatch.setattr(
inf,
"_detect_safetensors_features",
lambda *a, **k: {"supports_tools": supports_tools},
)
return monitor
def _request(**kwargs):
base = dict(model = "default", messages = [ChatMessage(role = "user", content = "hi")])
base.update(kwargs)
return ChatCompletionRequest(**base)
def _call(payload, monkeypatch, backend, **install_kwargs):
_install(monkeypatch, backend, **install_kwargs)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
return asyncio.run(_run())
def _json_body(response):
return json.loads(response.body if hasattr(response, "body") else response.content)
def _collect_sse(response):
async def _run():
return [c async for c in response.body_iterator]
return asyncio.run(_run())
def _sse_objects(chunks):
out = []
for chunk in chunks:
if isinstance(chunk, bytes):
chunk = chunk.decode()
for line in str(chunk).splitlines():
if line.startswith("data: "):
data = line.removeprefix("data: ")
if data != "[DONE]":
out.append(json.loads(data))
return out
# ── Non-streaming ─────────────────────────────────────────────────
def test_xml_healed_to_tool_calls_non_streaming(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["content"] is None
calls = choice["message"]["tool_calls"]
assert len(calls) == 1
assert calls[0]["function"]["name"] == "lookup"
assert json.loads(calls[0]["function"]["arguments"]) == {"q": "cats"}
# The client tools reached the generator (template injection).
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]
def test_undeclared_call_stays_text(monkeypatch):
xml = '<tool_call>{"name": "other", "arguments": {}}</tool_call>'
backend = _ScriptedBackend(_fixed(xml))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == xml
def test_opt_out_relays_verbatim(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False, auto_heal_tool_calls = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == _CALL_XML
def test_env_kill_switch_relays_verbatim(monkeypatch):
import core.inference.passthrough_healing as ph
monkeypatch.setattr(ph, "_HEALING_DISABLED", True)
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
assert choice["message"]["content"] == _CALL_XML
def test_no_tools_request_untouched(monkeypatch):
backend = _ScriptedBackend(_fixed("just a plain answer"))
payload = _request(stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
# No tools and no tool messages -> plain path, normal ChatCompletion.
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] == "just a plain answer"
assert choice["message"].get("tool_calls") is None
def test_prose_around_call_retained(monkeypatch):
text = "Let me look:\n" + _CALL_XML + "\ndone"
backend = _ScriptedBackend(_fixed(text))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["content"] == "Let me look:\n\ndone"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
def test_empty_output_is_valid_stop(monkeypatch):
backend = _ScriptedBackend(_fixed(""))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] in ("", None)
assert choice["message"].get("tool_calls") is None
def test_tool_role_follow_up_turn_preserves_history(monkeypatch):
backend = _ScriptedBackend(_fixed("The weather is sunny."))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "weather?"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "The weather is sunny."
# The tool history reached the generator intact (role=tool + assistant.tool_calls).
sent = backend.calls[0]["messages"]
roles = [m["role"] for m in sent]
assert "tool" in roles
assistant = next(m for m in sent if m["role"] == "assistant")
assert assistant.get("tool_calls")
def test_dict_arguments_history_does_not_crash(monkeypatch):
# Non-spec client: assistant tool_calls[].function.arguments as a dict.
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "hi"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": {"q": "x"}},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "ok"
def test_forced_tool_choice_narrows_promotion(monkeypatch):
# tool_choice forces `search`; a `lookup` text call must NOT promote.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(
tools = [LOOKUP_TOOL, SEARCH_TOOL],
stream = False,
tool_choice = {"type": "function", "function": {"name": "search"}},
)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"].get("tool_calls") is None
def test_parallel_cap_non_streaming(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = False, parallel_tool_calls = False)
body = _json_body(_call(payload, monkeypatch, backend))
calls = body["choices"][0]["message"]["tool_calls"]
assert len(calls) == 1
assert calls[0]["function"]["name"] == "lookup"
def test_usage_recorded_when_stats_present(monkeypatch):
stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
backend = _ScriptedBackend(_fixed(_CALL_XML), stats = stats)
payload = _request(tools = [LOOKUP_TOOL], stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
[entry] = monitor.snapshot()
assert entry["prompt_tokens"] == 7
assert entry["completion_tokens"] == 3
# ── Nudge ─────────────────────────────────────────────────────────
def test_nudge_default_off_single_generation(monkeypatch):
# Signal present but unparseable; without opt-in, no retry.
truncated = '<tool_call>{"name": "lookup"'
backend = _ScriptedBackend(_fixed(truncated))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
_call(payload, monkeypatch, backend)
assert len(backend.calls) == 1
def test_nudge_opt_in_retry_recovers(monkeypatch):
truncated = '<tool_call>{"name": "lookup"'
def responder(messages, tools):
nudged = any(
"native tool-call format" in (m.get("content") or "")
for m in messages
if m.get("role") == "user"
)
return [_CALL_XML] if nudged else [truncated]
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
body = _json_body(_call(payload, monkeypatch, backend))
assert len(backend.calls) == 2
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
def test_nudge_double_failure_relays_original(monkeypatch):
truncated = '<tool_call>{"name": "lookup"'
backend = _ScriptedBackend(_fixed(truncated))
payload = _request(tools = [LOOKUP_TOOL], stream = False, nudge_tool_calls = True)
body = _json_body(_call(payload, monkeypatch, backend))
assert len(backend.calls) == 2 # exactly one retry
choice = body["choices"][0]
assert choice["finish_reason"] == "stop"
assert choice["message"]["content"] == truncated
# ── Streaming ─────────────────────────────────────────────────────
def test_streaming_heals_split_call_into_one_delta(monkeypatch):
# Cumulative snapshots that build the call across many increments.
pieces = ["<tool", '<tool_call>{"name": "loo', '<tool_call>{"name": "lookup", "argum']
cumulative = pieces + [_CALL_XML]
backend = _ScriptedBackend(_fixed(*cumulative))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
assert tool_deltas[0]["function"]["name"] == "lookup"
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert finishes == ["tool_calls"]
def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch):
# A stream cancelled via the registry ("Stop") must NOT promote the
# buffered-but-unclosed tool markup at finalize, else it executes a tool
# the user just cancelled. Guarded on cancel_event at the finalize step.
import routes.inference as inf
cancel_id = "cancel-me-6870"
# Balanced JSON but no closing </tool_call> -> healer HOLDS it until finalize.
held = '<tool_call>{"name": "lookup", "arguments": {"q": "cats"}}'
class _CancelMidStream(_ScriptedBackend):
def __init__(self):
super().__init__(_fixed(held))
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
yield held # healer holds the unclosed call
inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF
backend = _CancelMidStream()
payload = _request(tools = [LOOKUP_TOOL], stream = True, cancel_id = cancel_id)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert tool_deltas == [] # no tool promoted after cancel
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert "tool_calls" not in finishes # ends with finish_reason=stop, not tool_calls
def test_streaming_no_tools_verbatim(monkeypatch):
backend = _ScriptedBackend(_fixed("hello ", "hello world"))
payload = _request(stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
text = "".join(
(o["choices"][0]["delta"].get("content") or "")
for o in objs
if o["choices"] and "delta" in o["choices"][0]
)
assert text == "hello world"
finishes = [
o["choices"][0]["finish_reason"]
for o in objs
if o["choices"] and o["choices"][0].get("finish_reason")
]
assert finishes == ["stop"]
def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch):
# Repeated then shrunk cumulative snapshots must not double-heal.
backend = _ScriptedBackend(_fixed(_CALL_XML, _CALL_XML, _CALL_XML[:5], _CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
def test_streaming_parallel_cap(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML + _SEARCH_XML))
payload = _request(tools = [LOOKUP_TOOL, SEARCH_TOOL], stream = True, parallel_tool_calls = False)
response = _call(payload, monkeypatch, backend)
objs = _sse_objects(_collect_sse(response))
tool_deltas = [
tc
for o in objs
for tc in (o.get("choices", [{}])[0].get("delta", {}) or {}).get("tool_calls", []) or []
]
assert len(tool_deltas) == 1
assert tool_deltas[0]["function"]["name"] == "lookup"
def test_streaming_generator_error_closes_cleanly(monkeypatch):
def responder(messages, tools):
raise RuntimeError("boom /secret/path")
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], stream = True)
response = _call(payload, monkeypatch, backend)
chunks = _collect_sse(response)
joined = "".join(c.decode() if isinstance(c, bytes) else c for c in chunks)
assert "An internal error occurred" in joined
assert "secret/path" not in joined # CWE-209: no path leak
assert backend.reset_count >= 1
def test_streaming_disconnect_resets_once(monkeypatch):
class _DisconnectRequest(_Request):
async def is_disconnected(self):
return True
backend = _ScriptedBackend(_fixed("a", "ab", "abc"))
payload = _request(tools = [LOOKUP_TOOL], stream = True)
_install(monkeypatch, backend)
async def _run():
resp = await openai_chat_completions(
payload, request = _DisconnectRequest(), current_subject = "u"
)
return [c async for c in resp.body_iterator]
asyncio.run(_run())
assert backend.reset_count == 1
def test_mlx_uses_same_path(monkeypatch):
# MLX and safetensors share get_inference_backend(); one scripted backend covers both.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["finish_reason"] == "tool_calls"
def test_tool_choice_none_does_not_advertise_tools(monkeypatch):
# tool_choice="none": no tools rendered into the template; history templating still applies.
backend = _ScriptedBackend(_fixed("plain answer"))
payload = _request(tools = [LOOKUP_TOOL], tool_choice = "none", stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "plain answer"
assert backend.calls[0]["tools"] is None
def test_developer_message_folded_into_system_prompt(monkeypatch):
# The "developer" role folds into one leading system message (local templates reject it).
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
messages = [
ChatMessage(role = "developer", content = "always be terse"),
ChatMessage(role = "user", content = "hi"),
],
tools = [LOOKUP_TOOL],
stream = False,
)
_call(payload, monkeypatch, backend)
sent = backend.calls[0]["messages"]
assert sent[0]["role"] == "system"
assert "always be terse" in sent[0]["content"]
assert all(m.get("role") != "developer" for m in sent)
def test_failed_nudge_retry_keeps_original_response(monkeypatch):
# A raising retry must not 500; the first response is returned.
state = {"n": 0}
def responder(messages, tools):
state["n"] += 1
if state["n"] == 1:
return ['<tool_call>{"name":"lookup"'] # unhealable signal
raise RuntimeError("retry blew up")
backend = _ScriptedBackend(responder)
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
body = _json_body(_call(payload, monkeypatch, backend))
assert state["n"] == 2
assert body["choices"][0]["finish_reason"] == "stop"
assert body["choices"][0]["message"]["content"] == '<tool_call>{"name":"lookup"'
def test_discarded_nudge_retry_reports_first_attempt_usage(monkeypatch):
# Double-failure nudge: the first response is delivered, but the retry's
# generate() overwrites stats_holder. The monitor must record the FIRST
# attempt's usage, not the discarded retry's.
first_stats = {"usage": {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}}
retry_stats = {"usage": {"prompt_tokens": 99, "completion_tokens": 99, "total_tokens": 198}}
class _PerCallStatsBackend(_ScriptedBackend):
def __init__(self):
# Unhealable truncated markup on both attempts -> retry is discarded.
super().__init__(lambda m, t: ['<tool_call>{"name":"lookup"'])
self._stats_seq = [first_stats, retry_stats]
def generate_chat_response(
self,
*,
messages,
tools = None,
stats_holder = None,
**kwargs,
):
self.calls.append({"messages": messages, "tools": tools, **kwargs})
stats = self._stats_seq[min(len(self.calls) - 1, len(self._stats_seq) - 1)]
if stats_holder is not None:
stats_holder["stats"] = stats
for snap in self._responder(messages, tools):
yield snap
backend = _PerCallStatsBackend()
payload = _request(tools = [LOOKUP_TOOL], nudge_tool_calls = True, stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
assert len(backend.calls) == 2 # first attempt + one discarded retry
[entry] = monitor.snapshot()
# The delivered response is the first attempt, so its usage must be reported.
assert entry["prompt_tokens"] == 7
assert entry["completion_tokens"] == 3
def test_monitor_records_healed_call_not_raw_xml(monkeypatch):
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
asyncio.run(_run())
snap = monitor.snapshot(include_details = True)
replies = json.dumps(snap)
assert "<tool_call>" not in replies
assert "lookup" in replies
def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch):
# Monitor mirrors what the client received, never the healed-away raw markup.
backend = _ScriptedBackend(
_fixed("Sure. ", 'Sure. <tool_call>{"name": "loo', "Sure. " + _CALL_XML)
)
payload = _request(tools = [LOOKUP_TOOL], stream = True)
monitor = _install(monkeypatch, backend)
async def _run():
return await openai_chat_completions(payload, request = _Request(), current_subject = "u")
response = asyncio.run(_run())
_collect_sse(response)
replies = json.dumps(monitor.snapshot(include_details = True))
assert "<tool_call>" not in replies
assert "Sure. " in replies
assert "[tool_calls] lookup(" in replies
def test_forced_tool_choice_narrows_templated_tools(monkeypatch):
# A forced function is the only schema rendered into the template.
backend = _ScriptedBackend(_fixed(_SEARCH_XML))
payload = _request(
tools = [LOOKUP_TOOL, SEARCH_TOOL],
stream = False,
tool_choice = {"type": "function", "function": {"name": "search"}},
)
body = _json_body(_call(payload, monkeypatch, backend))
templated = backend.calls[0]["tools"]
assert [t["function"]["name"] for t in templated] == ["search"]
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "search"
def test_multimodal_content_parts_flattened_for_local_template(monkeypatch):
# Remote image URLs leave image=None, so content arrives as a part LIST:
# text parts are kept, the image part dropped.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(
messages = [
ChatMessage(
role = "user",
content = [
{"type": "text", "text": "what is this?"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/cat.png"},
},
],
)
],
tools = [LOOKUP_TOOL],
stream = False,
)
body = _json_body(_call(payload, monkeypatch, backend))
templated = backend.calls[0]["messages"]
assert all(isinstance(m.get("content"), str) for m in templated)
assert any(m["content"] == "what is this?" for m in templated)
assert body["choices"][0]["finish_reason"] == "tool_calls"
def test_string_arguments_history_deserialized_for_template(monkeypatch):
# JSON-string tool_calls arguments become dicts in the templated copy;
# the HTTP response stays OpenAI-shaped.
backend = _ScriptedBackend(_fixed("done"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "weather?"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": '{"q": "weather"}'},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "sunny"),
],
)
_json_body(_call(payload, monkeypatch, backend))
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
assert assistant["tool_calls"][0]["function"]["arguments"] == {"q": "weather"}
def test_unparseable_arguments_string_left_untouched(monkeypatch):
backend = _ScriptedBackend(_fixed("ok"))
payload = _request(
tools = [LOOKUP_TOOL],
stream = False,
messages = [
ChatMessage(role = "user", content = "hi"),
ChatMessage(
role = "assistant",
content = None,
tool_calls = [
{
"id": "call_0",
"type": "function",
"function": {"name": "lookup", "arguments": "not json {"},
}
],
),
ChatMessage(role = "tool", tool_call_id = "call_0", content = "y"),
],
)
body = _json_body(_call(payload, monkeypatch, backend))
assert body["choices"][0]["message"]["content"] == "ok"
assistant = next(m for m in backend.calls[0]["messages"] if m["role"] == "assistant")
assert assistant["tool_calls"][0]["function"]["arguments"] == "not json {"
def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch):
# mcp_enabled=true with an empty registry must not silently drop the
# declared tools; the gate keys on the server-side path claiming the request.
backend = _ScriptedBackend(_fixed(_CALL_XML))
payload = _request(tools = [LOOKUP_TOOL], stream = False, mcp_enabled = True)
body = _json_body(_call(payload, monkeypatch, backend))
choice = body["choices"][0]
assert choice["finish_reason"] == "tool_calls"
assert choice["message"]["tool_calls"][0]["function"]["name"] == "lookup"
assert backend.calls[0]["tools"] == [LOOKUP_TOOL]