ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Improve local chat tool call flow (#5962)

Browse source 
Unify the Studio local tool-call loop (GGUF + safetensors) behind a shared ToolLoopController: ordered preface-then-tool-card rendering, duplicate-call de-looping with a forced final answer, XML-leak containment, and a parser fix that accepts closed <function=...> calls followed by trailing prose. Includes backend tests for the controller, strict parser, and GGUF route cursor reset.
This commit is contained in:
Wasim Yousef Said 2026-06-09 16:28:44 +02:00 • committed by GitHub
commit ccb471f5bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 3647 additions and 900 deletions
Download patch file Download diff file Expand all files Collapse all files

7
studio/backend/core/inference/inference.py
View file

@ -772,8 +772,11 @@ class InferenceBackend:
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
def _single_turn(conv: list):
def _single_turn(conv: list, *, active_tools: Optional[list[dict]] = None):
# conv already has the system message -- avoid double-prepend.
# `active_tools` is supplied by run_safetensors_tool_loop so one-shot
# tools such as render_html can be removed from later same-response prompts.
turn_tools = active_tools if active_tools is not None else tools
yield from self._generate_chat_response_inner(
messages = conv,
system_prompt = "",
@ -784,7 +787,7 @@ class InferenceBackend:
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
tools = turn_tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,

489
studio/backend/core/inference/llama_cpp.py
View file

@ -24,7 +24,6 @@ import threading
import time
from pathlib import Path
from typing import Generator, Iterable, List, Optional
from urllib.parse import urlparse
import httpx
@ -51,9 +50,13 @@ from utils.subprocess_compat import (
windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs,
)
from core.inference.tool_call_parser import (
RENDER_HTML_REPEAT_NUDGE,
TOOL_XML_SIGNALS,
parse_tool_calls_from_text as _shared_parse_tool_calls_from_text,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
tool_event_provenance,
)
logger = get_logger(__name__)
@ -76,7 +79,7 @@ _INTENT_SIGNAL = re.compile(
r"\b(?:now i|next i)\b"
r")"
)
_MAX_REPROMPTS = 3
_MAX_REPROMPTS = 1
# Without max_tokens, llama-server defaults n_predict = n_ctx (up to 262144 for
# Qwen3.5), causing many-minute zombie decodes when cancel fails.
@ -88,6 +91,30 @@ _MAX_REPROMPTS = 3
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_T_MAX_PREDICT_MS = 600_000 # 10 min
_REPROMPT_MAX_CHARS = 2000
_FORCED_REPEAT_PLAN_SIGNAL = re.compile(
r"\b(?:i\s+will|i'll|let\s+me|going\s+to|need\s+to|call|use|run|search|fetch|render)\b",
re.I,
)
_FINAL_ANSWER_SIGNAL = re.compile(
r"\b(?:final\s+answer|answer\s*:|here\s+is|here's|in\s+summary|result\s*:)\b",
re.I,
)
def _is_short_intent_without_action(text: str) -> bool:
stripped = text.strip()
return 0 < len(stripped) < _REPROMPT_MAX_CHARS and _INTENT_SIGNAL.search(stripped) is not None
def _should_suppress_forced_no_tool_output(text: str) -> bool:
"""Suppress only repeated forced-turn planning text, not final answers."""
stripped = text.strip()
if not stripped or len(stripped) >= _REPROMPT_MAX_CHARS:
return False
if _FINAL_ANSWER_SIGNAL.search(stripped):
return False
return _FORCED_REPEAT_PLAN_SIGNAL.search(stripped) is not None
# ── Pre-compiled patterns for GGUF shard detection ───────────
_SHARD_FULL_RE = re.compile(r"^(.*)-(\d{5})-of-(\d{5})\.gguf$")
@ -3933,10 +3960,13 @@ class LlamaCppBackend:
# ── Message building (OpenAI format) ──────────────────────────
@staticmethod
def _parse_tool_calls_from_text(content: str) -> list[dict]:
"""Thin wrapper around the shared tool_call_parser so safetensors
and llama_cpp pick up the same fixes."""
return _shared_parse_tool_calls_from_text(content)
def _parse_tool_calls_from_text(content: str, *, allow_incomplete: bool = True) -> list[dict]:
"""Thin wrapper around the shared parser in tool_call_parser
so safetensors and llama_cpp pick up the same fixes."""
return _shared_parse_tool_calls_from_text(
content,
allow_incomplete = allow_incomplete,
)
@staticmethod
def _build_openai_messages(messages: list[dict], image_b64: Optional[str] = None) -> list[dict]:
@ -4282,24 +4312,37 @@ class LlamaCppBackend:
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
def _strip_tool_markup(text: str, *, final: bool = False) -> str:
if not auto_heal_tool_calls:
def _strip_tool_markup(
text: str,
*,
final: bool = False,
force: bool = False,
) -> str:
if not (auto_heal_tool_calls or force):
return text
return strip_tool_call_markup(text, final = final)
# XML prefixes that signal a tool call in content. Empty when
# auto_heal is disabled so the buffer never speculatively holds
# content for XML detection.
_TOOL_XML_SIGNALS = ("<tool_call>", "<function=") if auto_heal_tool_calls else ()
_MAX_BUFFER_CHARS = 32
def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str:
if not (auto_heal_tool_calls or force):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
# ── Duplicate tool-call detection ────────────────────────
# Track recent (tool_name, arguments) hashes to detect loops where
# the model repeats the exact same call. Retries after a transient
# failure are allowed (only block when the prior identical call
# succeeded).
_tool_call_history: list[tuple[str, bool]] = [] # (key, failed)
_render_html_succeeded = False
tool_controller = ToolLoopController(
tools = tools,
auto_heal_tool_calls = auto_heal_tool_calls,
)
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
return any(
record.executed and not record.is_error and record.key.startswith(key_prefix)
for record in tool_controller.history
)
_MAX_BUFFER_CHARS = 32
_append_budget_exhausted_nudge = True
# ── Re-prompt on plan-without-action ─────────────────
# When the model describes what it intends to do (forward-looking
@ -4308,6 +4351,7 @@ class LlamaCppBackend:
# "Hello!" won't match. Pattern compiled at module level
# (_INTENT_SIGNAL).
_reprompt_count = 0
_forced_tool_call_pending = False
# Reserve extra iterations for re-prompts so they don't consume the
# caller's tool-call budget; only when tool iterations are allowed.
@ -4316,8 +4360,14 @@ class LlamaCppBackend:
if cancel_event is not None and cancel_event.is_set():
return
# stream: True so we detect tool signals in the first 1-2 chunks
# without a non-streaming penalty.
active_tools = tool_controller.active_tools()
if not active_tools:
_append_budget_exhausted_nudge = False
break
_tool_xml_signals = TOOL_XML_SIGNALS if active_tools else ()
# Build payload -- stream: True so we detect tool signals
# in the first 1-2 chunks without a non-streaming penalty.
payload = {
"messages": conversation,
"stream": True,
@ -4328,7 +4378,7 @@ class LlamaCppBackend:
"min_p": min_p,
"repeat_penalty": repetition_penalty,
"presence_penalty": presence_penalty,
"tools": tools,
"tools": active_tools,
"tool_choice": "auto",
}
_reasoning_kw = self._request_reasoning_kwargs(
@ -4372,6 +4422,7 @@ class LlamaCppBackend:
_stream_done = False
_last_emitted = ""
provisional_render_html_tool_call_ids = set()
_suppress_visible_output = _forced_tool_call_pending
stream_timeout = httpx.Timeout(
connect = 10,
@ -4413,19 +4464,21 @@ class LlamaCppBackend:
if detect_state == _S_STREAMING and in_thinking:
if has_content_tokens:
cumulative_display += "</think>"
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
else:
cumulative_display = reasoning_accum
yield {
"type": "content",
"text": cumulative_display,
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": cumulative_display,
}
_stream_done = True
break # exit inner while
if not line.startswith("data: "):
@ -4449,11 +4502,11 @@ class LlamaCppBackend:
# ── Structured tool_calls ──
tc_deltas = delta.get("tool_calls")
if tc_deltas:
# Once visible content has been
# emitted, don't reclassify this turn
# as a tool call.
if _last_emitted:
continue
# llama-server can emit visible assistant
# preface content before native structured
# tool_calls. Preserve content_accum as
# the assistant pre-tool text and still
# drain/execute the structured call.
has_structured_tc = True
detect_state = _S_DRAINING
for tc_d in tc_deltas:
@ -4491,8 +4544,16 @@ class LlamaCppBackend:
has_real_id = current_id != fallback_id
if (
current_name == "render_html"
and not _render_html_succeeded
and not _tool_succeeded("render_html")
and any(
(
(tool.get("function") or {}).get("name")
== "render_html"
)
for tool in active_tools
)
and not already_started
and not provisional_render_html_tool_call_ids
and has_real_id
):
provisional_render_html_tool_call_ids.add(
@ -4503,6 +4564,9 @@ class LlamaCppBackend:
"tool_name": "render_html",
"tool_call_id": current_id,
"arguments": {},
"provenance": tool_event_provenance(
provisional = True,
),
}
continue
@ -4520,10 +4584,11 @@ class LlamaCppBackend:
cumulative_display += "<think>"
in_thinking = True
cumulative_display += reasoning
yield {
"type": "content",
"text": cumulative_display,
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": cumulative_display,
}
# ── Content tokens ──
token = delta.get("content", "")
@ -4539,15 +4604,16 @@ class LlamaCppBackend:
cumulative_display += "</think>"
in_thinking = False
cumulative_display += token
cleaned = _strip_tool_markup(
cumulative_display,
cleaned = _strip_tool_markup_streaming(
cumulative_display
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
yield {
"type": "content",
"text": cleaned,
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
}
elif detect_state == _S_BUFFERING:
content_buffer += token
@ -4558,7 +4624,7 @@ class LlamaCppBackend:
# Check tool signal prefixes.
is_prefix = False
is_match = False
for sig in _TOOL_XML_SIGNALS:
for sig in _tool_xml_signals:
if stripped_buf.startswith(sig):
is_match = True
break
@ -4567,6 +4633,25 @@ class LlamaCppBackend:
break
if is_match:
# Tool signal -- flush any visible
# prefix before DRAINING so the
# route sends it before tool_start.
if reasoning_accum:
cumulative_display += "<think>"
cumulative_display += reasoning_accum
cumulative_display += "</think>"
cumulative_display += content_buffer
cleaned = _strip_tool_markup_streaming(
cumulative_display,
force = True,
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
}
detect_state = _S_DRAINING
elif (
is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS
@ -4587,10 +4672,11 @@ class LlamaCppBackend:
)
if len(cleaned) > len(_last_emitted):
_last_emitted = cleaned
yield {
"type": "content",
"text": cleaned,
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": cleaned,
}
except json.JSONDecodeError:
logger.debug(f"Skipping malformed SSE line: {line[:100]}")
@ -4600,11 +4686,7 @@ class LlamaCppBackend:
# ── Resolve BUFFERING at stream end ──
if detect_state == _S_BUFFERING:
stripped_buf = content_buffer.lstrip()
if (
stripped_buf
and auto_heal_tool_calls
and any(s in stripped_buf for s in _TOOL_XML_SIGNALS)
):
if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals):
detect_state = _S_DRAINING
elif content_accum or reasoning_accum:
detect_state = _S_STREAMING
@ -4615,22 +4697,24 @@ class LlamaCppBackend:
cumulative_display += reasoning_accum
cumulative_display += "</think>"
cumulative_display += content_buffer
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": _strip_tool_markup(
cumulative_display,
final = True,
),
}
elif reasoning_accum and not has_content_tokens:
# Reasoning-only response: show reasoning as plain
# text, matching the final streaming pass for
# models that put everything in reasoning.
cumulative_display = reasoning_accum
yield {
"type": "content",
"text": cumulative_display,
}
if not _suppress_visible_output:
yield {
"type": "content",
"text": cumulative_display,
}
else:
return
@ -4641,9 +4725,10 @@ class LlamaCppBackend:
# synthesis streams correctly even if content was emitted
# before the tool XML.
_safety_tc = None
if auto_heal_tool_calls and any(s in content_accum for s in _TOOL_XML_SIGNALS):
if any(s in content_accum for s in _tool_xml_signals):
_safety_tc = self._parse_tool_calls_from_text(
content_accum,
allow_incomplete = auto_heal_tool_calls,
)
if not _safety_tc:
# ── Re-prompt on plan-without-action ──
@ -4656,11 +4741,18 @@ class LlamaCppBackend:
_stripped = content_accum.strip()
if not _stripped:
_stripped = reasoning_accum.strip()
_render_html_already_done_intent = _tool_succeeded(
"render_html"
) and re.search(
r"(?i)\brender[_\s-]?html\b",
_stripped,
)
if (
tools
auto_heal_tool_calls
and active_tools
and not _render_html_already_done_intent
and _reprompt_count < _MAX_REPROMPTS
and 0 < len(_stripped) < _REPROMPT_MAX_CHARS
and _INTENT_SIGNAL.search(_stripped)
and _is_short_intent_without_action(_stripped)
):
_reprompt_count += 1
logger.info(
@ -4675,19 +4767,21 @@ class LlamaCppBackend:
}
)
available_tool_names = [
tool.get("function", {}).get("name")
for tool in tools
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
available_tool_names = [name for name in available_tool_names if name]
tool_hint = " or ".join(available_tool_names) or "an available tool"
_forced_tool_call_pending = True
conversation.append(
{
"role": "user",
"content": (
"STOP. Do NOT write code or explain. "
"You MUST call a tool NOW. "
f"Call {tool_hint} immediately."
"You have access to enabled tools. If a tool is needed to satisfy "
"the user's request or complete the action you described, call "
f"{tool_hint} now. If no tool is needed, provide the final answer "
"and follow the user's requested format."
),
}
)
@ -4700,7 +4794,28 @@ class LlamaCppBackend:
yield {"type": "status", "text": ""}
continue
# Content already streamed; yield metadata.
if _forced_tool_call_pending:
_forced_tool_call_pending = False
if not _should_suppress_forced_no_tool_output(_stripped):
if cumulative_display:
forced_visible_text = _strip_tool_markup(
cumulative_display,
final = True,
)
elif content_accum:
forced_visible_text = _strip_tool_markup(
content_accum,
final = True,
)
else:
forced_visible_text = reasoning_accum
if forced_visible_text:
yield {
"type": "content",
"text": forced_visible_text,
}
# Content was already streamed. Yield metadata.
yield {"type": "status", "text": ""}
_fu = _backfill_usage_from_timings(_iter_usage, _iter_timings) or {}
_fc = _fu.get("completion_tokens", 0)
@ -4733,6 +4848,7 @@ class LlamaCppBackend:
content_text = _strip_tool_markup(
content_accum,
final = True,
force = True,
)
logger.info(
f"Safety net: parsed {len(tool_calls)} tool call(s) "
@ -4750,18 +4866,16 @@ class LlamaCppBackend:
for i in sorted(tool_calls_acc)
if (tool_calls_acc[i].get("function", {}).get("name", "").strip())
] or None
if (
not tool_calls
and auto_heal_tool_calls
and any(s in content_accum for s in _TOOL_XML_SIGNALS)
):
if not tool_calls and any(s in content_accum for s in _tool_xml_signals):
tool_calls = self._parse_tool_calls_from_text(
content_accum,
allow_incomplete = auto_heal_tool_calls,
)
if tool_calls and not has_structured_tc:
content_text = _strip_tool_markup(
content_text,
final = True,
force = True,
)
if tool_calls:
logger.info(
@ -4812,163 +4926,68 @@ class LlamaCppBackend:
_accumulated_predicted_ms += _it.get("predicted_ms", 0)
_accumulated_predicted_n += _it.get("predicted_n", 0)
assistant_msg = {"role": "assistant", "content": content_text}
if tool_calls:
assistant_msg["tool_calls"] = tool_calls
conversation.append(assistant_msg)
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
for tc in tool_calls or []:
func = tc.get("function", {})
tool_name = func.get("name", "")
raw_args = func.get("arguments", {})
if isinstance(raw_args, str):
try:
arguments = json.loads(raw_args)
except (json.JSONDecodeError, ValueError):
if auto_heal_tool_calls:
heal_key = {
"python": "code",
"terminal": "command",
"render_html": "code",
}.get(tool_name, "query")
arguments = {heal_key: raw_args}
else:
arguments = {"raw": raw_args}
else:
arguments = raw_args
if tool_name == "web_search":
_ws_url = (arguments.get("url") or "").strip()
if _ws_url:
_parsed = urlparse(_ws_url)
if _parsed.scheme in ("http", "https") and _parsed.hostname:
_ws_host = _parsed.hostname
if _ws_host.startswith("www."):
_ws_host = _ws_host[4:]
status_text = f"Reading: {_ws_host}"
else:
status_text = "Reading page..."
else:
status_text = f"Searching: {arguments.get('query', '')}"
elif tool_name == "python":
preview = (arguments.get("code") or "").strip().split("\n")[0][:60]
status_text = (
f"Running Python: {preview}" if preview else "Running Python..."
)
elif tool_name == "terminal":
cmd_preview = (arguments.get("command") or "")[:60]
status_text = (
f"Running: {cmd_preview}" if cmd_preview else "Running command..."
)
else:
status_text = f"Calling: {tool_name}"
_repeat_render_html = tool_name == "render_html" and _render_html_succeeded
if not _repeat_render_html:
yield {"type": "status", "text": status_text}
yield {
"type": "tool_start",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"arguments": arguments,
}
# ── Duplicate call detection ──────────────
# str(dict) is stable here: arguments always come from
# json.loads on the same model output within one request,
# so insertion order is deterministic (Python 3.7+).
_tc_key = tool_name + str(arguments)
_prev = _tool_call_history[-1] if _tool_call_history else None
if _repeat_render_html:
result = RENDER_HTML_REPEAT_NUDGE
elif _prev and _prev[0] == _tc_key and not _prev[1]:
result = (
"You already made this exact call. "
"Do not repeat the same tool call. "
"Try a different approach: fetch a URL "
"from previous results, use Python to "
"process data you already have, or "
"provide your final answer now."
)
else:
_effective_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
# Guard against a tool not in the per-request
# advertised set: filtered MCP names, an opted-out
# built-in, or a stale name from a prior turn. Mirrors
# the safetensors loop's allowed_tool_names check.
_allowed = {
(t.get("function") or {}).get("name")
for t in (tools or [])
if (t.get("function") or {}).get("name")
}
if _allowed and tool_name not in _allowed:
result = (
f"Error: tool '{tool_name}' is not enabled "
"for this request. Use one of the enabled "
"tools or provide a final answer."
)
else:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
)
if not _repeat_render_html:
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
# Nudge the model toward a different approach on errors.
_error_prefixes = (
"Error",
"Search failed",
"Execution error",
"Blocked:",
"Exit code",
"Failed to fetch",
"Failed to resolve",
"No query provided",
provisional_render_html_match = (
tool_name == "render_html"
and tc.get("id") in provisional_render_html_tool_call_ids
)
_is_error = isinstance(result, str) and result.lstrip().startswith(
_error_prefixes
decision = tool_controller.prepare_call(
tc,
forced = _forced_tool_call_pending,
provisional = provisional_render_html_match,
)
if tool_name == "render_html" and not _is_error:
_render_html_succeeded = True
_tool_call_history.append((_tc_key, _is_error))
# Strip image sentinel before feeding the result to the
# LLM (the full result with sentinel is still yielded via
# tool_end so the frontend can extract image paths).
_result_content = result
if "\n__IMAGES__:" in _result_content:
_result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
if _is_error:
_result_content = (
_result_content + "\n\nThe tool call encountered an issue. "
"Please try a different approach or rephrase your request."
if not decision.should_execute:
if content_text and not assistant_appended:
conversation.append(assistant_msg)
assistant_appended = True
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
if _forced_tool_call_pending:
_forced_tool_call_pending = False
logger.info(
"Suppressed local GGUF tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
conversation.append(assistant_msg)
assistant_appended = True
else:
assistant_msg.setdefault("tool_calls", []).append(
decision.as_assistant_tool_call()
)
tool_msg = {
"role": "tool",
"name": tool_name,
"content": _result_content,
}
tool_call_id = tc.get("id")
if tool_call_id:
tool_msg["tool_call_id"] = tool_call_id
conversation.append(tool_msg)
yield {"type": "status", "text": decision.status_text}
yield decision.tool_start_event()
# Clear tool status badge before the next generation.
_effective_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
result = execute_tool(
decision.tool_name,
decision.arguments,
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
)
completion = tool_controller.record_result(decision, result)
yield completion.tool_end_event()
conversation.append(completion.tool_message())
if _forced_tool_call_pending:
_forced_tool_call_pending = False
# Clear tool status badge before next generation/final pass.
yield {"type": "status", "text": ""}
# Continue so the model responds with tool context.
if tool_controller.force_final_answer or not tool_controller.active_tools():
_append_budget_exhausted_nudge = False
break
continue
except httpx.ConnectError:
@ -4982,7 +5001,7 @@ class LlamaCppBackend:
# The model used all iterations without a final text response. Nudge
# the final streaming pass to produce a useful answer instead of
# continuing to request tools.
if max_tool_iterations > 0:
if max_tool_iterations > 0 and _append_budget_exhausted_nudge:
conversation.append(
{
"role": "user",

9
studio/backend/core/inference/orchestrator.py
View file

@ -836,8 +836,11 @@ class InferenceOrchestrator:
max_new_tokens = max_tokens if max_tokens and max_tokens > 0 else 2048
def _single_turn(conv: list):
# ``conv`` already carries any system message.
def _single_turn(conv: list, *, active_tools: Optional[list[dict]] = None):
# ``conv`` already carries any system message. ``active_tools`` lets
# run_safetensors_tool_loop drop one-shot tools (e.g. render_html) from
# later same-response prompts.
turn_tools = active_tools if active_tools is not None else tools
common_kwargs = dict(
messages = conv,
system_prompt = "",
@ -849,7 +852,7 @@ class InferenceOrchestrator:
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
tools = turn_tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,

385
studio/backend/core/inference/safetensors_agentic.py
View file

@ -14,25 +14,25 @@ parses tool calls from the cumulative text and dispatches via
``core.inference.tools``.
"""
import json
import re
import threading
from typing import Callable, Generator, Optional
from urllib.parse import urlparse
from loggers import get_logger
from core.inference.tool_call_parser import (
_TOOL_ALL_PATS,
BUDGET_EXHAUSTED_NUDGE,
DUPLICATE_CALL_NUDGE,
RENDER_HTML_REPEAT_NUDGE,
TOOL_ERROR_NUDGE,
TOOL_ERROR_PREFIXES,
TOOL_XML_SIGNALS,
has_tool_signal,
parse_tool_calls_from_text,
strip_tool_markup,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
coerce_tool_arguments,
status_for_tool,
tool_event_provenance,
)
logger = get_logger(__name__)
@ -42,34 +42,34 @@ logger = get_logger(__name__)
_MAX_BUFFER_CHARS = 32
def strip_tool_markup_streaming(
text: str,
*,
auto_heal_tool_calls: bool = True,
tool_protocol_active: bool = False,
) -> str:
"""Strip open-ended tool XML from display text without trimming whitespace."""
if not (auto_heal_tool_calls or tool_protocol_active):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
def _strip_tool_markup_final(
text: str,
*,
auto_heal_tool_calls: bool,
tool_protocol_active: bool = False,
) -> str:
if not (auto_heal_tool_calls or tool_protocol_active):
return text
return strip_tool_markup(text, final = True)
def _status_for_tool(tool_name: str, arguments: dict) -> str:
"""Return a human-readable status line matching the GGUF path."""
if tool_name == "web_search":
url = (arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):
host = host[4:]
return f"Reading: {host}"
return "Reading page..."
query = arguments.get("query", "")
return f"Searching: {query}"
if tool_name == "python":
preview = (arguments.get("code") or "").strip().split("\n")[0][:60]
return f"Running Python: {preview}" if preview else "Running Python..."
if tool_name == "terminal":
preview = (arguments.get("command") or "")[:60]
return f"Running: {preview}" if preview else "Running command..."
return f"Calling: {tool_name}"
_CANONICAL_HEAL_ARG = {
"python": "code",
"terminal": "command",
"render_html": "code",
}
return status_for_tool(tool_name, arguments)
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
@ -93,34 +93,43 @@ def _detect_render_html_tool_start(content: str) -> bool:
return False
def _coerce_arguments_with_provenance(
raw_args,
*,
heal: bool,
tool_name: str = "",
):
"""Normalise tool ``arguments`` and report whether healing was applied."""
coerced = coerce_tool_arguments(raw_args, heal = heal, tool_name = tool_name)
return coerced.arguments, coerced.healed
def _coerce_arguments(
raw_args,
*,
heal: bool,
tool_name: str = "",
) -> dict:
"""Normalise tool ``arguments`` to a dict.
arguments, _ = _coerce_arguments_with_provenance(
raw_args,
heal = heal,
tool_name = tool_name,
)
return arguments
Some templates emit a JSON string, others a bare query string. With
``heal=True`` a bare string becomes ``{<canonical_key>: ...}`` so a
Hermes-style call without proper JSON still runs. Canonical key per
tool: ``code`` for python, ``command`` for terminal, ``query`` otherwise
(e.g. web_search).
"""
if isinstance(raw_args, dict):
return raw_args
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, dict):
return parsed
except (json.JSONDecodeError, ValueError):
pass
if heal:
key = _CANONICAL_HEAL_ARG.get(tool_name, "query")
return {key: raw_args}
return {"raw": raw_args}
return {}
def _tool_event_provenance(**flags: object) -> dict[str, object]:
return tool_event_provenance(**flags)
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
"""Call a single-turn generator with active tool schemas when supported."""
try:
return single_turn(conversation, active_tools = active_tools)
except TypeError as exc:
if "active_tools" not in str(exc):
raise
return single_turn(conversation)
def run_safetensors_tool_loop(
@ -158,16 +167,21 @@ def run_safetensors_tool_loop(
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
"""
conversation = list(messages)
tool_call_history: list[tuple[str, bool]] = []
render_html_succeeded = False
unrestricted_tools = not tools
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
)
final_attempt_done = False
allowed_tool_names = {
(tool.get("function") or {}).get("name")
for tool in (tools or [])
if (tool.get("function") or {}).get("name")
}
next_call_id = 0
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
return any(
record.executed and not record.is_error and record.key.startswith(key_prefix)
for record in tool_controller.history
)
if max_tool_iterations <= 0:
# 0 = disabled (same contract as the GGUF loop).
yield {"type": "status", "text": ""}
@ -181,6 +195,17 @@ def run_safetensors_tool_loop(
if cancel_event is not None and cancel_event.is_set():
return
if final_attempt_done:
active_tools: list[dict] = []
else:
active_tools = tool_controller.active_tools()
if not active_tools and not unrestricted_tools:
final_attempt_done = True
active_tools = []
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
detect_state = _state_buffering
content_buffer = ""
content_accum = ""
@ -189,7 +214,7 @@ def run_safetensors_tool_loop(
provisional_render_html_started = False
provisional_render_html_id = f"call_{next_call_id}"
gen = single_turn(conversation)
gen = _call_single_turn(single_turn, conversation, active_tools)
prev_cumulative = ""
for cumulative in gen:
@ -207,7 +232,11 @@ def run_safetensors_tool_loop(
if detect_state == _state_draining:
if (
not render_html_succeeded
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
@ -217,26 +246,35 @@ def run_safetensors_tool_loop(
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
continue
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in TOOL_XML_SIGNALS:
for sig in tool_xml_signals:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup(before_tool)
cleaned_before = strip_tool_markup_streaming(
before_tool,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
yield {"type": "content", "text": cleaned_before}
cumulative_display = candidate
detect_state = _state_draining
if (
not render_html_succeeded
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
@ -246,10 +284,15 @@ def run_safetensors_tool_loop(
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
continue
cumulative_display = candidate
cleaned = strip_tool_markup(cumulative_display)
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
@ -263,7 +306,7 @@ def run_safetensors_tool_loop(
is_match = False
is_prefix = False
for sig in TOOL_XML_SIGNALS:
for sig in tool_xml_signals:
if stripped.startswith(sig):
is_match = True
break
@ -272,9 +315,24 @@ def run_safetensors_tool_loop(
break
if is_match:
# Tool signal -- flush any visible prefix before DRAINING
# so the route sends it before tool_start.
cumulative_display += content_buffer
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
detect_state = _state_draining
if (
not render_html_succeeded
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
@ -284,13 +342,18 @@ def run_safetensors_tool_loop(
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
continue
else:
detect_state = _state_streaming
cumulative_display += content_buffer
cleaned = strip_tool_markup(cumulative_display)
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
@ -302,14 +365,22 @@ def run_safetensors_tool_loop(
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content?
stripped = content_buffer.lstrip()
if stripped and has_tool_signal(stripped):
if (
stripped
and tool_protocol_active
and any(sig in stripped for sig in tool_xml_signals)
):
detect_state = _state_draining
else:
if content_buffer:
cumulative_display += content_buffer
yield {
"type": "content",
"text": strip_tool_markup(cumulative_display, final = True),
"text": _strip_tool_markup_final(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
yield {"type": "status", "text": ""}
return
@ -317,19 +388,30 @@ def run_safetensors_tool_loop(
if detect_state == _state_streaming:
# No tool detected mid-stream -- check for late tool XML.
safety_tc = None
if has_tool_signal(content_accum):
saw_tool_signal = tool_protocol_active and any(
sig in content_accum for sig in tool_xml_signals
)
if saw_tool_signal:
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not safety_tc:
# Final answer: streaming already emitted content. Skip the
# final=True re-strip so literal "<tool_call>" in prose
# survives when no real tool call parsed.
# Final answer: if a literal tool marker in prose was stripped
# during streaming but did not parse as a real call, restore the
# raw cumulative text for core callers. Route-level cleanup can
# still apply the Auto-Heal display policy.
if saw_tool_signal and content_accum:
yield {"type": "content", "text": content_accum}
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
content_text = strip_tool_markup(content_accum, final = True)
content_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
@ -339,22 +421,39 @@ def run_safetensors_tool_loop(
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not tool_calls and auto_heal_tool_calls:
# Parser found nothing -- surface raw content so literal
# "<tool_call>" prose is preserved.
if not tool_calls:
# Parser found nothing. Auto-Heal-enabled display cleanup
# strips unparseable tool XML; disabled Auto-Heal preserves
# the raw text so literal/malformed markup stays visible.
if content_accum:
yield {"type": "content", "text": content_accum}
yield {
"type": "content",
"text": _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
if provisional_render_html_started:
yield {
"type": "tool_end",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"result": "Error: render_html tool call could not be parsed.",
"provenance": _tool_event_provenance(provisional = True),
}
yield {"type": "status", "text": ""}
return
content_text = strip_tool_markup(content_accum, final = True)
content_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
)
if tool_calls:
next_call_id += len(tool_calls)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
@ -364,99 +463,69 @@ def run_safetensors_tool_loop(
return
assistant_msg: dict = {"role": "assistant", "content": content_text}
if tool_calls:
assistant_msg["tool_calls"] = tool_calls
next_call_id += len(tool_calls)
conversation.append(assistant_msg)
assistant_appended = False
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
tool_name = func.get("name", "") or ""
arguments = _coerce_arguments(
func.get("arguments", {}),
heal = auto_heal_tool_calls,
tool_name = tool_name,
provisional_match = (
provisional_render_html_started
and tool_name == "render_html"
and tc.get("id", "") == provisional_render_html_id
)
decision = tool_controller.prepare_call(tc, provisional = provisional_match)
repeat_render_html = tool_name == "render_html" and render_html_succeeded
if not repeat_render_html:
yield {"type": "status", "text": _status_for_tool(tool_name, arguments)}
yield {
"type": "tool_start",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"arguments": arguments,
}
tc_key = tool_name + str(arguments)
if repeat_render_html:
result = RENDER_HTML_REPEAT_NUDGE
elif allowed_tool_names and tool_name not in allowed_tool_names:
result = (
f"Error: tool '{tool_name}' is not enabled for this "
"request. Use one of the enabled tools or provide a "
"final answer."
if not decision.should_execute:
if content_text and not assistant_appended:
conversation.append(assistant_msg)
assistant_appended = True
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
logger.info(
"Suppressed local safetensors tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
conversation.append(assistant_msg)
assistant_appended = True
else:
already_ran_ok = any(k == tc_key and not err for k, err in tool_call_history)
if already_ran_ok:
result = DUPLICATE_CALL_NUDGE
else:
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
try:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
if not repeat_render_html:
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
yield {"type": "status", "text": decision.status_text}
yield decision.tool_start_event()
is_error = isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
if tool_name == "render_html" and not is_error:
render_html_succeeded = True
tool_call_history.append((tc_key, is_error))
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
try:
result = execute_tool(
decision.tool_name,
decision.arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
# Strip frontend image sentinel from the model's view. Cut at the
# first occurrence so leading and consecutive sentinels both go.
result_for_model = result
if isinstance(result_for_model, str) and "__IMAGES__:" in result_for_model:
result_for_model = result_for_model.split("__IMAGES__:", 1)[0].rstrip()
if is_error:
result_for_model = result_for_model + TOOL_ERROR_NUDGE
tool_msg: dict = {
"role": "tool",
"name": tool_name,
"content": result_for_model,
}
tool_call_id = tc.get("id")
if tool_call_id:
tool_msg["tool_call_id"] = tool_call_id
conversation.append(tool_msg)
completion = tool_controller.record_result(decision, result)
yield completion.tool_end_event()
conversation.append(completion.tool_message())
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}
if tool_controller.force_final_answer:
final_attempt_done = True
continue
if not unrestricted_tools and not tool_controller.active_tools():
final_attempt_done = True
continue
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
# Budget exhausted; nudge a final plain answer.
final_attempt_done = True
conversation.append(
{
"role": "user",
"content": BUDGET_EXHAUSTED_NUDGE,
}
)
conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})
yield {"type": "status", "text": ""}

86
studio/backend/core/inference/tool_call_parser.py
View file

@ -102,7 +102,12 @@ def strip_tool_markup(text: str, *, final: bool = False) -> str:
return text.strip() if final else text
def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict]:
def parse_tool_calls_from_text(
content: str,
*,
id_offset: int = 0,
allow_incomplete: bool = True,
) -> list[dict]:
"""Parse OpenAI-format ``tool_calls`` from model text.
Returns a list of ``{"id", "type", "function": {"name", "arguments"}}``
@ -116,8 +121,10 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
- XML-style function blocks:
``<function=name><parameter=k>v</parameter></function>``
Closing tags (``</tool_call>``, ``</function>``, ``</parameter>``)
are all optional since models frequently omit them.
``allow_incomplete=True`` keeps the historical healing behavior for
missing closing tags. ``allow_incomplete=False`` accepts only
well-formed wrappers so disabled Auto-Heal can still parse valid
local tool protocol without repairing truncated output.
"""
tool_calls: list[dict] = []
@ -144,23 +151,28 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"])
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
if depth != 0:
continue
if not allow_incomplete:
tail_after_json = content[i + 1 :].lstrip()
if _TC_END_TAG_RE.match(tail_after_json) is None:
continue
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"])
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: <function=name><parameter=k>v... -- closing tags optional;
# </function> isn't a body boundary since code values can contain it.
@ -181,7 +193,19 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
body = _TC_FUNC_CLOSE_RE.sub("", body)
if not allow_incomplete:
# Bound the body at the closing </function> tag rather than
# the end of the response, so a complete call followed by
# trailing prose is still accepted (matching the JSON-style
# <tool_call> path, which already tolerates trailing text).
# rfind picks the last </function>, so a literal </function>
# inside a code parameter value stays in the body.
close_idx = body.rfind(_FUNC_CLOSE_TAG)
if close_idx < 0:
continue
body = body[:close_idx]
else:
body = _TC_FUNC_CLOSE_RE.sub("", body)
arguments: dict = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
@ -190,9 +214,16 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
# </parameter> in code strings is preserved.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
continue
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
valid_params = True
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
@ -202,8 +233,17 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
else len(body)
)
val = body[val_start:next_param]
val = _TC_PARAM_CLOSE_RE.sub("", val)
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
valid_params = False
break
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
if not valid_params:
continue
tc = {
"id": f"call_{id_offset + len(tool_calls)}",

410
studio/backend/core/inference/tool_loop_controller.py Normal file
View file

@ -0,0 +1,410 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared controller state for Studio local agentic tool loops.
This module is intentionally dependency-light: it owns only per-response
ledger state and value objects used by the GGUF and safetensors loops.
Route/SSE conversion, tool execution, and model streaming stay in the
backend-specific modules.
"""
from __future__ import annotations
import copy
import json
from dataclasses import dataclass, field
from typing import Any, Literal, Mapping, Sequence
from urllib.parse import urlparse
from core.inference.tool_call_parser import TOOL_ERROR_NUDGE, TOOL_ERROR_PREFIXES
_CANONICAL_HEAL_ARG = {
"python": "code",
"terminal": "command",
"render_html": "code",
}
_ONE_SHOT_TOOLS = frozenset({"render_html"})
NoopReason = Literal["duplicate", "disabled", "render_html_repeat"]
ToolAction = Literal["execute", "duplicate", "disabled", "render_html_repeat"]
@dataclass(frozen = True)
class CoercedArguments:
"""Normalized tool arguments plus whether healing changed the shape."""
arguments: dict[str, Any]
healed: bool = False
@dataclass(frozen = True)
class ToolCallDecision:
"""Decision made before any visible tool event is emitted."""
action: ToolAction
tool_name: str
arguments: dict[str, Any]
tool_call_id: str = ""
key: str = ""
provenance: dict[str, Any] = field(default_factory = dict)
status_text: str = ""
noop_result: str = ""
@property
def should_execute(self) -> bool:
return self.action == "execute"
@property
def emit_visible_events(self) -> bool:
"""Only real executions should become frontend-visible tool cards."""
return self.should_execute
@property
def noop_reason(self) -> NoopReason | None:
if self.action == "execute":
return None
return self.action
def tool_start_payload(self) -> dict[str, Any]:
"""Build the payload fields for a real tool_start event."""
return {
"tool_name": self.tool_name,
"tool_call_id": self.tool_call_id,
"arguments": self.arguments,
"provenance": self.provenance,
}
def tool_start_event(self) -> dict[str, Any]:
"""Build the existing backend event shape for a real execution."""
return {"type": "tool_start", **self.tool_start_payload()}
def as_assistant_tool_call(self) -> dict[str, Any]:
"""Return an OpenAI-style tool_call with normalized arguments."""
tool_call: dict[str, Any] = {
"type": "function",
"function": {
"name": self.tool_name,
"arguments": json.dumps(
self.arguments,
ensure_ascii = False,
sort_keys = True,
separators = (",", ":"),
),
},
}
if self.tool_call_id:
tool_call["id"] = self.tool_call_id
return tool_call
@dataclass(frozen = True)
class ToolCallCompletion:
"""Result/nudge that should be fed back to the next model turn."""
decision: ToolCallDecision
result: str
is_error: bool = False
executed: bool = False
def tool_end_payload(self) -> dict[str, Any]:
"""Build the payload fields for a real tool_end event."""
return {
"tool_name": self.decision.tool_name,
"tool_call_id": self.decision.tool_call_id,
"result": self.result,
"provenance": self.decision.provenance,
}
def tool_end_event(self) -> dict[str, Any]:
"""Build the existing backend event shape for a real execution result."""
return {"type": "tool_end", **self.tool_end_payload()}
def tool_message(self) -> dict[str, Any]:
"""Return the OpenAI-compatible tool message for a real execution."""
if not self.executed:
raise ValueError("No-op completions are internal nudges, not tool messages")
return self.model_message()
def model_message(self) -> dict[str, Any]:
"""Return the internal message appended before the next generation.
Executed calls keep the existing OpenAI-compatible ``role=tool``
continuation. No-op controller decisions are not real tool output, so
they are fed back as a hidden user nudge rather than a normal tool
result.
"""
if not self.executed:
return {"role": "user", "content": self.result}
content = strip_result_for_model(self.result)
if self.is_error:
content = content + TOOL_ERROR_NUDGE
message: dict[str, Any] = {
"role": "tool",
"name": self.decision.tool_name,
"content": content,
}
if self.decision.tool_call_id:
message["tool_call_id"] = self.decision.tool_call_id
return message
@dataclass(frozen = True)
class _ToolCallRecord:
key: str
is_error: bool
executed: bool
action: ToolAction
def _json_default(value: Any) -> str:
return str(value)
def canonical_tool_call_key(tool_name: str, arguments: Mapping[str, Any]) -> str:
"""Return a stable key for duplicate detection."""
canonical_args = json.dumps(
dict(arguments),
ensure_ascii = False,
sort_keys = True,
separators = (",", ":"),
default = _json_default,
)
return f"{tool_name}:{canonical_args}"
def coerce_tool_arguments(
raw_args: Any,
*,
heal: bool,
tool_name: str = "",
) -> CoercedArguments:
"""Normalize model-emitted ``function.arguments`` to a dictionary."""
if isinstance(raw_args, Mapping):
return CoercedArguments(dict(raw_args), False)
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, Mapping):
return CoercedArguments(dict(parsed), False)
except (json.JSONDecodeError, ValueError):
pass
if heal:
key = _CANONICAL_HEAL_ARG.get(tool_name, "query")
return CoercedArguments({key: raw_args}, True)
return CoercedArguments({"raw": raw_args}, False)
return CoercedArguments({}, False)
def tool_event_provenance(**flags: object) -> dict[str, object]:
"""Return provenance metadata with falsey flags omitted."""
provenance: dict[str, object] = {"source": "local"}
for key, value in flags.items():
if value is not None and value is not False:
provenance[key] = value
return provenance
def status_for_tool(tool_name: str, arguments: Mapping[str, Any]) -> str:
"""Return the status text already used by local tool streams."""
if tool_name == "web_search":
url = str(arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):
host = host[4:]
return f"Reading: {host}"
return "Reading page..."
return f"Searching: {arguments.get('query', '')}"
if tool_name == "python":
preview = str(arguments.get("code") or "").strip().split("\n")[0][:60]
return f"Running Python: {preview}" if preview else "Running Python..."
if tool_name == "terminal":
preview = str(arguments.get("command") or "")[:60]
return f"Running: {preview}" if preview else "Running command..."
return f"Calling: {tool_name}"
def is_tool_error(result: str) -> bool:
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
def strip_result_for_model(result: str) -> str:
"""Remove frontend-only image sentinels before feeding the model."""
if "__IMAGES__:" in result:
return result.split("__IMAGES__:", 1)[0].rstrip()
return result
def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
function = tool.get("function")
if not isinstance(function, Mapping):
return ""
name = function.get("name")
return str(name or "")
def _noop_result(reason: NoopReason, tool_name: str) -> str:
if reason == "duplicate":
return (
"The previous tool request was not executed because this exact "
"tool call already completed successfully. Do not repeat the same "
"tool call. Continue with a different enabled tool if that would "
"materially help, or provide the final answer if you have enough "
"information."
)
if reason == "render_html_repeat":
return (
"render_html completed successfully earlier in this assistant "
"response. Do not call render_html again unless the user asks for "
"changes. Do not mention this internal instruction. Provide only "
"the requested final note or answer."
)
return (
f"The previous tool request was not executed because tool "
f"'{tool_name}' is not enabled for this request. Provide the "
"final answer now without calling more tools."
)
class ToolLoopController:
"""Per-response ledger for local agentic tool loops."""
def __init__(
self,
*,
tools: Sequence[Mapping[str, Any]] | None,
auto_heal_tool_calls: bool = True,
one_shot_tools: frozenset[str] = _ONE_SHOT_TOOLS,
duplicate_noop_limit: int = 2,
) -> None:
self._restrict_to_allowed = tools is not None
self._tools = [copy.deepcopy(dict(tool)) for tool in (tools or [])]
self._allowed_tool_names = {
name for name in (_tool_name_from_schema(tool) for tool in self._tools) if name
}
self._auto_heal_tool_calls = auto_heal_tool_calls
self._one_shot_tools = one_shot_tools
self._completed_one_shot_tools: set[str] = set()
self._successful_keys: set[str] = set()
self._duplicate_noop_counts: dict[str, int] = {}
self._duplicate_noop_limit = max(1, duplicate_noop_limit)
self._history: list[_ToolCallRecord] = []
self._force_final_answer = False
@property
def history(self) -> tuple[_ToolCallRecord, ...]:
return tuple(self._history)
@property
def force_final_answer(self) -> bool:
"""True once a terminal no-op should transition to a no-tools pass."""
return self._force_final_answer
def active_tools(self) -> list[dict[str, Any]]:
"""Return tools still worth advertising to the next model call."""
if self._force_final_answer:
return []
active: list[dict[str, Any]] = []
for tool in self._tools:
name = _tool_name_from_schema(tool)
if name in self._completed_one_shot_tools:
continue
active.append(copy.deepcopy(tool))
return active
def prepare_call(
self,
tool_call: Mapping[str, Any],
*,
forced: bool = False,
provisional: bool = False,
) -> ToolCallDecision:
"""Classify a parsed tool call before any visible event is yielded."""
function = tool_call.get("function")
function = function if isinstance(function, Mapping) else {}
tool_name = str(function.get("name") or "").strip()
coerced = coerce_tool_arguments(
function.get("arguments", {}),
heal = self._auto_heal_tool_calls,
tool_name = tool_name,
)
key = canonical_tool_call_key(tool_name, coerced.arguments)
provenance = tool_event_provenance(
healed = coerced.healed,
forced = forced,
provisional = provisional,
)
action: ToolAction = "execute"
noop = ""
if tool_name in self._completed_one_shot_tools:
action = "render_html_repeat"
noop = _noop_result("render_html_repeat", tool_name)
elif self._restrict_to_allowed and tool_name not in self._allowed_tool_names:
action = "disabled"
noop = _noop_result("disabled", tool_name)
elif key in self._successful_keys:
action = "duplicate"
noop = _noop_result("duplicate", tool_name)
return ToolCallDecision(
action = action,
tool_name = tool_name,
arguments = coerced.arguments,
tool_call_id = str(tool_call.get("id") or ""),
key = key,
provenance = provenance,
status_text = status_for_tool(tool_name, coerced.arguments),
noop_result = noop,
)
def record_result(self, decision: ToolCallDecision, result: Any) -> ToolCallCompletion:
"""Record a real tool execution and return model/frontend payload helpers."""
result_text = result if isinstance(result, str) else str(result)
failed = is_tool_error(result_text)
self._history.append(
_ToolCallRecord(
key = decision.key,
is_error = failed,
executed = True,
action = decision.action,
)
)
if not failed:
self._successful_keys.add(decision.key)
if decision.tool_name in self._one_shot_tools:
self._completed_one_shot_tools.add(decision.tool_name)
return ToolCallCompletion(
decision = decision,
result = result_text,
is_error = failed,
executed = True,
)
def record_noop(self, decision: ToolCallDecision) -> ToolCallCompletion:
"""Record a controller no-op without creating visible tool output."""
self._history.append(
_ToolCallRecord(
key = decision.key,
is_error = False,
executed = False,
action = decision.action,
)
)
if decision.action == "duplicate":
duplicate_count = self._duplicate_noop_counts.get(decision.key, 0) + 1
self._duplicate_noop_counts[decision.key] = duplicate_count
if duplicate_count >= self._duplicate_noop_limit:
self._force_final_answer = True
elif decision.action in ("disabled", "render_html_repeat"):
self._force_final_answer = True
return ToolCallCompletion(
decision = decision,
result = decision.noop_result,
is_error = False,
executed = False,
)

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

@ -538,27 +538,63 @@ async def _await_cancel_then_close(cancel_event, resp) -> None:
return
# Appended to tool-use nudge to discourage plan-without-action.
# Gate render_html guidance to turns where the artifact tool is in the schema;
# otherwise small local models hallucinate a missing tool call instead of the
# fenced-HTML fallback prompt.
_TOOL_ACTION_NUDGE = (
" IMPORTANT: Always call tools directly -- never write code yourself."
" Never describe what you plan to do -- just call the tool immediately."
" For non-artifact code requests, call the python tool when it is available."
" For factual questions that require current information, call web_search when it is available."
" Do NOT output raw code blocks when an enabled tool can satisfy the request."
# Centralized local/server tool nudge. Keep render_html guidance gated to turns
# where the artifact tool is actually present in the tool schema; otherwise
# small local models can hallucinate a missing tool call instead of following
# the fenced-HTML fallback prompt.
_TOOL_BASE_NUDGE = (
"Tools are available when they materially improve the answer. Use an enabled "
"tool for current facts, calculations, code execution, or artifacts when it "
"materially helps; otherwise answer normally and follow the user's requested "
"format."
)
_ARTIFACT_TOOL_ACTION_NUDGE = (
" For HTML, CSS, or JavaScript artifact requests, call render_html once when "
"it is available. After render_html succeeds, do not call it again in the "
"same response unless the user asks for changes. Future user requests for "
"new artifacts may call render_html once."
_TOOL_WEB_COMPACT_TIP = "When using web_search, do not repeat the same search query."
_TOOL_WEB_EXPANDED_TIP = (
"When using web_search and a result URL is relevant, fetch its full content "
"by calling web_search with the url parameter. Do not repeat the same search "
"query. If a search returns no useful results, try rephrasing or fetching a "
"result URL directly."
)
_TOOL_CODE_TIP = (
"Use code execution for math, calculations, data processing, or to parse "
"and analyze information from tool results."
)
_TOOL_ARTIFACT_TIP = (
"For HTML, CSS, or JavaScript artifact requests, call render_html once when "
"it is available with one complete self-contained HTML document in the code "
"argument. After render_html succeeds, do not call it again in the same "
"response unless the user asks for changes. Future user requests for new "
"artifacts may call render_html once."
)
def _tool_action_nudge(has_artifact: bool) -> str:
return _TOOL_ACTION_NUDGE + (_ARTIFACT_TOOL_ACTION_NUDGE if has_artifact else "")
def _build_tool_action_nudge(*, tools: list[dict], model_name: str) -> str:
tool_names = {
(tool.get("function") or {}).get("name")
for tool in tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
}
has_web = "web_search" in tool_names
has_code = "python" in tool_names or "terminal" in tool_names
has_artifact = "render_html" in tool_names
if not (has_web or has_code or has_artifact):
return ""
model_size_b = _extract_model_size_b(model_name)
compact_web_tip = model_size_b is not None and model_size_b < 9
tool_tip_parts: list[str] = []
if has_web:
tool_tip_parts.append(_TOOL_WEB_COMPACT_TIP if compact_web_tip else _TOOL_WEB_EXPANDED_TIP)
if has_code:
tool_tip_parts.append(_TOOL_CODE_TIP)
if has_artifact:
tool_tip_parts.append(_TOOL_ARTIFACT_TIP)
return (
f"The current date is {_date.today().isoformat()}. "
+ _TOOL_BASE_NUDGE
+ " "
+ " ".join(tool_tip_parts)
)
# Strip tool-call XML the speculative buffer in core/inference/llama_cpp.py
@ -576,6 +612,15 @@ _TOOL_XML_RE = _re.compile(
r"|</parameter>\s*\Z",
_re.DOTALL,
)
def _strip_tool_xml_for_display(text: str, *, auto_heal_tool_calls: bool) -> str:
"""Apply route-level XML leak cleanup only when Auto-Heal is enabled."""
if not auto_heal_tool_calls:
return text
return _TOOL_XML_RE.sub("", text)
logger = get_logger(__name__)
@ -2722,69 +2767,30 @@ async def openai_chat_completions(
if use_tools:
# ── Tool-use system prompt nudge ──────────────────────
_tool_names = {t["function"]["name"] for t in tools_to_use}
_has_web = "web_search" in _tool_names
_has_code = "python" in _tool_names or "terminal" in _tool_names
_has_artifact = "render_html" in _tool_names
_date_line = f"The current date is {_date.today().isoformat()}."
# Small models (<9B) struggle with multi-step search plans, so
# simplify the web tips to avoid plan-then-stall behavior.
_model_size_b = _extract_model_size_b(model_name)
_is_small_model = _model_size_b is not None and _model_size_b < 9
if _is_small_model:
_web_tips = "Do not repeat the same search query."
else:
_web_tips = (
"When you search and find a relevant URL in the results, "
"fetch its full content by calling web_search with the url parameter. "
"Do not repeat the same search query. If a search returns "
"no useful results, try rephrasing or fetching a result URL directly."
)
_code_tips = (
"Use code execution for math, calculations, data processing, "
"or to parse and analyze information from tool results."
_nudge = _build_tool_action_nudge(
tools = tools_to_use,
model_name = model_name,
)
_artifact_tips = (
"Use render_html for HTML, CSS, or JavaScript artifact requests "
"with one complete self-contained HTML document in the code argument. "
"Call it once, then do not call it again in the same response unless "
"the user asks for changes. Future user requests for new artifacts may "
"call render_html once."
)
_tool_tip_parts = []
if _has_web:
_tool_tip_parts.append(_web_tips)
if _has_code:
_tool_tip_parts.append(_code_tips)
if _has_artifact:
_tool_tip_parts.append(_artifact_tips)
if _tool_tip_parts:
_nudge = (
_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. " + " ".join(_tool_tip_parts)
)
else:
_nudge = ""
if _nudge:
_nudge += _tool_action_nudge(_has_artifact)
# Append nudge to system prompt (preserve the user's prompt)
# Append nudge to system prompt (preserve user's prompt)
if system_prompt:
system_prompt = system_prompt.rstrip() + "\n\n" + _nudge
else:
system_prompt = _nudge
gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt)
_gguf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# ── Strip stale tool-call XML from conversation history ─
for _msg in gguf_messages:
if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str):
_msg["content"] = _TOOL_XML_RE.sub("", _msg["content"]).strip()
_msg["content"] = _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
).strip()
def gguf_generate_with_tools():
return llama_backend.generate_chat_completion_with_tools(
@ -2801,9 +2807,7 @@ async def openai_chat_completions(
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
preserve_thinking = payload.preserve_thinking,
auto_heal_tool_calls = payload.auto_heal_tool_calls
if payload.auto_heal_tool_calls is not None
else True,
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
max_tool_iterations = payload.max_tool_calls_per_message
if payload.max_tool_calls_per_message is not None
else 25,
@ -2884,7 +2888,10 @@ async def openai_chat_completions(
# cumulative then diff against the last sanitized
# snapshot so cross-chunk XML tags are handled correctly.
raw_cumulative = event.get("text", "")
clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _gguf_auto_heal_tool_calls,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
@ -3228,61 +3235,22 @@ async def openai_chat_completions(
_sf_use_tools = False
if _sf_use_tools:
_sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use}
_sf_has_web = "web_search" in _sf_tool_names
_sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names
_sf_has_artifact = "render_html" in _sf_tool_names
_sf_date_line = f"The current date is {_date.today().isoformat()}."
_sf_model_size_b = _extract_model_size_b(model_name)
_sf_is_small_model = _sf_model_size_b is not None and _sf_model_size_b < 9
if _sf_is_small_model:
_sf_web_tips = "Do not repeat the same search query."
else:
_sf_web_tips = (
"When you search and find a relevant URL in the results, "
"fetch its full content by calling web_search with the url parameter. "
"Do not repeat the same search query. If a search returns "
"no useful results, try rephrasing or fetching a result URL directly."
)
_sf_code_tips = (
"Use code execution for math, calculations, data processing, "
"or to parse and analyze information from tool results."
_sf_nudge = _build_tool_action_nudge(
tools = _sf_tools_to_use,
model_name = model_name,
)
_sf_artifact_tips = (
"Use render_html for HTML, CSS, or JavaScript artifact requests "
"with one complete self-contained HTML document in the code argument. "
"Call it once, then do not call it again in the same response unless "
"the user asks for changes. Future user requests for new artifacts may "
"call render_html once."
)
_sf_tool_tip_parts = []
if _sf_has_web:
_sf_tool_tip_parts.append(_sf_web_tips)
if _sf_has_code:
_sf_tool_tip_parts.append(_sf_code_tips)
if _sf_has_artifact:
_sf_tool_tip_parts.append(_sf_artifact_tips)
if _sf_tool_tip_parts:
_sf_nudge = (
_sf_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. " + " ".join(_sf_tool_tip_parts)
)
else:
_sf_nudge = ""
_sf_system_prompt = system_prompt
if _sf_nudge:
_sf_nudge += _tool_action_nudge(_sf_has_artifact)
if _sf_system_prompt:
_sf_system_prompt = _sf_system_prompt.rstrip() + "\n\n" + _sf_nudge
else:
_sf_system_prompt = _sf_nudge
_sf_auto_heal_tool_calls = (
payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True
)
# Strip stale tool-call XML from prior assistant turns.
_sf_chat_messages = []
for _msg in chat_messages:
@ -3290,7 +3258,10 @@ async def openai_chat_completions(
_sf_chat_messages.append(
{
**_msg,
"content": _TOOL_XML_RE.sub("", _msg["content"]).strip(),
"content": _strip_tool_xml_for_display(
_msg["content"],
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
).strip(),
}
)
else:
@ -3314,9 +3285,7 @@ async def openai_chat_completions(
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
preserve_thinking = payload.preserve_thinking,
auto_heal_tool_calls = payload.auto_heal_tool_calls
if payload.auto_heal_tool_calls is not None
else True,
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
max_tool_iterations = _sf_tool_budget,
tool_call_timeout = payload.tool_call_timeout
if payload.tool_call_timeout is not None
@ -3381,7 +3350,10 @@ async def openai_chat_completions(
# Diff cumulative cleaned text against last snapshot.
raw_cumulative = event.get("text", "")
clean_cumulative = _TOOL_XML_RE.sub("", raw_cumulative)
clean_cumulative = _strip_tool_xml_for_display(
raw_cumulative,
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
)
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
@ -3472,7 +3444,10 @@ async def openai_chat_completions(
if cancel_event.is_set():
break
if event.get("type") == "content":
full_text = _TOOL_XML_RE.sub("", event.get("text", ""))
full_text = _strip_tool_xml_for_display(
event.get("text", ""),
auto_heal_tool_calls = _sf_auto_heal_tool_calls,
)
return full_text
content_text = await asyncio.to_thread(_drain_to_text)
@ -4942,56 +4917,13 @@ async def anthropic_messages(
payload.enabled_tools,
)
# Build tool-use system prompt nudge (same as /chat/completions)
_tool_names = {t["function"]["name"] for t in openai_tools}
_has_web = "web_search" in _tool_names
_has_code = "python" in _tool_names or "terminal" in _tool_names
_has_artifact = "render_html" in _tool_names
_date_line = f"The current date is {_date.today().isoformat()}."
_model_size_b = _extract_model_size_b(model_name)
_is_small_model = _model_size_b is not None and _model_size_b < 9
if _is_small_model:
_web_tips = "Do not repeat the same search query."
else:
_web_tips = (
"When you search and find a relevant URL in the results, "
"fetch its full content by calling web_search with the url parameter. "
"Do not repeat the same search query. If a search returns "
"no useful results, try rephrasing or fetching a result URL directly."
)
_code_tips = (
"Use code execution for math, calculations, data processing, "
"or to parse and analyze information from tool results."
# Build tool-use system prompt nudge (same logic as /chat/completions)
_nudge = _build_tool_action_nudge(
tools = openai_tools,
model_name = model_name,
)
_artifact_tips = (
"Use render_html for HTML, CSS, or JavaScript artifact requests "
"with one complete self-contained HTML document in the code argument. "
"Call it once, then do not call it again in the same response unless "
"the user asks for changes. Future user requests for new artifacts may "
"call render_html once."
)
_tool_tip_parts = []
if _has_web:
_tool_tip_parts.append(_web_tips)
if _has_code:
_tool_tip_parts.append(_code_tips)
if _has_artifact:
_tool_tip_parts.append(_artifact_tips)
if _tool_tip_parts:
_nudge = (
_date_line + " "
"You have access to tools. When appropriate, prefer using "
"tools rather than answering from memory. " + " ".join(_tool_tip_parts)
)
else:
_nudge = ""
if _nudge:
_nudge += _tool_action_nudge(_has_artifact)
# Inject into system prompt
if openai_messages and openai_messages[0].get("role") == "system":
openai_messages[0]["content"] = (

46
studio/backend/tests/test_anthropic_messages.py
View file

@ -32,6 +32,7 @@ from core.inference.anthropic_compat import (
AnthropicPassthroughEmitter,
)
from routes.inference import (
_build_tool_action_nudge,
_normalize_anthropic_openai_images,
_select_anthropic_server_tools,
_anthropic_requested_studio_tools,
@ -46,6 +47,51 @@ from io import BytesIO as _BytesIO
from types import SimpleNamespace
# =====================================================================
# Tool nudge tests
# =====================================================================
class TestToolActionNudge:
def test_balanced_nudge_uses_expanded_web_and_code_tips(self):
nudge = _build_tool_action_nudge(
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "python"}},
],
model_name = "Llama-3.1-70B-Instruct",
)
assert nudge.startswith("The current date is ")
assert "Tools are available when they materially improve" in nudge
assert "prefer using tools rather than answering from memory" not in nudge
assert "fetch its full content by calling web_search with the url parameter" in nudge
assert "Use code execution for math" in nudge
assert "render_html" not in nudge
def test_balanced_nudge_preserves_compact_web_tip_and_artifact_gate(self):
nudge = _build_tool_action_nudge(
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "render_html"}},
],
model_name = "Llama-3.1-8B-Instruct",
)
assert "When using web_search, do not repeat the same search query." in nudge
assert "fetch its full content" not in nudge
assert "call render_html once" in nudge
def test_balanced_nudge_empty_without_known_tool_categories(self):
assert (
_build_tool_action_nudge(
tools = [],
model_name = "Llama-3.1-8B-Instruct",
)
== ""
)
# =====================================================================
# Pydantic model tests
# =====================================================================

224
studio/backend/tests/test_gguf_route_cursor_reset.py Normal file
View file

@ -0,0 +1,224 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGUF SSE cursor reset across an internal no-op tool decision.
When a GGUF turn streams visible preface text ("I'll render it again...")
and then hits an internal no-op (duplicate / disabled / repeat render_html),
no ``tool_start`` event is emitted. The GGUF SSE route diffs each cumulative
``content`` event against a ``prev_text`` cursor and resets that cursor on
``tool_start`` *and* on an empty ``status`` event. The generator must emit an
empty status between the preface turn and the final no-tools pass, otherwise
the final answer would be diffed against the stale preface and truncated or
dropped entirely.
This test drives the real generator and replays its events through a faithful
copy of the route's cursor loop (studio/backend/routes/inference.py), asserting
the final answer survives in full and the no-op produces no phantom tool card.
"""
from __future__ import annotations
import contextlib
import copy
import json
import sys
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.llama_cpp import LlamaCppBackend
def _sse(delta: dict) -> str:
return "data: " + json.dumps({"choices": [{"index": 0, "delta": delta}]}) + "\n"
def _done() -> str:
return "data: [DONE]\n"
def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
backend._port = 48848
backend._api_key = None
backend._effective_context_length = 4096
backend._supports_reasoning = False
backend._reasoning_always_on = False
backend._reasoning_style = "enable_thinking"
backend._supports_preserve_thinking = False
@contextlib.contextmanager
def fake_stream_with_retry(
_client,
_url,
payload,
_cancel_event,
headers = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
def fake_iter_text_cancellable(response, _cancel_event):
yield from response.chunks
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable)
return backend
def _replay_route_cursor(events: list[dict]) -> dict:
"""Replicate the GGUF SSE route's cumulative-cursor loop.
Mirrors studio/backend/routes/inference.py: reset ``prev_text`` on empty
status and on ``tool_start``; otherwise diff each cumulative ``content``
snapshot against the cursor and stream the delta. The preface/final text
here carry no tool XML, so the display strip is the identity -- the cursor
reset is the behaviour under test.
"""
prev_text = ""
visible_deltas: list[str] = []
tool_starts: list[dict] = []
statuses: list[str] = []
for event in events:
etype = event["type"]
if etype == "status":
if not event["text"]:
prev_text = ""
statuses.append(event["text"])
continue
if etype in ("tool_start", "tool_end"):
if etype == "tool_start":
prev_text = ""
tool_starts.append(event)
continue
if etype == "metadata":
continue
clean_cumulative = event.get("text", "")
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
continue
visible_deltas.append(new_text)
return {
"visible": "".join(visible_deltas),
"tool_starts": tool_starts,
"statuses": statuses,
}
def _replay_route_cursor_without_status_reset(events: list[dict]) -> dict:
"""Pre-fix control: identical to the route loop but never resets the
cursor on an empty status (only on ``tool_start``)."""
prev_text = ""
visible_deltas: list[str] = []
for event in events:
etype = event["type"]
if etype == "status":
continue
if etype in ("tool_start", "tool_end"):
if etype == "tool_start":
prev_text = ""
continue
if etype == "metadata":
continue
clean_cumulative = event.get("text", "")
new_text = clean_cumulative[len(prev_text) :]
prev_text = clean_cumulative
if not new_text:
continue
visible_deltas.append(new_text)
return {"visible": "".join(visible_deltas)}
def _web_search_tool() -> dict:
return {
"type": "function",
"function": {
"name": "web_search",
"description": "Search the web.",
"parameters": {
"type": "object",
"properties": {"query": {"type": "string"}},
"required": ["query"],
},
},
}
def test_final_answer_survives_preface_then_disabled_tool_noop(monkeypatch):
"""Preface text, then a call to a disabled tool (internal no-op).
A disabled-tool decision emits no ``tool_start`` and forces the final
no-tools pass. The route cursor must be reset before that pass so the
short final answer is not diffed away against the longer preface.
"""
preface = "Let me run a quick command to double-check."
final = "All set." # deliberately shorter than the preface -> truncation is visible
# Single turn: visible preface + a call to `terminal`, which is NOT in the
# enabled tool list, so the controller marks it disabled -> internal no-op.
turn_stream = [
_sse({"content": preface}),
_sse(
{
"tool_calls": [
{
"index": 0,
"id": "call_disabled",
"type": "function",
"function": {
"name": "terminal",
"arguments": json.dumps({"command": "ls"}),
},
}
]
}
),
_done(),
]
final_stream = [_sse({"content": final}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [turn_stream, final_stream], payloads)
executed: list[str] = []
monkeypatch.setattr(
"core.inference.tools.execute_tool",
lambda name, arguments, **_kw: executed.append(name) or "should-not-run",
)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "answer me"}],
tools = [_web_search_tool()], # terminal intentionally absent
temperature = 0.0,
max_tool_iterations = 5,
)
)
replay = _replay_route_cursor(events)
# Disabled tool is an internal no-op: never executed, no visible card.
assert executed == []
assert replay["tool_starts"] == []
# The generator must emit an empty status that resets the route cursor
# before the final pass; otherwise `final` (shorter than `preface`) would
# be diffed to nothing and dropped.
assert "" in replay["statuses"], "no cursor-resetting empty status emitted"
# Both the preface and the final answer survive, in order, untruncated.
assert preface in replay["visible"], replay["visible"]
assert final in replay["visible"], replay["visible"]
assert replay["visible"].index(preface) < replay["visible"].index(final)
assert replay["visible"].count(preface) == 1
# Negative control: a route loop that does NOT reset on empty status (the
# pre-fix behaviour) would diff `final` against the stale preface cursor
# and drop it -- proving the empty status is load-bearing here.
no_reset = _replay_route_cursor_without_status_reset(events)
assert final not in no_reset["visible"], no_reset["visible"]

1151
studio/backend/tests/test_llama_cpp_tool_loop.py Normal file
View file

File diff suppressed because it is too large Load diff

398
studio/backend/tests/test_safetensors_tool_loop.py
View file

@ -20,6 +20,7 @@ from core.inference.safetensors_agentic import (
_coerce_arguments,
_detect_render_html_tool_start,
run_safetensors_tool_loop,
strip_tool_markup_streaming,
)
from core.inference.tool_call_parser import (
has_tool_signal,
@ -53,6 +54,11 @@ class TestParser:
assert len(result) == 1
assert result[0]["function"]["name"] == "python"
def test_json_tool_call_unclosed_requires_healing(self):
text = '<tool_call>{"name":"python","arguments":{"code":"print(1)"}}'
assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
def test_xml_function_call(self):
text = "<function=python><parameter=code>print('hi')</parameter></function>"
result = parse_tool_calls_from_text(text)
@ -68,6 +74,11 @@ class TestParser:
assert result[0]["function"]["name"] == "terminal"
assert "ls -la" in result[0]["function"]["arguments"]
def test_xml_unclosed_requires_healing(self):
text = "<function=terminal><parameter=command>ls -la"
assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "terminal"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
def test_code_with_embedded_xml(self):
# A code parameter with a literal </parameter> must not truncate: the
# parser uses end-of-body as the only boundary for single-param calls.
@ -134,6 +145,23 @@ class TestParser:
# Without final=True the unclosed run is preserved.
assert "partial" in strip_tool_markup(text)
def test_streaming_strip_respects_disabled_healing(self):
raw = 'before <tool_call>{"name":"web_search"'
assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw
assert strip_tool_markup_streaming(raw) == "before "
def test_streaming_strip_respects_disabled_healing_without_tool_protocol(self):
raw = 'before <tool_call>{"name":"web_search"'
assert strip_tool_markup_streaming(raw, auto_heal_tool_calls = False) == raw
assert (
strip_tool_markup_streaming(
raw,
auto_heal_tool_calls = False,
tool_protocol_active = True,
)
== "before "
)
# ────────────────────────────────────────────────────────────────────
# run_safetensors_tool_loop
@ -230,6 +258,41 @@ def _make_loop(
), exec_fn
def test_active_tools_are_passed_to_single_turn_after_render_html_success():
captured_tool_names: list[list[str]] = []
exec_fn = FakeExecuteTool(["Rendered HTML artifact."])
def fake_single_turn(_messages, *, active_tools = None):
captured_tool_names.append(
[
(tool.get("function") or {}).get("name")
for tool in (active_tools or [])
if (tool.get("function") or {}).get("name")
]
)
if len(captured_tool_names) == 1:
yield '<tool_call>{"name":"render_html","arguments":{"code":"<html>one</html>"}}</tool_call>'
else:
yield "Done."
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "make html"}],
tools = [
{"type": "function", "function": {"name": "render_html"}},
{"type": "function", "function": {"name": "web_search"}},
],
execute_tool = exec_fn,
max_tool_iterations = 3,
)
)
assert exec_fn.calls == [("render_html", {"code": "<html>one</html>"})]
assert captured_tool_names == [["render_html", "web_search"], ["web_search"]]
assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events)
class TestLoopBasic:
def test_plain_answer(self):
# No tool XML; loop should yield content then status="".
@ -409,23 +472,158 @@ class TestLoopBasic:
class TestLoopBehaviour:
def test_duplicate_tool_call_synthetic_result(self):
# Two identical calls in a row: the second short-circuits with a "do not
# repeat" message; execute_tool runs once.
loop, exec_fn = _make_loop(
turns = [
def test_duplicate_tool_call_internal_noop(self):
captured_messages: list[list[dict]] = []
turns = iter(
[
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
["final"],
],
exec_results = ["search-result-1"],
]
)
def fake_single_turn(messages):
captured_messages.append([dict(message) for message in messages])
chunks = next(turns)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["search-result-1"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 3,
)
)
assert exec_fn.calls == [("web_search", {"query": "x"})]
assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == ["call_0"]
assert not [
e
for e in events
if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"}
]
duplicate_nudges = [
message
for message in captured_messages[-1]
if message.get("role") == "user"
and "already completed successfully" in message.get("content", "")
]
assert len(duplicate_nudges) == 1
def test_duplicate_tool_call_internal_noop_allows_distinct_followup_tool(self):
captured_messages: list[list[dict]] = []
captured_tool_names: list[list[str]] = []
turns = iter(
[
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>'],
["final"],
]
)
def fake_single_turn(messages, active_tools = None):
captured_messages.append([dict(message) for message in messages])
captured_tool_names.append(
[
tool["function"]["name"]
for tool in (active_tools or [])
if tool.get("function", {}).get("name")
]
)
chunks = next(turns)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["search-result-1", "python-result"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [
{"type": "function", "function": {"name": "web_search"}},
{"type": "function", "function": {"name": "python"}},
],
execute_tool = exec_fn,
max_tool_iterations = 4,
)
)
assert exec_fn.calls == [
("web_search", {"query": "x"}),
("python", {"code": "print(1)"}),
]
assert [e["tool_call_id"] for e in events if e["type"] == "tool_end"] == [
"call_0",
"call_2",
]
assert not [
e
for e in events
if e.get("tool_call_id") == "call_1" and e.get("type") in {"tool_start", "tool_end"}
]
duplicate_nudges = [
message
for message in captured_messages[2]
if message.get("role") == "user"
and "already completed successfully" in message.get("content", "")
]
assert len(duplicate_nudges) == 1
assert captured_tool_names[2] == ["web_search", "python"]
def test_repeated_duplicate_noop_transitions_to_final_attempt(self):
captured_tool_names: list[list[str]] = []
turns = iter(
[
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
["final from first result"],
]
)
def fake_single_turn(messages, active_tools = None):
captured_tool_names.append(
[
(tool.get("function") or {}).get("name")
for tool in (active_tools or [])
if (tool.get("function") or {}).get("name")
]
)
chunks = next(turns)
acc = ""
for chunk in chunks:
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["search-result"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 10,
)
)
assert exec_fn.calls == [("web_search", {"query": "x"})]
assert [
event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
] == ["call_0"]
assert captured_tool_names[-1] == []
assert any(
event.get("type") == "content" and "final from first result" in event.get("text", "")
for event in events
)
events = _collect_events(loop)
# Only one real call.
assert len(exec_fn.calls) == 1
tool_end_events = [e for e in events if e["type"] == "tool_end"]
assert len(tool_end_events) == 2
assert "do not repeat" in tool_end_events[1]["result"].lower()
def test_image_sentinel_stripped_from_model_feed(self):
# The image sentinel is stripped before the next turn, but tool_end still
@ -695,34 +893,62 @@ class TestChatTemplateHelper:
class TestGuardrails:
def test_disabled_tool_is_not_executed(self):
exec_fn = FakeExecuteTool([])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
['<tool_call>{"name":"terminal","arguments":{"command":"echo bypass"}}</tool_call>']
),
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
events = _collect_events(loop)
assert exec_fn.calls == []
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert tool_ends and "not enabled" in tool_ends[0]["result"].lower()
captured_messages: list[list[dict]] = []
def test_empty_tools_list_does_not_enforce_allowlist(self):
exec_fn = FakeExecuteTool(["OK"])
loop = run_safetensors_tool_loop(
single_turn = _fake_stream(
['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>']
),
messages = [{"role": "user", "content": "hi"}],
tools = [],
execute_tool = exec_fn,
max_tool_iterations = 2,
def fake_single_turn(messages):
captured_messages.append([dict(message) for message in messages])
if len(captured_messages) == 1:
yield '<tool_call>{"name":"terminal","arguments":{"command":"echo bypass"}}</tool_call>'
else:
yield "final"
exec_fn = FakeExecuteTool([])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
)
assert exec_fn.calls == []
assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}]
disabled_nudges = [
message
for message in captured_messages[-1]
if message.get("role") == "user" and "not enabled" in message.get("content", "")
]
assert len(disabled_nudges) == 1
def test_empty_tools_list_means_allow_all_in_core_loop(self):
turns = iter(
[
['<tool_call>{"name":"python","arguments":{"code":"print(1)"}}</tool_call>'],
["done"],
]
)
def fake_single_turn(_messages, active_tools = None):
assert active_tools == []
acc = ""
for chunk in next(turns):
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["OK"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "hi"}],
tools = [],
execute_tool = exec_fn,
max_tool_iterations = 2,
)
)
_collect_events(loop)
assert exec_fn.calls == [("python", {"code": "print(1)"})]
assert any(event.get("type") == "tool_end" for event in events)
def test_max_iterations_zero_executes_no_tools(self):
loop, exec_fn = _make_loop(
@ -767,6 +993,67 @@ class TestGuardrails:
_collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "x"})]
def test_auto_heal_disabled_preserves_xml_on_final_no_tools_pass(self):
turns = iter(
[
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'],
['<tool_call>{"name":"web_search","arguments":{"query":"literal"}}</tool_call>'],
]
)
def fake_single_turn(_messages, active_tools = None):
acc = ""
for chunk in next(turns):
acc += chunk
yield acc
exec_fn = FakeExecuteTool(["OK"])
events = _collect_events(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "show literal"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
execute_tool = exec_fn,
max_tool_iterations = 1,
auto_heal_tool_calls = False,
)
)
assert exec_fn.calls == [("web_search", {"query": "x"})]
assert any(
event.get("type") == "content" and "<tool_call>" in event.get("text", "")
for event in events
)
def test_auto_heal_disabled_does_not_repair_unclosed_tool_call(self):
loop, exec_fn = _make_loop(
turns = [
['<tool_call>{"name":"web_search","arguments":{"query":"x"}}'],
],
exec_results = ["OK"],
auto_heal_tool_calls = False,
max_tool_iterations = 1,
)
events = _collect_events(loop)
assert exec_fn.calls == []
assert any(
event.get("type") == "content" and "<tool_call>" in event.get("text", "")
for event in events
)
def test_auto_heal_enabled_strips_unparseable_xml_tool_call(self):
loop, exec_fn = _make_loop(
turns = [["<tool_call>{not valid json}</tool_call>"]],
exec_results = ["OK"],
auto_heal_tool_calls = True,
max_tool_iterations = 1,
)
events = _collect_events(loop)
assert exec_fn.calls == []
assert not any(
event.get("type") == "content" and "<tool_call>" in event.get("text", "")
for event in events
)
def test_non_consecutive_duplicate_is_short_circuited(self):
loop, exec_fn = _make_loop(
turns = [
@ -780,8 +1067,39 @@ class TestGuardrails:
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "A"}), ("web_search", {"query": "B"})]
tool_ends = [e for e in events if e["type"] == "tool_end"]
assert "already made this exact call" in tool_ends[-1]["result"]
assert [
event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
] == ["call_0", "call_1"]
assert not [
event
for event in events
if event.get("tool_call_id") == "call_2"
and event.get("type") in {"tool_start", "tool_end"}
]
def test_same_turn_duplicate_is_short_circuited(self):
loop, exec_fn = _make_loop(
turns = [
[
'<tool_call>{"name":"web_search","arguments":{"query":"A"}}</tool_call>'
'<tool_call>{"name":"web_search","arguments":{"query":"A"}}</tool_call>'
],
["final"],
],
exec_results = ["res-A"],
max_tool_iterations = 2,
)
events = _collect_events(loop)
assert exec_fn.calls == [("web_search", {"query": "A"})]
assert [
event.get("tool_call_id") for event in events if event.get("type") == "tool_end"
] == ["call_0"]
assert not [
event
for event in events
if event.get("tool_call_id") == "call_1"
and event.get("type") in {"tool_start", "tool_end"}
]
def test_coerce_string_args_python_uses_code_key(self):
assert _coerce_arguments("print(1)", heal = True, tool_name = "python") == {"code": "print(1)"}

114
studio/backend/tests/test_tool_call_parser_strict.py Normal file
View file

@ -0,0 +1,114 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Strict-mode (Auto-Heal disabled) tool-call parsing.
With ``allow_incomplete=False`` the parser must accept a well-formed
``<function=...>...</function>`` call even when the model appends prose
after the closing tag -- matching the JSON-style ``<tool_call>...`` path,
which already tolerates trailing text -- while still rejecting genuinely
truncated calls that never close.
"""
from __future__ import annotations
import json
import sys
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.tool_call_parser import parse_tool_calls_from_text
def _only(text: str) -> dict:
calls = parse_tool_calls_from_text(text, allow_incomplete = False)
assert len(calls) == 1, f"expected exactly one call, got {len(calls)}: {calls!r}"
fn = calls[0]["function"]
return {"name": fn["name"], "arguments": json.loads(fn["arguments"])}
class TestFunctionStyleTrailingText:
def test_closed_function_with_trailing_prose_is_accepted(self):
text = (
"<function=web_search><parameter=query>weather london</parameter></function>"
" Let me check that for you."
)
call = _only(text)
assert call == {"name": "web_search", "arguments": {"query": "weather london"}}
def test_closed_function_with_trailing_whitespace_is_accepted(self):
text = "<function=web_search><parameter=query>cats</parameter></function> \n\n"
call = _only(text)
assert call == {"name": "web_search", "arguments": {"query": "cats"}}
def test_closed_function_without_trailing_text_still_parses(self):
text = "<function=web_search><parameter=query>cats</parameter></function>"
call = _only(text)
assert call == {"name": "web_search", "arguments": {"query": "cats"}}
def test_multi_param_with_trailing_prose(self):
text = (
"<function=terminal><parameter=command>ls -la</parameter>"
"<parameter=workdir>home</parameter></function> running it now"
)
call = _only(text)
assert call == {
"name": "terminal",
"arguments": {"command": "ls -la", "workdir": "home"},
}
def test_code_value_containing_literal_close_tag_is_preserved(self):
# The real closing </function> is the last one; the literal inside
# the code argument must survive (rfind, not the first match).
text = (
"<function=python><parameter=code>"
'print("</function>")'
"</parameter></function> all done"
)
call = _only(text)
assert call == {"name": "python", "arguments": {"code": 'print("</function>")'}}
def test_incomplete_function_without_close_is_still_rejected(self):
text = "<function=web_search><parameter=query>weather london"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
def test_param_without_close_tag_is_rejected_in_strict_mode(self):
# Closing </function> present, but the single parameter never closes.
text = "<function=web_search><parameter=query>weather london</function>"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
class TestParityWithJsonStyle:
def test_json_tool_call_with_trailing_prose_is_accepted(self):
text = (
'<tool_call>{"name":"web_search","arguments":{"query":"weather london"}}</tool_call>'
" Let me check that for you."
)
calls = parse_tool_calls_from_text(text, allow_incomplete = False)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "web_search"
def test_function_and_json_styles_agree_on_trailing_text(self):
q = "weather london"
func = parse_tool_calls_from_text(
f"<function=web_search><parameter=query>{q}</parameter></function> trailing",
allow_incomplete = False,
)
js = parse_tool_calls_from_text(
f'<tool_call>{{"name":"web_search","arguments":{{"query":"{q}"}}}}</tool_call> trailing',
allow_incomplete = False,
)
assert len(func) == len(js) == 1
assert json.loads(func[0]["function"]["arguments"]) == {"query": q}
assert json.loads(js[0]["function"]["arguments"]) == {"query": q}
class TestHealingPathUnaffected:
def test_auto_heal_still_repairs_unclosed_function(self):
text = "<function=web_search><parameter=query>cats"
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "web_search"

212
studio/backend/tests/test_tool_loop_controller.py Normal file
View file

@ -0,0 +1,212 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import json
import sys
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference.tool_loop_controller import (
ToolLoopController,
canonical_tool_call_key,
coerce_tool_arguments,
status_for_tool,
strip_result_for_model,
tool_event_provenance,
)
def _tool(name: str) -> dict:
return {"type": "function", "function": {"name": name}}
def _call(
name: str,
args,
call_id: str = "call_0",
) -> dict:
return {
"id": call_id,
"type": "function",
"function": {
"name": name,
"arguments": json.dumps(args) if isinstance(args, dict) else args,
},
}
def test_canonical_tool_call_key_sorts_arguments():
a = canonical_tool_call_key("web_search", {"query": "gpu", "limit": 5})
b = canonical_tool_call_key("web_search", {"limit": 5, "query": "gpu"})
c = canonical_tool_call_key("python", {"limit": 5, "query": "gpu"})
assert a == b
assert a != c
assert a == 'web_search:{"limit":5,"query":"gpu"}'
def test_coerce_tool_arguments_parses_json_and_heals_raw_strings():
parsed = coerce_tool_arguments('{"query":"gpu prices"}', heal = True)
healed = coerce_tool_arguments("print(1)", heal = True, tool_name = "python")
raw = coerce_tool_arguments("not-json", heal = False, tool_name = "python")
assert parsed.arguments == {"query": "gpu prices"}
assert not parsed.healed
assert healed.arguments == {"code": "print(1)"}
assert healed.healed
assert raw.arguments == {"raw": "not-json"}
assert not raw.healed
def test_status_and_provenance_match_local_event_conventions():
assert status_for_tool("web_search", {"query": "gpus"}) == "Searching: gpus"
assert (
status_for_tool("web_search", {"url": "https://www.example.com/a"})
== "Reading: example.com"
)
assert status_for_tool("python", {"code": "print(1)\nprint(2)"}) == "Running Python: print(1)"
assert tool_event_provenance(healed = True, forced = False, provisional = None) == {
"source": "local",
"healed": True,
}
def test_prepare_execute_builds_visible_events_and_model_tool_message():
controller = ToolLoopController(tools = [_tool("web_search")])
decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
assert decision.should_execute
assert decision.emit_visible_events
assert decision.status_text == "Searching: gpu prices"
assert decision.tool_start_payload()["arguments"] == {"query": "gpu prices"}
assert decision.tool_start_event()["type"] == "tool_start"
assert decision.as_assistant_tool_call()["function"]["arguments"] == '{"query":"gpu prices"}'
completion = controller.record_result(decision, "Search result\n__IMAGES__:{...}")
assert completion.tool_end_payload()["result"] == "Search result\n__IMAGES__:{...}"
assert completion.tool_end_event()["type"] == "tool_end"
assert completion.tool_message() == {
"role": "tool",
"name": "web_search",
"content": "Search result",
"tool_call_id": "call_0",
}
def test_successful_duplicate_is_internal_noop_and_keeps_remaining_tools():
controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")])
first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a"))
controller.record_result(first, "ok")
duplicate = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b"))
completion = controller.record_noop(duplicate)
assert duplicate.action == "duplicate"
assert not duplicate.should_execute
assert not duplicate.emit_visible_events
duplicate_nudge = completion.model_message()["content"]
assert "already completed successfully" in duplicate_nudge
assert "different enabled tool" in duplicate_nudge
assert completion.model_message()["role"] == "user"
assert not controller.force_final_answer
assert [tool["function"]["name"] for tool in controller.active_tools()] == [
"web_search",
"python",
]
def test_repeated_successful_duplicate_becomes_terminal_after_one_recovery_nudge():
controller = ToolLoopController(tools = [_tool("web_search"), _tool("python")])
first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_a"))
controller.record_result(first, "ok")
duplicate_one = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_b"))
completion_one = controller.record_noop(duplicate_one)
assert duplicate_one.action == "duplicate"
assert "already completed successfully" in completion_one.model_message()["content"]
assert not controller.force_final_answer
assert [tool["function"]["name"] for tool in controller.active_tools()] == [
"web_search",
"python",
]
duplicate_two = controller.prepare_call(_call("web_search", {"query": "gpu prices"}, "call_c"))
completion_two = controller.record_noop(duplicate_two)
assert duplicate_two.action == "duplicate"
assert "already completed successfully" in completion_two.model_message()["content"]
assert controller.force_final_answer
assert controller.active_tools() == []
def test_failed_call_does_not_block_retry():
controller = ToolLoopController(tools = [_tool("web_search")])
first = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
controller.record_result(first, "Error: temporary failure")
retry = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
assert retry.should_execute
assert retry.action == "execute"
def test_empty_enabled_tool_list_blocks_all_tool_calls():
controller = ToolLoopController(tools = [])
decision = controller.prepare_call(_call("web_search", {"query": "gpu prices"}))
completion = controller.record_noop(decision)
assert decision.action == "disabled"
assert not decision.emit_visible_events
assert completion.model_message()["role"] == "user"
assert "not enabled" in completion.model_message()["content"]
assert controller.force_final_answer
assert controller.active_tools() == []
def test_disabled_tool_is_internal_noop_not_visible_tool_error():
controller = ToolLoopController(tools = [_tool("web_search")])
decision = controller.prepare_call(_call("python", {"code": "print(1)"}))
completion = controller.record_noop(decision)
assert decision.action == "disabled"
assert not decision.emit_visible_events
assert completion.model_message()["role"] == "user"
assert "not enabled" in completion.model_message()["content"]
assert controller.force_final_answer
assert controller.active_tools() == []
def test_render_html_success_filters_active_tools_and_repeat_is_internal():
controller = ToolLoopController(tools = [_tool("render_html"), _tool("web_search")])
assert [t["function"]["name"] for t in controller.active_tools()] == [
"render_html",
"web_search",
]
first = controller.prepare_call(_call("render_html", {"code": "<html></html>"}, "call_html_1"))
controller.record_result(first, "Rendered HTML artifact: Demo")
assert [t["function"]["name"] for t in controller.active_tools()] == ["web_search"]
repeat = controller.prepare_call(_call("render_html", {"code": "<html></html>"}, "call_html_2"))
completion = controller.record_noop(repeat)
assert repeat.action == "render_html_repeat"
assert not repeat.emit_visible_events
assert completion.model_message()["role"] == "user"
assert "Do not call render_html again" in completion.model_message()["content"]
assert controller.force_final_answer
assert controller.active_tools() == []
def test_strip_result_for_model_removes_frontend_image_sentinel():
assert strip_result_for_model('text\n__IMAGES__:{"paths":[]}') == "text"
assert strip_result_for_model("text __IMAGES__:payload") == "text"
assert strip_result_for_model("plain text") == "plain text"

14
studio/backend/tests/test_tool_xml_strip.py
View file

@ -27,11 +27,25 @@ assert _m, "could not extract _TOOL_XML_RE source"
_ns = {"_re": _re}
exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns)
_TOOL_XML_RE = _ns["_TOOL_XML_RE"]
_helper = _re.search(
r"def _strip_tool_xml_for_display\(text: str, \*, auto_heal_tool_calls: bool\) -> str:\n"
r"(?: .+\n)+",
_src,
)
assert _helper, "could not extract _strip_tool_xml_for_display source"
exec(_helper.group(0), _ns)
_strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"]
# ── Well-formed pairs ─────────────────────────────────────────────
def test_route_display_strip_respects_disabled_auto_heal_contract():
text = 'literal <tool_call>{"name":"web_search"}</tool_call> survives'
assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text
assert "<tool_call>" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True)
def test_strips_well_formed_tool_call():
text = (
"Let me search.\n"

736
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -252,8 +252,7 @@ function documentCitationToSource(
title: string;
metadata?: { description: string };
} | null {
const source =
typeof cit.source === "string" && cit.source ? cit.source : "";
const source = typeof cit.source === "string" && cit.source ? cit.source : "";
const docTitle =
(typeof cit.document_title === "string" && cit.document_title) ||
(typeof cit.title === "string" && cit.title) ||
@ -264,13 +263,12 @@ function documentCitationToSource(
// search_result_location can carry a free-form id (e.g. ``kb-doc-42``)
// or a hostile scheme. Fall back to a stable doc anchor otherwise.
const url =
isSafeNavigableSourceUrl(source) || `#anthropic-doc-${docIndex ?? fallbackIdx}`;
isSafeNavigableSourceUrl(source) ||
`#anthropic-doc-${docIndex ?? fallbackIdx}`;
const title = docTitle || source || `Document ${fallbackIdx + 1}`;
const cited =
typeof cit.cited_text === "string" ? cit.cited_text.trim() : "";
const cited = typeof cit.cited_text === "string" ? cit.cited_text.trim() : "";
// Trim the cited snippet so the Sources panel stays scannable.
const description =
cited.length > 240 ? `${cited.slice(0, 240)}...` : cited;
const description = cited.length > 240 ? `${cited.slice(0, 240)}...` : cited;
// Anthropic numbers inline [N] per citation, not per source URL.
// Fold citation type + position-bearing fields into the id so distinct
// citations on the same source keep separate Sources entries.
@ -565,154 +563,6 @@ function isAnthropicRefusalMessage(message: RunMessage): boolean {
return metadata?.custom?.anthropicRefusal === true;
}
function collectAssistantToolCalls(
message: RunMessage,
): Array<{
id: string;
type: "function";
function: { name: string; arguments: string };
extra_content?: unknown;
}> {
const out: Array<{
id: string;
type: "function";
function: { name: string; arguments: string };
extra_content?: unknown;
}> = [];
for (const part of message.content ?? []) {
if (part.type !== "tool-call") continue;
const tc = part as ToolCallMessagePart & {
argsText?: string;
extra_content?: unknown;
};
const toolNameLower = (tc.toolName ?? "").toLowerCase();
const argsObj =
tc.args && typeof tc.args === "object"
? (tc.args as Record<string, unknown>)
: null;
const argsGoogle =
argsObj && typeof argsObj.google === "object" && argsObj.google !== null
? (argsObj.google as Record<string, unknown>)
: null;
const hasNativePart = Boolean(
argsGoogle &&
typeof argsGoogle.native_part === "object" &&
argsGoogle.native_part !== null,
);
const hasServerToolMarker = Boolean(
argsObj && (argsObj as Record<string, unknown>)._server_tool === true,
);
const isServerSideBuiltin = isServerSideBuiltinToolPart(
toolNameLower,
argsObj,
hasServerToolMarker,
hasNativePart,
);
if (isServerSideBuiltin) {
// Gemini code_execution / image_generation must round-trip the
// native_part payload for native replay; drop the rest.
if (!hasNativePart) continue;
}
const argumentsStr =
typeof tc.argsText === "string" && tc.argsText.length > 0
? tc.argsText
: JSON.stringify(tc.args ?? {});
const entry: {
id: string;
type: "function";
function: { name: string; arguments: string };
extra_content?: unknown;
} = {
id: tc.toolCallId,
type: "function" as const,
function: {
name: tc.toolName ?? "",
arguments: argumentsStr,
},
};
// Promote args.google to extra_content.google so the backend
// native_part replay branch finds it (it only inspects
// extra_content, not function.arguments).
if (tc.extra_content !== undefined) {
entry.extra_content = tc.extra_content;
} else if (argsGoogle) {
entry.extra_content = { google: argsGoogle };
}
out.push(entry);
}
return out;
}
function collectToolResultMessages(
message: RunMessage,
): Array<{
role: "tool";
content: string;
tool_call_id: string;
name?: string;
}> {
const out: Array<{
role: "tool";
content: string;
tool_call_id: string;
name?: string;
}> = [];
for (const part of message.content ?? []) {
if (part.type !== "tool-call") continue;
const tc = part as ToolCallMessagePart;
const result = (tc as { result?: unknown }).result;
// Skip provider-side builtins; see isServerSideBuiltinToolPart().
const argsObj =
tc.args && typeof tc.args === "object"
? (tc.args as Record<string, unknown>)
: null;
const argsGoogle =
argsObj && typeof argsObj.google === "object" && argsObj.google !== null
? (argsObj.google as Record<string, unknown>)
: null;
const toolNameLower = (tc.toolName ?? "").toLowerCase();
const hasServerToolMarker = Boolean(
argsObj && argsObj._server_tool === true,
);
const hasNativePart = Boolean(
argsGoogle &&
typeof argsGoogle.native_part === "object" &&
argsGoogle.native_part !== null,
);
if (
isServerSideBuiltinToolPart(
toolNameLower,
argsObj,
hasServerToolMarker,
hasNativePart,
)
) {
continue;
}
if (result === undefined || result === null) continue;
let content: string;
if (typeof result === "string") {
// Backend ChatMessage validator rejects role="tool" with empty
// content; serialise sentinel JSON so legitimately empty tool
// outputs still round-trip to the provider.
content = result.length > 0 ? result : JSON.stringify({ result: "" });
} else {
try {
content = JSON.stringify(result);
} catch {
content = String(result);
}
}
out.push({
role: "tool",
content,
tool_call_id: tc.toolCallId,
...(tc.toolName ? { name: tc.toolName } : {}),
});
}
return out;
}
type SerializedMessage = {
role: "system" | "user" | "assistant" | "tool";
content: OpenAIMessageContent | null;
@ -733,6 +583,185 @@ type SerializedMessage = {
extra_content?: unknown;
};
type SerializedToolCall = NonNullable<SerializedMessage["tool_calls"]>[number];
type SerializedToolResult = {
role: "tool";
content: string;
tool_call_id: string;
name?: string;
};
type ToolPartReplayMetadata = {
argsObj: Record<string, unknown> | null;
argsGoogle: Record<string, unknown> | null;
hasNativePart: boolean;
isServerSideBuiltin: boolean;
};
function getToolPartReplayMetadata(
tc: ToolCallMessagePart,
): ToolPartReplayMetadata {
const toolNameLower = (tc.toolName ?? "").toLowerCase();
const argsObj =
tc.args && typeof tc.args === "object"
? (tc.args as Record<string, unknown>)
: null;
const argsGoogle =
argsObj && typeof argsObj.google === "object" && argsObj.google !== null
? (argsObj.google as Record<string, unknown>)
: null;
const hasNativePart = Boolean(
argsGoogle &&
typeof argsGoogle.native_part === "object" &&
argsGoogle.native_part !== null,
);
const hasServerToolMarker = Boolean(
argsObj && (argsObj as Record<string, unknown>)._server_tool === true,
);
return {
argsObj,
argsGoogle,
hasNativePart,
isServerSideBuiltin: isServerSideBuiltinToolPart(
toolNameLower,
argsObj,
hasServerToolMarker,
hasNativePart,
),
};
}
type ToolReplayProvenance = {
source?: string;
[key: string]: unknown;
};
function getToolReplayProvenance(
part: ToolCallMessagePart,
): ToolReplayProvenance | null {
const provenance = (part as { provenance?: unknown }).provenance;
if (
!provenance ||
typeof provenance !== "object" ||
Array.isArray(provenance)
) {
return null;
}
return provenance as ToolReplayProvenance;
}
function hasToolReplayResult(part: ToolCallMessagePart): boolean {
const result = (part as { result?: unknown }).result;
return result !== undefined && result !== null;
}
function shouldFlushCompletedLocalToolPair(part: ToolCallMessagePart): boolean {
const provenance = getToolReplayProvenance(part);
if (provenance?.source !== "local") {
return false;
}
if (getToolPartReplayMetadata(part).isServerSideBuiltin) {
return false;
}
return hasToolReplayResult(part);
}
function serializeAssistantToolCallPart(
part: ToolCallMessagePart,
): SerializedToolCall | null {
const tc = part as ToolCallMessagePart & {
argsText?: string;
extra_content?: unknown;
};
const { argsGoogle, hasNativePart, isServerSideBuiltin } =
getToolPartReplayMetadata(tc);
if (isServerSideBuiltin && !hasNativePart) {
return null;
}
const argumentsStr =
typeof tc.argsText === "string" && tc.argsText.length > 0
? tc.argsText
: JSON.stringify(tc.args ?? {});
const entry: SerializedToolCall = {
id: tc.toolCallId,
type: "function" as const,
function: {
name: tc.toolName ?? "",
arguments: argumentsStr,
},
};
// Promote args.google to extra_content.google so the backend
// native_part replay branch can find it. The backend only inspects
// extra_content, not function.arguments.
if (tc.extra_content !== undefined) {
entry.extra_content = tc.extra_content;
} else if (argsGoogle) {
entry.extra_content = { google: argsGoogle };
}
return entry;
}
function serializeToolResultPart(
part: ToolCallMessagePart,
): SerializedToolResult | null {
const tc = part as ToolCallMessagePart;
const result = (tc as { result?: unknown }).result;
const { isServerSideBuiltin } = getToolPartReplayMetadata(tc);
// Skip provider-side builtins; see isServerSideBuiltinToolPart.
if (isServerSideBuiltin) {
return null;
}
if (result === undefined || result === null) return null;
let content: string;
if (typeof result === "string") {
// Backend ChatMessage validator rejects role="tool" with empty
// content; serialise a sentinel JSON so legitimately empty tool
// outputs still round-trip the follow-up turn to the provider.
content = result.length > 0 ? result : JSON.stringify({ result: "" });
} else {
try {
content = JSON.stringify(result);
} catch {
content = String(result);
}
}
return {
role: "tool" as const,
content,
tool_call_id: tc.toolCallId,
...(tc.toolName ? { name: tc.toolName } : {}),
};
}
function canReplayToolCallWithoutRoleTool(part: ToolCallMessagePart): boolean {
// Gemini/OpenAI provider-native builtin cards replay through
// extra_content/native parts and intentionally do not produce role="tool"
// messages. Local/user tool calls must have a concrete tool result before
// they are replayed ahead of later assistant text.
return getToolPartReplayMetadata(part).isServerSideBuiltin;
}
function sanitizeAssistantReplayText(text: string): string {
return text.replace(
/data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g,
"[audio]",
);
}
function buildReplayContent(
textContent: string,
imageParts: Array<{ type: "image_url"; image_url: { url: string } }>,
): OpenAIMessageContent {
return imageParts.length > 0
? [{ type: "text", text: textContent }, ...imageParts]
: textContent;
}
function collectAssistantTextThoughtSignature(
message: RunMessage,
): string | undefined {
@ -749,6 +778,132 @@ function collectAssistantTextThoughtSignature(
return undefined;
}
function attachAssistantThoughtSignature(
messages: SerializedMessage[],
thoughtSignature: string | undefined,
): void {
if (!thoughtSignature) return;
for (let i = messages.length - 1; i >= 0; i -= 1) {
const message = messages[i];
if (message.role !== "assistant") continue;
const extra =
message.extra_content &&
typeof message.extra_content === "object" &&
!Array.isArray(message.extra_content)
? (message.extra_content as Record<string, unknown>)
: {};
const google =
extra.google &&
typeof extra.google === "object" &&
!Array.isArray(extra.google)
? (extra.google as Record<string, unknown>)
: {};
message.extra_content = {
...extra,
google: { ...google, thought_signature: thoughtSignature },
};
return;
}
}
function serializeAssistantReplayMessages(
message: RunMessage,
): SerializedMessage[] {
if (isAnthropicRefusalMessage(message)) {
// Prune refused assistant turn from outbound history; the
// rendered transcript still shows the user-visible notice.
return [];
}
const imageParts = collectImageParts(message);
const messages: SerializedMessage[] = [];
const pendingTextParts: string[] = [];
let pendingToolCalls: SerializedToolCall[] = [];
let pendingToolResults: SerializedToolResult[] = [];
let imagePartsPending = imageParts.length > 0;
const flushAssistantAndToolResults = (force = false): void => {
const textContent = sanitizeAssistantReplayText(
pendingTextParts.join("\n"),
);
const includeImageParts = imagePartsPending ? imageParts : [];
const hasContent = textContent.length > 0 || includeImageParts.length > 0;
const hasToolCalls = pendingToolCalls.length > 0;
if (!force && !hasContent && !hasToolCalls) {
return;
}
const assistantMessage: SerializedMessage = {
role: "assistant",
content: hasContent
? buildReplayContent(textContent, includeImageParts)
: "",
};
if (hasToolCalls) {
assistantMessage.tool_calls = pendingToolCalls;
// OpenAI requires content === null on assistant turns whose
// payload is entirely tool_calls (matches the wire shape Gemini
// expects for the next functionCall replay).
if (!hasContent) {
assistantMessage.content = null;
}
}
messages.push(assistantMessage);
if (pendingToolResults.length > 0) {
messages.push(...pendingToolResults);
}
pendingTextParts.length = 0;
pendingToolCalls = [];
pendingToolResults = [];
imagePartsPending = false;
};
for (const part of message.content ?? []) {
if (part.type === "text") {
if (pendingToolCalls.length > 0) {
flushAssistantAndToolResults();
}
pendingTextParts.push(part.text);
continue;
}
if (part.type === "tool-call") {
const toolPart = part as ToolCallMessagePart;
const toolCall = serializeAssistantToolCallPart(toolPart);
if (!toolCall) continue;
const toolResult = serializeToolResultPart(toolPart);
if (!toolResult && !canReplayToolCallWithoutRoleTool(toolPart)) {
continue;
}
const flushLocalPair = shouldFlushCompletedLocalToolPair(toolPart);
if (flushLocalPair && pendingToolCalls.length > 0) {
flushAssistantAndToolResults();
}
pendingToolCalls.push(toolCall);
if (toolResult) {
pendingToolResults.push(toolResult);
}
if (flushLocalPair) {
flushAssistantAndToolResults();
}
}
}
flushAssistantAndToolResults(messages.length === 0);
attachAssistantThoughtSignature(
messages,
collectAssistantTextThoughtSignature(message),
);
return messages;
}
function toOpenAIMessages(message: RunMessage): SerializedMessage[] {
if (
message.role !== "system" &&
@ -758,49 +913,18 @@ function toOpenAIMessages(message: RunMessage): SerializedMessage[] {
return [];
}
let textContent = collectTextParts(message).join("\n");
if (message.role === "assistant") {
textContent = textContent.replace(
/data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g,
"[audio]",
);
if (isAnthropicRefusalMessage(message)) {
// Prune refused assistant turn from outbound history; the rendered
// transcript still shows the user-visible notice.
return [];
}
return serializeAssistantReplayMessages(message);
}
const textContent = collectTextParts(message).join("\n");
const imageParts = collectImageParts(message);
const toolCalls =
message.role === "assistant" ? collectAssistantToolCalls(message) : [];
const toolResults =
message.role === "assistant" ? collectToolResultMessages(message) : [];
const base: SerializedMessage = {
role: message.role,
content:
imageParts.length > 0
? [{ type: "text", text: textContent }, ...imageParts]
: textContent,
};
if (toolCalls.length > 0) {
base.tool_calls = toolCalls;
// OpenAI requires content === null on assistant turns that are
// entirely tool_calls (matches the wire shape Gemini expects for
// the next functionCall replay).
if (!textContent && imageParts.length === 0) {
base.content = null;
}
}
if (message.role === "assistant") {
const sig = collectAssistantTextThoughtSignature(message);
if (sig) {
base.extra_content = { google: { thought_signature: sig } };
}
}
return toolResults.length > 0 ? [base, ...toolResults] : [base];
return [
{
role: message.role,
content: buildReplayContent(textContent, imageParts),
},
];
}
function extractImageBase64(input: string): string | undefined {
@ -1310,7 +1434,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
let loaded: boolean;
let blockedByTrustRemoteCode: boolean;
try {
({ loaded, blockedByTrustRemoteCode } = await autoLoadSmallestModel());
({ loaded, blockedByTrustRemoteCode } =
await autoLoadSmallestModel());
} catch (error) {
clearSelectedImageEditReference();
throw error;
@ -1772,8 +1897,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
// <think>...</think> for parseAssistantContent. Lives outside the
// SSE loop because the close tag fires when content arrives.
let reasoningContentOpen = false;
type ToolCallProvenance = {
source?: string;
healed?: boolean;
forced?: boolean;
provisional?: boolean;
duplicate?: boolean;
reason?: string;
[key: string]: unknown;
};
type PositionedToolCallPart = ToolCallMessagePart & {
textCursor?: number;
_delta_index?: number;
extra_content?: unknown;
provenance?: ToolCallProvenance;
};
// Tool call parts, cumulative; result lands on tool_end.
const toolCallParts: ToolCallMessagePart[] = [];
const toolCallParts: PositionedToolCallPart[] = [];
// Latest Gemini text-part thoughtSignature; pinned onto the final
// text MessagePart so next-turn replay carries it.
let latestTextThoughtSignature: string | undefined;
@ -1792,16 +1932,81 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
return parts;
};
const orderAssistantContent = (
textParts: ReturnType<typeof parseAssistantContent>,
) => {
const imageToolParts = toolCallParts.filter(
(part) => part.toolName === "image_generation",
);
const otherToolParts = toolCallParts.filter(
(part) => part.toolName !== "image_generation",
);
return [...otherToolParts, ...textParts, ...imageToolParts];
const buildAssistantContent = (rawText: string) => {
const positionedTools = toolCallParts
.map((part, index) => {
const cursor = (part as PositionedToolCallPart).textCursor;
return {
part,
index,
cursor:
typeof cursor === "number" && Number.isFinite(cursor)
? Math.min(Math.max(cursor, 0), rawText.length)
: 0,
};
})
.sort((a, b) => a.cursor - b.cursor || a.index - b.index);
const assembled: Array<
ReturnType<typeof parseAssistantContent>[number] | ToolCallMessagePart
> = [];
let textCursor = 0;
let toolIndex = 0;
const appendTextThrough = (nextCursor: number) => {
if (nextCursor <= textCursor) return;
assembled.push(
...parseAssistantContent(rawText.slice(textCursor, nextCursor)),
);
textCursor = nextCursor;
};
while (toolIndex < positionedTools.length) {
const cursor = positionedTools[toolIndex].cursor;
appendTextThrough(cursor);
while (
toolIndex < positionedTools.length &&
positionedTools[toolIndex].cursor === cursor
) {
assembled.push(positionedTools[toolIndex].part);
toolIndex += 1;
}
}
appendTextThrough(rawText.length);
return pinTextThoughtSignature(assembled);
};
const parseToolProvenance = (
value: unknown,
): ToolCallProvenance | undefined => {
if (!value || typeof value !== "object" || Array.isArray(value)) {
return undefined;
}
return { ...(value as Record<string, unknown>) } as ToolCallProvenance;
};
const mergeToolProvenance = (
existing: ToolCallProvenance | undefined,
incoming: ToolCallProvenance | undefined,
): ToolCallProvenance | undefined => {
if (!incoming) return existing;
if (!existing) return incoming;
const merged: ToolCallProvenance = { ...existing, ...incoming };
for (const key of [
"healed",
"forced",
"provisional",
"duplicate",
] as const) {
if (existing[key] === true || incoming[key] === true) {
merged[key] = true;
}
}
return merged;
};
const closeReasoningContent = () => {
if (!reasoningContentOpen) return;
cumulativeText += "</think>";
reasoningContentOpen = false;
};
// Anthropic document_citations payload, converted to Sources-panel
// parts at end-of-stream so inline [N] markers have matching entries.
@ -2189,7 +2394,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
try {
let requestPayload: OpenAIChatCompletionsRequest;
try {
requestPayload = await buildRequestPayload(retriedWithRefreshedKey);
requestPayload = await buildRequestPayload(
retriedWithRefreshedKey,
);
} catch (error) {
clearSelectedImageEditReference();
throw error;
@ -2270,6 +2477,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
anthropicRefusalSeen = true;
continue;
}
closeReasoningContent();
const toolProvenance = parseToolProvenance(
toolEvent.provenance,
);
if (toolEvent.type === "tool_start") {
const id =
(toolEvent.tool_call_id as string) ||
@ -2280,11 +2491,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
(p) => p.toolCallId === id,
);
if (idx !== -1) {
const existing = toolCallParts[
idx
] as PositionedToolCallPart;
toolCallParts[idx] = {
...toolCallParts[idx],
...existing,
toolName: toolEvent.tool_name as string,
argsText: JSON.stringify(toolArgs),
args: toolArgs,
provenance: mergeToolProvenance(
existing.provenance,
toolProvenance,
),
};
} else {
toolCallParts.push({
@ -2293,7 +2511,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
toolName: toolEvent.tool_name as string,
argsText: JSON.stringify(toolArgs),
args: toolArgs,
});
textCursor: cumulativeText.length,
...(toolProvenance ? { provenance: toolProvenance } : {}),
} as PositionedToolCallPart);
}
} else if (toolEvent.type === "tool_end") {
const id =
@ -2437,21 +2657,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
native_part: { parts: mergedParts },
};
}
const existing = toolCallParts[
idx
] as PositionedToolCallPart;
toolCallParts[idx] = {
...toolCallParts[idx],
...existing,
args: mergedArgs,
argsText: JSON.stringify(mergedArgs ?? {}),
result: parsedResult,
provenance: mergeToolProvenance(
existing.provenance,
toolProvenance,
),
};
}
}
// Cumulative yield; orderAssistantContent puts search/code
// before text and generated images after.
const textParts = pinTextThoughtSignature(
parseAssistantContent(cumulativeText),
);
yield {
content: orderAssistantContent(textParts),
content: buildAssistantContent(cumulativeText),
metadata: {
timing: buildTiming(
streamStartTime,
@ -2504,10 +2726,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
| { extra_content?: unknown }
| undefined
)?.extra_content;
if (
deltaExtraContent &&
typeof deltaExtraContent === "object"
) {
if (deltaExtraContent && typeof deltaExtraContent === "object") {
const eGoogle = (deltaExtraContent as Record<string, unknown>)
.google;
if (eGoogle && typeof eGoogle === "object") {
@ -2556,6 +2775,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
Array.isArray(rawDeltaToolCalls) &&
rawDeltaToolCalls.length > 0
) {
closeReasoningContent();
for (const tc of rawDeltaToolCalls) {
if (!tc || typeof tc !== "object") continue;
const call = tc as {
@ -2575,20 +2795,16 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
: undefined;
if (!existing && idx !== undefined) {
existing = toolCallParts.find(
(p) =>
(
p as ToolCallMessagePart & { _delta_index?: number }
)._delta_index === idx,
(p) => (p as PositionedToolCallPart)._delta_index === idx,
);
}
const argsFragment = call.function?.arguments ?? "";
if (existing) {
const prevName = existing.toolName ?? "";
const nextName = call.function?.name ?? prevName;
const merged =
(existing.argsText ?? "") + argsFragment;
let parsedArgs:
ToolCallMessagePart["args"] = existing.args ?? {};
const merged = (existing.argsText ?? "") + argsFragment;
let parsedArgs: ToolCallMessagePart["args"] =
existing.args ?? {};
if (merged) {
try {
parsedArgs = JSON.parse(
@ -2600,24 +2816,18 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
} as ToolCallMessagePart["args"];
}
}
const prevExtra = (
existing as ToolCallMessagePart & {
extra_content?: unknown;
}
).extra_content;
const updated: ToolCallMessagePart & {
_delta_index?: number;
extra_content?: unknown;
} = {
...(existing as ToolCallMessagePart),
const prevExtra = (existing as PositionedToolCallPart)
.extra_content;
const updated: PositionedToolCallPart = {
...(existing as PositionedToolCallPart),
toolName: nextName,
argsText: merged,
args: parsedArgs,
...(call.extra_content !== undefined
? { extra_content: call.extra_content }
: prevExtra !== undefined
? { extra_content: prevExtra }
: {}),
? { extra_content: prevExtra }
: {}),
...(idx !== undefined ? { _delta_index: idx } : {}),
};
const replaceIdx = toolCallParts.indexOf(existing);
@ -2626,8 +2836,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
} else {
const callId =
stableId ||
`tool_call_${idx ?? toolCallParts.length}`;
stableId || `tool_call_${idx ?? toolCallParts.length}`;
const argsText = argsFragment;
let parsedArgs: ToolCallMessagePart["args"] = {};
if (argsText) {
@ -2641,15 +2850,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
} as ToolCallMessagePart["args"];
}
}
const fresh: ToolCallMessagePart & {
_delta_index?: number;
extra_content?: unknown;
} = {
const fresh: PositionedToolCallPart = {
type: "tool-call" as const,
toolCallId: callId,
toolName: call.function?.name ?? "",
argsText,
args: parsedArgs,
textCursor: cumulativeText.length,
...(call.extra_content !== undefined
? { extra_content: call.extra_content }
: {}),
@ -2659,12 +2866,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
}
yield {
content: [
...toolCallParts,
...pinTextThoughtSignature(
parseAssistantContent(cumulativeText),
),
],
content: buildAssistantContent(cumulativeText),
metadata: {
timing: buildTiming(
streamStartTime,
@ -2695,10 +2897,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
}
if (delta) {
if (reasoningContentOpen) {
cumulativeText += "</think>";
reasoningContentOpen = false;
}
closeReasoningContent();
cumulativeText += delta;
}
// Strip a trailing ${...} template-literal artifact from
@ -2709,12 +2908,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
"",
);
}
const parts = pinTextThoughtSignature(
parseAssistantContent(cumulativeText),
);
const textParts = parseAssistantContent(cumulativeText);
if (
parts.some((part) => part.type === "reasoning") &&
textParts.some((part) => part.type === "reasoning") &&
!reasoningStartAt
) {
reasoningStartAt = Date.now();
@ -2729,9 +2926,9 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
);
}
if (parts.length > 0 || toolCallParts.length > 0) {
if (textParts.length > 0 || toolCallParts.length > 0) {
yield {
content: orderAssistantContent(parts),
content: buildAssistantContent(cumulativeText),
metadata: {
timing: buildTiming(
streamStartTime,
@ -2756,13 +2953,10 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
throw streamError;
}
}
// If the stream ended inside a delta.reasoning_content block
// (Kimi / DeepSeek), close the open <think> tag so the reasoning
// panel parses cleanly.
if (reasoningContentOpen) {
cumulativeText += "</think>";
reasoningContentOpen = false;
}
// If the stream ended while we were still inside a
// delta.reasoning_content block (Kimi / DeepSeek path), close
// the open <think> tag so the reasoning panel parses cleanly.
closeReasoningContent();
settleFirstTokenOk();
// Extract source parts from completed web_search and web_fetch
@ -2826,9 +3020,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
yield {
content: [
...orderAssistantContent(
pinTextThoughtSignature(parseAssistantContent(cumulativeText)),
),
...buildAssistantContent(cumulativeText),
...sourceParts,
...documentCitationParts,
],