From 5d87168478611e1641fb547a87194a6ce389d340 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 09:26:13 +0000 Subject: [PATCH 01/14] 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). --- studio/backend/models/inference.py | 4 +- studio/backend/routes/inference.py | 160 +++++- .../tests/test_sf_client_tools_passthrough.py | 513 ++++++++++++++++++ 3 files changed, 668 insertions(+), 9 deletions(-) create mode 100644 studio/backend/tests/test_sf_client_tools_passthrough.py diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 31c100dbec1..1041531d321 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1146,7 +1146,9 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - content: str + # ``None`` on a pure tool-call turn (OpenAI sets content=null when the + # assistant only returns tool_calls); a string on every text answer. + content: Optional[str] = None refusal: Optional[str] = None reasoning_content: Optional[str] = None tool_calls: Optional[list[dict]] = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 17be222d936..2c69a58d2b5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -614,6 +614,51 @@ 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: + """A delta chunk carrying one or more OpenAI tool-call deltas -- the sibling + of ``_chat_content_chunk`` for the in-process streamers that promote + text-form calls (see ``_sf_heal_events_to_sse``).""" + 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): + """Serialize ``StreamToolCallHealer`` events into chat-stream SSE lines, + reusing the shared chunk builders. ``state["idx"]`` tracks the healed + tool-call index across ``feed``/``finalize`` calls; ``parallel_tool_calls is + False`` caps promotion to a single call (mirrors the GGUF passthrough).""" + lines = [] + for kind, value in events: + if kind == "text": + if value: + lines.append(_chat_content_chunk(completion_id, created, model_name, 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"], + } + ], + ) + ) + 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 @@ -6829,25 +6874,67 @@ def _drain_to_text(): if payload.preserve_thinking is not None: gen_kwargs["preserve_thinking"] = payload.preserve_thinking + # ── Client-tool passthrough (safetensors + MLX) ────────────── + # A request that declares client tools (or carries tool-result history) but + # does NOT enable server-side tools is the local-model equivalent of the + # GGUF client-tool passthrough: generate one turn with the tools rendered + # into the template, then repair text-form calls with the shared #6801 + # healer. Vision / gpt-oss / server-side-tool / MCP turns are handled + # elsewhere, and a template that cannot parse the emitted markup + # (supports_tools False) falls through to plain relay with tools dropped, + # matching the GGUF gate. MLX rides this same path (the route only talks to + # the orchestrator; MLX runs inside the worker). + _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 + 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 assistant.tool_calls and role="tool" + # history survive templating (_extract_content_parts drops them). The + # system message rides inside the list, so clear the separate prompt. + gen_kwargs["messages"] = _openai_messages_for_passthrough(payload) + gen_kwargs["system_prompt"] = "" + 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 ──────────────────────────────────────── @@ -6863,6 +6950,11 @@ async def stream_chunks(): try: yield _chat_role_chunk(completion_id, created, model_name) + # Client-tool passthrough: repair text-form calls on the fly with + # the shared #6801 healer (None => relay verbatim, as before). + healer = StreamToolCallHealer(_sf_heal, payload.tools) if _sf_heal else None + heal_state = {"idx": 0} + prev_text = "" # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive @@ -6893,9 +6985,32 @@ async def stream_chunks(): if not new_text: continue api_monitor.append_reply(monitor_id, new_text) - yield _chat_content_chunk(completion_id, created, model_name, new_text) + if healer is None: + yield _chat_content_chunk(completion_id, created, model_name, new_text) + else: + for line in _sf_heal_events_to_sse( + healer.feed(new_text), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + ): + yield line - 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(), + completion_id, + created, + model_name, + heal_state, + payload.parallel_tool_calls, + ): + yield line + + _finish = "tool_calls" if (healer is not None 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 @@ -6956,14 +7071,43 @@ async def stream_chunks(): for token in generate(): full_text = token + # Client-tool passthrough: promote text-form calls, and (opt-in) a + # single nudge retry when the model tried to call a tool but emitted + # unparseable markup. + _msg = {"role": "assistant", "content": full_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": full_text}}]} + if nudge_should_retry(_data, _sf_heal, payload.tools): + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + retry_msg = {"role": "assistant", "content": retry_text} + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + full_text, _msg, _finish = retry_text, retry_msg, "tool_calls" + # Honor parallel_tool_calls=false (best-effort) by capping to one + # call, matching the GGUF passthrough (covers the nudge retry too). + 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(content = full_text), - finish_reason = "stop", + message = CompletionMessage( + content = _msg["content"], + tool_calls = _msg.get("tool_calls"), + ), + finish_reason = _finish, ) ], ) diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py new file mode 100644 index 00000000000..332851e4e59 --- /dev/null +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -0,0 +1,513 @@ +# 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 = '{"name": "lookup", "arguments": {"q": "cats"}}' +_SEARCH_XML = '{"name": "search", "arguments": {"query": "dogs"}}' + + +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 whose ``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": " 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 = '{"name": "other", "arguments": {}}' + 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 declared and no tool messages -> the branch is skipped and the + # plain path returns a 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 = '{"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 = '{"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 = '{"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 = ["{"name": "loo', '{"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_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): + # Same cumulative snapshot twice, then a shrunk one, must not double-heal + # or negative-slice. + 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 both dispatch through get_inference_backend(); the same + # scripted backend + branch cover both. A healed call proves the shared path. + 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" From f9be9e08684378978971e8831ae429337d33a17c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 11:12:40 +0000 Subject: [PATCH 02/14] 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. --- studio/backend/routes/inference.py | 53 +++++++++++---- .../tests/test_sf_client_tools_passthrough.py | 65 +++++++++++++++++++ 2 files changed, 105 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2c69a58d2b5..3fedcb1fd4e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6900,11 +6900,21 @@ def _drain_to_text(): ) if _sf_client_tools: # Re-derive from payload.messages so assistant.tool_calls and role="tool" - # history survive templating (_extract_content_parts drops them). The - # system message rides inside the list, so clear the separate prompt. - gen_kwargs["messages"] = _openai_messages_for_passthrough(payload) + # history survive templating (_extract_content_parts drops them). Fold + # system/developer turns into one leading system message (the extracted + # system_prompt already collects both) -- local templates reject the + # OpenAI "developer" role and the fallback formatter drops it -- then + # clear the separate prompt so the worker does not prepend a duplicate. + gen_kwargs["messages"] = _set_or_prepend_system_message( + _openai_messages_for_passthrough(payload), system_prompt + ) gen_kwargs["system_prompt"] = "" - gen_kwargs["tools"] = payload.tools + # tool_choice="none" forces a final-answer turn: keep the tool-history + # templating but do not advertise the tools, otherwise the model is + # 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 # Request-scoped usage/timings receptacle (filled at gen_done). stats_holder: dict = {} @@ -7082,14 +7092,20 @@ async def stream_chunks(): elif nudge_enabled(payload.nudge_tool_calls): _data = {"choices": [{"message": {"role": "assistant", "content": full_text}}]} if nudge_should_retry(_data, _sf_heal, payload.tools): - retry_text = "" - for token in generate( - [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] - ): - retry_text = token - retry_msg = {"role": "assistant", "content": retry_text} - if heal_openai_message(retry_msg, _sf_heal, payload.tools): - full_text, _msg, _finish = retry_text, retry_msg, "tool_calls" + # The original answer is already in hand: a retry that + # fails or is cancelled must not turn the request into a + # 500 (the GGUF nudge path keeps the first response too). + try: + retry_text = "" + for token in generate( + [*gen_kwargs["messages"], *nudge_messages(_data, _sf_heal)] + ): + retry_text = token + retry_msg = {"role": "assistant", "content": retry_text} + if heal_openai_message(retry_msg, _sf_heal, payload.tools): + full_text, _msg, _finish = retry_text, retry_msg, "tool_calls" + except Exception as retry_exc: + logger.debug("Nudge retry failed; keeping first response: %s", retry_exc) # Honor parallel_tool_calls=false (best-effort) by capping to one # call, matching the GGUF passthrough (covers the nudge retry too). if payload.parallel_tool_calls is False: @@ -7111,7 +7127,18 @@ async def stream_chunks(): ) ], ) - api_monitor.set_reply(monitor_id, full_text) + _monitor_reply = full_text + 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")) diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 332851e4e59..d007c7547b5 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -511,3 +511,68 @@ def test_mlx_uses_same_path(monkeypatch): 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" forces a final-answer turn: the template must NOT be + # prompted with the tools (heal_gate is off, so any emitted markup would + # relay as prose). 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 OpenAI "developer" role must not reach local templating verbatim + # (templates reject it / the fallback formatter drops it); it folds into a + # single leading system message. + 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 retry that raises after the original answer exists must not become a + # 500; the first response is returned (GGUF nudge parity). + state = {"n": 0} + + def responder(messages, tools): + state["n"] += 1 + if state["n"] == 1: + return ['{"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"] == '{"name":"lookup"' + + +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 "" not in replies + assert "lookup" in replies From ab2dec5582b170e8cc76688ab8c839e8fde04e6e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 11:13:15 +0000 Subject: [PATCH 03/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3fedcb1fd4e..eae92664db8 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7105,7 +7105,9 @@ async def stream_chunks(): if heal_openai_message(retry_msg, _sf_heal, payload.tools): full_text, _msg, _finish = retry_text, retry_msg, "tool_calls" except Exception as retry_exc: - logger.debug("Nudge retry failed; keeping first response: %s", retry_exc) + logger.debug( + "Nudge retry failed; keeping first response: %s", retry_exc + ) # Honor parallel_tool_calls=false (best-effort) by capping to one # call, matching the GGUF passthrough (covers the nudge retry too). if payload.parallel_tool_calls is False: From 1f7c065e0d5ff1e5d445debea17d6e056ef6238e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 11:40:50 +0000 Subject: [PATCH 04/14] 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. --- studio/backend/routes/inference.py | 87 +++++++++++++++++-- .../tests/test_sf_client_tools_passthrough.py | 68 +++++++++++++++ 2 files changed, 148 insertions(+), 7 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index eae92664db8..2006dcc9b30 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -627,16 +627,31 @@ def _chat_tool_calls_chunk(completion_id, created, model_name, tool_calls) -> st ) -def _sf_heal_events_to_sse(events, completion_id, created, model_name, state, parallel_tool_calls): +def _sf_heal_events_to_sse( + events, + completion_id, + created, + model_name, + state, + parallel_tool_calls, + monitor_id = None, +): """Serialize ``StreamToolCallHealer`` events into chat-stream SSE lines, reusing the shared chunk builders. ``state["idx"]`` tracks the healed tool-call index across ``feed``/``finalize`` calls; ``parallel_tool_calls is - False`` caps promotion to a single call (mirrors the GGUF passthrough).""" + False`` caps promotion to a single call (mirrors the GGUF passthrough). + + When ``monitor_id`` is given, the monitor entry is fed from the SAME events + the client receives -- relayed text verbatim, promoted calls as the + ``[tool_calls] name(args)`` summary the non-streaming path records -- so + /api/monitor never shows healed-away markup (or capped-off calls) that were + not actually returned.""" 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 @@ -655,6 +670,12 @@ def _sf_heal_events_to_sse(events, completion_id, created, model_name, state, pa ], ) ) + _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 @@ -6906,15 +6927,36 @@ def _drain_to_text(): # OpenAI "developer" role and the fallback formatter drops it -- then # clear the separate prompt so the worker does not prepend a duplicate. gen_kwargs["messages"] = _set_or_prepend_system_message( - _openai_messages_for_passthrough(payload), system_prompt + _flatten_content_parts_for_local_template( + _openai_messages_for_passthrough(payload) + ), + system_prompt, ) gen_kwargs["system_prompt"] = "" # tool_choice="none" forces a final-answer turn: keep the tool-history # templating but do not advertise the tools, otherwise the model is # 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 + # "none") then relays as ordinary content. A forced function narrows + # templating to that one schema, so the model is never prompted with + # tools the healer (correctly) refuses to promote. Both mirror the GGUF + # passthrough, where llama-server receives and 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 = {} @@ -6994,8 +7036,12 @@ async def stream_chunks(): prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) if healer is None: + # No healing: the monitor mirrors the verbatim relay. + # With healing on, _sf_heal_events_to_sse records the + # healed events instead, so the monitor never shows raw + # markup the client did not receive. + api_monitor.append_reply(monitor_id, new_text) yield _chat_content_chunk(completion_id, created, model_name, new_text) else: for line in _sf_heal_events_to_sse( @@ -7005,6 +7051,7 @@ async def stream_chunks(): model_name, heal_state, payload.parallel_tool_calls, + monitor_id, ): yield line @@ -7016,6 +7063,7 @@ async def stream_chunks(): model_name, heal_state, payload.parallel_tool_calls, + monitor_id, ): yield line @@ -10933,6 +10981,31 @@ 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 for the + safetensors/MLX client-tools passthrough. + + Local text templates take string content; forwarding a part list (e.g. a + remote ``image_url`` the local backends cannot fetch -- + ``_extract_content_parts`` only decodes ``data:`` URLs, so such requests + reach this path with ``image is None``) raises inside + ``apply_chat_template``. Mirror the flattening the plain non-GGUF path + performs: keep the text parts, drop the rest. The GGUF passthrough keeps + the original parts (llama-server accepts them).""" + 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 _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict], bool]: """Build llama-server messages for the standard GGUF chat path. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index d007c7547b5..c8614078f34 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -576,3 +576,71 @@ async def _run(): replies = json.dumps(snap) assert "" not in replies assert "lookup" in replies + + +def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): + # The monitor mirrors what the client received: relayed prose verbatim, + # promoted calls as the [tool_calls] summary -- never the raw markup the + # healer consumed (parity with the non-streaming set_reply rewrite). + backend = _ScriptedBackend( + _fixed("Sure. ", 'Sure. {"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 "" 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 local template, so + # the model is never prompted with tools the healer refuses to promote + # (llama-server enforces tool_choice itself on the GGUF path). + 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 (only data: URLs are decoded), so the + # request reaches this path with a content-part LIST. Local text templates + # take string content: keep the text parts, drop the image part, exactly + # like the plain non-GGUF path flattens in _extract_content_parts. + 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" From 32070dc350ea272cb067e1778b61220ae06f5227 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 11:41:36 +0000 Subject: [PATCH 05/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2006dcc9b30..fd07b319683 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6927,9 +6927,7 @@ def _drain_to_text(): # OpenAI "developer" role and the fallback formatter drops it -- then # clear the separate prompt so the worker does not prepend a duplicate. gen_kwargs["messages"] = _set_or_prepend_system_message( - _flatten_content_parts_for_local_template( - _openai_messages_for_passthrough(payload) - ), + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)), system_prompt, ) gen_kwargs["system_prompt"] = "" From a5ef36576784e3ce675b6032f496c650c600c62d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 12:52:41 +0000 Subject: [PATCH 06/14] 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. --- studio/backend/routes/inference.py | 47 ++++++++++++- .../tests/test_sf_client_tools_passthrough.py | 69 +++++++++++++++++++ 2 files changed, 114 insertions(+), 2 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fd07b319683..7d3639c7870 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6906,9 +6906,16 @@ def _drain_to_text(): # matching the GGUF gate. MLX rides this same path (the route only talks to # the orchestrator; MLX runs inside the worker). _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) + # Gate on whether the server-side tool path ACTUALLY claimed the request + # (_sf_use_tools is final by here and its branch always returns), not on + # the raw mcp_enabled flag: a client that sets mcp_enabled globally while + # declaring its own tools must still get the passthrough when no MCP tool + # survived (empty registry, CLI --disable-tools policy) instead of the + # silent tool-drop this branch exists to fix. The GGUF passthrough gate + # has no mcp_enabled clause either. _sf_client_tools = ( not _effective_enable_tools(payload) - and not payload.mcp_enabled + and not _sf_use_tools and image is None and not _sf_is_gptoss and _sf_features.get("supports_tools", False) @@ -6927,7 +6934,11 @@ def _drain_to_text(): # OpenAI "developer" role and the fallback formatter drops it -- then # clear the separate prompt so the worker does not prepend a duplicate. gen_kwargs["messages"] = _set_or_prepend_system_message( - _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)), + _structured_tool_history_for_local_template( + _flatten_content_parts_for_local_template( + _openai_messages_for_passthrough(payload) + ) + ), system_prompt, ) gen_kwargs["system_prompt"] = "" @@ -11004,6 +11015,38 @@ def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict 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 safetensors/MLX templating. + + Standard OpenAI clients send prior-turn arguments as JSON strings (the same + shape this endpoint returns), but local chat templates take mappings -- + ``arguments | items`` iteration, or an explicit raise on strings -- so the + string shape crashes or misrenders the templated tool history. Only the + internal ``gen_kwargs["messages"]`` copy is rewritten; the HTTP response + stays OpenAI-shaped, and the GGUF passthrough keeps strings (llama-server + expects the wire shape). 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. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index c8614078f34..6920b41a508 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -644,3 +644,72 @@ def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): 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): + # Spec-compliant clients send prior tool_calls arguments as JSON strings; + # local templates take mappings (some raise on strings), so the templated + # copy carries dicts while 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 MCP registry must not fall through to + # plain generation with the declared tools silently dropped: the gate keys + # on whether the server-side path actually claimed 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] From 16e5e9273a479695e1ab137ba3f5361fa6a1d76f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 12:54:32 +0000 Subject: [PATCH 07/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7d3639c7870..4acf4937c65 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6935,9 +6935,7 @@ def _drain_to_text(): # clear the separate prompt so the worker does not prepend a duplicate. 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) - ) + _flatten_content_parts_for_local_template(_openai_messages_for_passthrough(payload)) ), system_prompt, ) From 923abf4fbe1d70b078cd2240b8177b8aac5a4c6e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:18:14 +0000 Subject: [PATCH 08/14] Tighten comments and docstrings in the client-tools passthrough --- studio/backend/models/inference.py | 3 +- studio/backend/routes/inference.py | 118 +++++++----------- .../tests/test_sf_client_tools_passthrough.py | 46 +++---- 3 files changed, 58 insertions(+), 109 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 1041531d321..1e7770a7f88 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1146,8 +1146,7 @@ class CompletionMessage(BaseModel): """The assistant's complete response message.""" role: Literal["assistant"] = "assistant" - # ``None`` on a pure tool-call turn (OpenAI sets content=null when the - # assistant only returns tool_calls); a string on every text answer. + # ``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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4acf4937c65..2dc49bbf823 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -615,9 +615,7 @@ 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: - """A delta chunk carrying one or more OpenAI tool-call deltas -- the sibling - of ``_chat_content_chunk`` for the in-process streamers that promote - text-form calls (see ``_sf_heal_events_to_sse``).""" + """Delta chunk carrying OpenAI tool-call deltas (sibling of ``_chat_content_chunk``).""" return _chat_chunk_sse( completion_id, created, @@ -636,16 +634,12 @@ def _sf_heal_events_to_sse( parallel_tool_calls, monitor_id = None, ): - """Serialize ``StreamToolCallHealer`` events into chat-stream SSE lines, - reusing the shared chunk builders. ``state["idx"]`` tracks the healed - tool-call index across ``feed``/``finalize`` calls; ``parallel_tool_calls is - False`` caps promotion to a single call (mirrors the GGUF passthrough). - - When ``monitor_id`` is given, the monitor entry is fed from the SAME events - the client receives -- relayed text verbatim, promoted calls as the - ``[tool_calls] name(args)`` summary the non-streaming path records -- so - /api/monitor never shows healed-away markup (or capped-off calls) that were - not actually returned.""" + """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": @@ -6896,23 +6890,14 @@ def _drain_to_text(): gen_kwargs["preserve_thinking"] = payload.preserve_thinking # ── Client-tool passthrough (safetensors + MLX) ────────────── - # A request that declares client tools (or carries tool-result history) but - # does NOT enable server-side tools is the local-model equivalent of the - # GGUF client-tool passthrough: generate one turn with the tools rendered - # into the template, then repair text-form calls with the shared #6801 - # healer. Vision / gpt-oss / server-side-tool / MCP turns are handled - # elsewhere, and a template that cannot parse the emitted markup - # (supports_tools False) falls through to plain relay with tools dropped, - # matching the GGUF gate. MLX rides this same path (the route only talks to - # the orchestrator; MLX runs inside the worker). + # Client tools (or tool-result history) without server-side tools: render + # the tools into the template, generate one turn, repair text-form calls + # with the shared #6801 healer. supports_tools=False falls through to plain + # relay (GGUF gate parity); MLX rides the same orchestrator path. _sf_has_tool_msgs = any(m.role == "tool" or m.tool_calls for m in payload.messages) - # Gate on whether the server-side tool path ACTUALLY claimed the request - # (_sf_use_tools is final by here and its branch always returns), not on - # the raw mcp_enabled flag: a client that sets mcp_enabled globally while - # declaring its own tools must still get the passthrough when no MCP tool - # survived (empty registry, CLI --disable-tools policy) instead of the - # silent tool-drop this branch exists to fix. The GGUF passthrough gate - # has no mcp_enabled clause either. + # Gate on whether the server-side path actually claimed the request + # (_sf_use_tools), not raw mcp_enabled: with an empty MCP registry the + # declared client tools must still get the passthrough, not a silent drop. _sf_client_tools = ( not _effective_enable_tools(payload) and not _sf_use_tools @@ -6927,12 +6912,10 @@ def _drain_to_text(): else None ) if _sf_client_tools: - # Re-derive from payload.messages so assistant.tool_calls and role="tool" - # history survive templating (_extract_content_parts drops them). Fold - # system/developer turns into one leading system message (the extracted - # system_prompt already collects both) -- local templates reject the - # OpenAI "developer" role and the fallback formatter drops it -- then - # clear the separate prompt so the worker does not prepend a duplicate. + # Re-derive from payload.messages so tool_calls / role="tool" history + # survive templating; fold system/developer turns into one leading + # system message (local templates reject "developer") and clear the + # separate prompt so the worker does not prepend a duplicate. 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)) @@ -6940,14 +6923,10 @@ def _drain_to_text(): system_prompt, ) gen_kwargs["system_prompt"] = "" - # tool_choice="none" forces a final-answer turn: keep the tool-history - # templating but do not advertise the tools, otherwise the model is - # prompted to emit tool markup that heal_gate (correctly off for - # "none") then relays as ordinary content. A forced function narrows - # templating to that one schema, so the model is never prompted with - # tools the healer (correctly) refuses to promote. Both mirror the GGUF - # passthrough, where llama-server receives and honors tool_choice - # itself. + # 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): @@ -7009,8 +6988,8 @@ async def stream_chunks(): try: yield _chat_role_chunk(completion_id, created, model_name) - # Client-tool passthrough: repair text-form calls on the fly with - # the shared #6801 healer (None => relay verbatim, as before). + # 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} @@ -7044,10 +7023,8 @@ async def stream_chunks(): if not new_text: continue if healer is None: - # No healing: the monitor mirrors the verbatim relay. - # With healing on, _sf_heal_events_to_sse records the - # healed events instead, so the monitor never shows raw - # markup the client did not receive. + # Monitor mirrors the verbatim relay; with healing on, + # _sf_heal_events_to_sse records the healed events instead. api_monitor.append_reply(monitor_id, new_text) yield _chat_content_chunk(completion_id, created, model_name, new_text) else: @@ -7136,9 +7113,8 @@ async def stream_chunks(): for token in generate(): full_text = token - # Client-tool passthrough: promote text-form calls, and (opt-in) a - # single nudge retry when the model tried to call a tool but emitted - # unparseable markup. + # Client-tool passthrough: promote text-form calls; opt-in single + # nudge retry on unparseable tool markup. _msg = {"role": "assistant", "content": full_text} _finish = "stop" if _sf_heal: @@ -7147,9 +7123,8 @@ async def stream_chunks(): elif nudge_enabled(payload.nudge_tool_calls): _data = {"choices": [{"message": {"role": "assistant", "content": full_text}}]} if nudge_should_retry(_data, _sf_heal, payload.tools): - # The original answer is already in hand: a retry that - # fails or is cancelled must not turn the request into a - # 500 (the GGUF nudge path keeps the first response too). + # A failed retry must not 500 the request; keep the + # first response (GGUF nudge parity). try: retry_text = "" for token in generate( @@ -7163,8 +7138,7 @@ async def stream_chunks(): logger.debug( "Nudge retry failed; keeping first response: %s", retry_exc ) - # Honor parallel_tool_calls=false (best-effort) by capping to one - # call, matching the GGUF passthrough (covers the nudge retry too). + # 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: @@ -10989,16 +10963,11 @@ def _openai_messages_for_passthrough(payload) -> list[dict]: def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict]: - """Flatten OpenAI content-part lists to plain strings for the - safetensors/MLX client-tools passthrough. - - Local text templates take string content; forwarding a part list (e.g. a - remote ``image_url`` the local backends cannot fetch -- - ``_extract_content_parts`` only decodes ``data:`` URLs, so such requests - reach this path with ``image is None``) raises inside - ``apply_chat_template``. Mirror the flattening the plain non-GGUF path - performs: keep the text parts, drop the rest. The GGUF passthrough keeps - the original parts (llama-server accepts them).""" + """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") @@ -11015,15 +10984,12 @@ def _flatten_content_parts_for_local_template(messages: list[dict]) -> list[dict def _structured_tool_history_for_local_template(messages: list[dict]) -> list[dict]: """Deserialize assistant ``tool_calls[].function.arguments`` JSON strings to - mappings for safetensors/MLX templating. - - Standard OpenAI clients send prior-turn arguments as JSON strings (the same - shape this endpoint returns), but local chat templates take mappings -- - ``arguments | items`` iteration, or an explicit raise on strings -- so the - string shape crashes or misrenders the templated tool history. Only the - internal ``gen_kwargs["messages"]`` copy is rewritten; the HTTP response - stays OpenAI-shaped, and the GGUF passthrough keeps strings (llama-server - expects the wire shape). Unparseable strings are left untouched.""" + 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") diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 6920b41a508..5d14693dcf3 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -57,7 +57,7 @@ async def is_disconnected(self): class _ScriptedBackend: - """Non-GGUF backend whose ``generate_chat_response`` replays scripted + """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.""" @@ -230,8 +230,7 @@ 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 declared and no tool messages -> the branch is skipped and the - # plain path returns a normal ChatCompletion. + # 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" @@ -443,8 +442,7 @@ def test_streaming_no_tools_verbatim(monkeypatch): def test_streaming_repeated_snapshot_no_duplicate_call(monkeypatch): - # Same cumulative snapshot twice, then a shrunk one, must not double-heal - # or negative-slice. + # 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) @@ -505,8 +503,7 @@ async def _run(): def test_mlx_uses_same_path(monkeypatch): - # MLX and safetensors both dispatch through get_inference_backend(); the same - # scripted backend + branch cover both. A healed call proves the shared path. + # 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)) @@ -514,9 +511,7 @@ def test_mlx_uses_same_path(monkeypatch): def test_tool_choice_none_does_not_advertise_tools(monkeypatch): - # tool_choice="none" forces a final-answer turn: the template must NOT be - # prompted with the tools (heal_gate is off, so any emitted markup would - # relay as prose). History templating still applies. + # 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)) @@ -525,9 +520,7 @@ def test_tool_choice_none_does_not_advertise_tools(monkeypatch): def test_developer_message_folded_into_system_prompt(monkeypatch): - # The OpenAI "developer" role must not reach local templating verbatim - # (templates reject it / the fallback formatter drops it); it folds into a - # single leading system message. + # The "developer" role folds into one leading system message (local templates reject it). backend = _ScriptedBackend(_fixed("ok")) payload = _request( messages = [ @@ -545,8 +538,7 @@ def test_developer_message_folded_into_system_prompt(monkeypatch): def test_failed_nudge_retry_keeps_original_response(monkeypatch): - # A retry that raises after the original answer exists must not become a - # 500; the first response is returned (GGUF nudge parity). + # A raising retry must not 500; the first response is returned. state = {"n": 0} def responder(messages, tools): @@ -579,9 +571,7 @@ async def _run(): def test_streaming_monitor_records_healed_call_not_raw_xml(monkeypatch): - # The monitor mirrors what the client received: relayed prose verbatim, - # promoted calls as the [tool_calls] summary -- never the raw markup the - # healer consumed (parity with the non-streaming set_reply rewrite). + # Monitor mirrors what the client received, never the healed-away raw markup. backend = _ScriptedBackend( _fixed("Sure. ", 'Sure. {"name": "loo', "Sure. " + _CALL_XML) ) @@ -600,9 +590,7 @@ async def _run(): def test_forced_tool_choice_narrows_templated_tools(monkeypatch): - # A forced function is the only schema rendered into the local template, so - # the model is never prompted with tools the healer refuses to promote - # (llama-server enforces tool_choice itself on the GGUF path). + # A forced function is the only schema rendered into the template. backend = _ScriptedBackend(_fixed(_SEARCH_XML)) payload = _request( tools = [LOOKUP_TOOL, SEARCH_TOOL], @@ -618,10 +606,8 @@ def test_forced_tool_choice_narrows_templated_tools(monkeypatch): def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): - # Remote image URLs leave image=None (only data: URLs are decoded), so the - # request reaches this path with a content-part LIST. Local text templates - # take string content: keep the text parts, drop the image part, exactly - # like the plain non-GGUF path flattens in _extract_content_parts. + # 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 = [ @@ -647,9 +633,8 @@ def test_multimodal_content_parts_flattened_for_local_template(monkeypatch): def test_string_arguments_history_deserialized_for_template(monkeypatch): - # Spec-compliant clients send prior tool_calls arguments as JSON strings; - # local templates take mappings (some raise on strings), so the templated - # copy carries dicts while the HTTP response stays OpenAI-shaped. + # 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], @@ -703,9 +688,8 @@ def test_unparseable_arguments_string_left_untouched(monkeypatch): def test_mcp_enabled_without_server_tools_uses_passthrough(monkeypatch): - # mcp_enabled=true with an empty MCP registry must not fall through to - # plain generation with the declared tools silently dropped: the gate keys - # on whether the server-side path actually claimed the request. + # 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)) From 6369a3a7e62ba0b085dbedf1d36cb6d6c24b346d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 01:12:47 +0000 Subject: [PATCH 09/14] 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. --- studio/backend/routes/inference.py | 9 ++++ .../tests/test_sf_client_tools_passthrough.py | 43 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2dc49bbf823..1ffebc742c0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7125,6 +7125,11 @@ async def stream_chunks(): 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 preserve + # the first attempt's stats: if the retry is discarded we still + # deliver the first response and must report ITS usage, not the + # unseen retry's. + _first_stats = stats_holder.get("stats") try: retry_text = "" for token in generate( @@ -7134,10 +7139,14 @@ async def stream_chunks(): retry_msg = {"role": "assistant", "content": retry_text} if heal_openai_message(retry_msg, _sf_heal, payload.tools): full_text, _msg, _finish = retry_text, 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") diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 5d14693dcf3..64ec5f87305 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -555,6 +555,49 @@ def responder(messages, tools): assert body["choices"][0]["message"]["content"] == '{"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: ['{"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) From 57428a0eb576a5e16756371b2a7ddbe23632cd11 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jul 2026 01:55:45 +0000 Subject: [PATCH 10/14] 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. --- studio/backend/routes/inference.py | 15 +++++++- .../tests/test_sf_client_tools_passthrough.py | 37 +++++++++++++++++++ 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1ffebc742c0..db3028d6bd2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7039,7 +7039,14 @@ async def stream_chunks(): ): yield line - if healer is not None: + # A cancelled stream must not promote buffered-but-incomplete tool + # markup: finalize() heals partial markup at EOF (allow_incomplete), + # which would emit a tool_calls delta + finish_reason=tool_calls and make + # the client execute a tool the user just cancelled. The disconnect path + # already returns before here; the registry "Stop" path only sets + # cancel_event and breaks, so guard finalize/_finish on it 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, @@ -7051,7 +7058,11 @@ async def stream_chunks(): ): yield line - _finish = "tool_calls" if (healer is not None and healer.healed) else "stop" + _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. diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 64ec5f87305..07010455483 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -422,6 +422,43 @@ def test_streaming_heals_split_call_into_one_delta(monkeypatch): assert finishes == ["tool_calls"] +def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): + # A stream cancelled via the registry ("Stop") mid-emission must NOT promote the + # buffered-but-unclosed tool markup at finalize -- otherwise the client executes a + # tool the user just cancelled. Guarded on cancel_event at the finalize/_finish step. + import routes.inference as inf + + cancel_id = "cancel-me-6870" + # Balanced JSON but no closing -> healer HOLDS it until finalize. + held = '{"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) From 958eb021a14b4372af7f0a7654e96eaa67995c54 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 6 Jul 2026 01:56:58 +0000 Subject: [PATCH 11/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../backend/tests/test_sf_client_tools_passthrough.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index 07010455483..d7d7f685377 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -436,9 +436,16 @@ class _CancelMidStream(_ScriptedBackend): def __init__(self): super().__init__(_fixed(held)) - def generate_chat_response(self, *, messages, tools = None, stats_holder = None, **kwargs): + 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 + yield held # healer holds the unclosed call inf._cancel_by_cancel_id_or_stash(cancel_id) # user hits Stop before EOF backend = _CancelMidStream() From 8e3b37cf45da8eb876042a981ab16b63e5478814 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 08:51:26 +0000 Subject: [PATCH 12/14] Trim comments in the client-tools passthrough --- studio/backend/routes/inference.py | 37 +++++++++++++----------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index db3028d6bd2..2b7f2426f00 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6891,13 +6891,11 @@ def _drain_to_text(): # ── Client-tool passthrough (safetensors + MLX) ────────────── # Client tools (or tool-result history) without server-side tools: render - # the tools into the template, generate one turn, repair text-form calls - # with the shared #6801 healer. supports_tools=False falls through to plain - # relay (GGUF gate parity); MLX rides the same orchestrator path. + # 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 whether the server-side path actually claimed the request - # (_sf_use_tools), not raw mcp_enabled: with an empty MCP registry the - # declared client tools must still get the passthrough, not a silent drop. + # 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 @@ -6913,9 +6911,9 @@ def _drain_to_text(): ) if _sf_client_tools: # Re-derive from payload.messages so tool_calls / role="tool" history - # survive templating; fold system/developer turns into one leading - # system message (local templates reject "developer") and clear the - # separate prompt so the worker does not prepend a duplicate. + # survives templating; fold system/developer into one leading system + # message (templates reject "developer") and clear the separate prompt + # so the worker does not prepend a duplicate. 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)) @@ -7039,12 +7037,11 @@ async def stream_chunks(): ): yield line - # A cancelled stream must not promote buffered-but-incomplete tool - # markup: finalize() heals partial markup at EOF (allow_incomplete), - # which would emit a tool_calls delta + finish_reason=tool_calls and make - # the client execute a tool the user just cancelled. The disconnect path - # already returns before here; the registry "Stop" path only sets - # cancel_event and breaks, so guard finalize/_finish on it too. + # A cancelled stream must not promote buffered-but-incomplete + # markup: finalize()'s allow_incomplete heal would emit a + # tool_calls delta + finish_reason=tool_calls, executing a tool + # the user just cancelled. The disconnect path returns earlier; + # the "Stop" path 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( @@ -7134,12 +7131,10 @@ async def stream_chunks(): elif nudge_enabled(payload.nudge_tool_calls): _data = {"choices": [{"message": {"role": "assistant", "content": full_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 preserve - # the first attempt's stats: if the retry is discarded we still - # deliver the first response and must report ITS usage, not the - # unseen retry's. + # 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 (report ITS usage). _first_stats = stats_holder.get("stats") try: retry_text = "" From 3a098c664df40397b17c5f3c3c5c8951015e4a4a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 16:47:03 +0000 Subject: [PATCH 13/14] Trim client-tools passthrough comments further --- studio/backend/routes/inference.py | 12 +++++------- .../tests/test_sf_client_tools_passthrough.py | 6 +++--- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 2b7f2426f00..28a3907ccd2 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6912,8 +6912,7 @@ def _drain_to_text(): 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 the separate prompt - # so the worker does not prepend a duplicate. + # 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)) @@ -7038,10 +7037,9 @@ async def stream_chunks(): yield line # A cancelled stream must not promote buffered-but-incomplete - # markup: finalize()'s allow_incomplete heal would emit a - # tool_calls delta + finish_reason=tool_calls, executing a tool - # the user just cancelled. The disconnect path returns earlier; - # the "Stop" path only sets cancel_event, so guard on it here too. + # 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( @@ -7134,7 +7132,7 @@ async def stream_chunks(): # 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 (report ITS usage). + # and restore them if the retry is discarded. _first_stats = stats_holder.get("stats") try: retry_text = "" diff --git a/studio/backend/tests/test_sf_client_tools_passthrough.py b/studio/backend/tests/test_sf_client_tools_passthrough.py index d7d7f685377..01905b712cd 100644 --- a/studio/backend/tests/test_sf_client_tools_passthrough.py +++ b/studio/backend/tests/test_sf_client_tools_passthrough.py @@ -423,9 +423,9 @@ def test_streaming_heals_split_call_into_one_delta(monkeypatch): def test_streaming_cancel_does_not_finalize_tool_call(monkeypatch): - # A stream cancelled via the registry ("Stop") mid-emission must NOT promote the - # buffered-but-unclosed tool markup at finalize -- otherwise the client executes a - # tool the user just cancelled. Guarded on cancel_event at the finalize/_finish step. + # 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" From 1b3514422c958280e691497c381f10deb5fdfc61 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 7 Jul 2026 02:47:59 +0000 Subject: [PATCH 14/14] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 86f61c3c10b..d78ba422e09 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7480,7 +7480,9 @@ async def stream_chunks(): 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}}]} + _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() @@ -7504,7 +7506,11 @@ async def stream_chunks(): 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" + _visible_text, _msg, _finish = ( + _retry_visible, + retry_msg, + "tool_calls", + ) else: # Retry produced no healable call -> first response wins. stats_holder["stats"] = _first_stats