diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index c44db7485d..8e202d161a 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 4d4ae57afe..d12eb6e29e 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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 = ("", " 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 += "" - 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 += "" 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 += "" 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 += "" + cumulative_display += reasoning_accum + cumulative_display += "" + 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 += "" 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", diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 0cbdea1337..74a011093a 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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, diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index acf6bd9eda..7d498a170d 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -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"") @@ -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 ``{: ...}`` 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 "" 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 - # "" 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": ""} diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 601cd9f92b..ccb777db1c 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -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: ``v`` - Closing tags (````, ````, ````) - 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: v... -- closing tags optional; # 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 tag rather than + # the end of the response, so a complete call followed by + # trailing prose is still accepted (matching the JSON-style + # path, which already tolerates trailing text). + # rfind picks the last , so a literal + # 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 # 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)}", diff --git a/studio/backend/core/inference/tool_loop_controller.py b/studio/backend/core/inference/tool_loop_controller.py new file mode 100644 index 0000000000..667ec7a33e --- /dev/null +++ b/studio/backend/core/inference/tool_loop_controller.py @@ -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, + ) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index cf264d0dc9..28634396f5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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"|\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"] = ( diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 77b0b21212..207c4b3805 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -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 # ===================================================================== diff --git a/studio/backend/tests/test_gguf_route_cursor_reset.py b/studio/backend/tests/test_gguf_route_cursor_reset.py new file mode 100644 index 0000000000..fdbf31f9de --- /dev/null +++ b/studio/backend/tests/test_gguf_route_cursor_reset.py @@ -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"] diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py new file mode 100644 index 0000000000..fa583ef53d --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -0,0 +1,1151 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Focused tests for the GGUF llama.cpp agentic tool loop. + +These tests drive ``LlamaCppBackend.generate_chat_completion_with_tools`` +with fake llama-server SSE streams. They require no model, subprocess, GPU, +or network access. +""" + +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 = 48847 + 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 _tool_names(payload: dict) -> list[str]: + return [ + (tool.get("function") or {}).get("name") + for tool in payload.get("tools", []) + if (tool.get("function") or {}).get("name") + ] + + +def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): + """llama-server may emit content first and then native delta.tool_calls. + + Studio must not drop that tool call after it has streamed the preface. + """ + + tool_call_id = "call_render_late" + first_stream = [ + _sse({"content": "Here is the artifact.\n\n"}), + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": tool_call_id, + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "
red
", + "title": "Simple Red Square", + } + ), + }, + } + ] + } + ), + _done(), + ] + second_stream = [ + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: Simple Red Square." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + content_events = [e for e in events if e.get("type") == "content"] + assert content_events[0]["text"] == "Here is the artifact.\n\n" + + first_content_index = next( + i for i, event in enumerate(events) if event.get("type") == "content" + ) + actual_tool_start_index = next( + i + for i, event in enumerate(events) + if event.get("type") == "tool_start" and event.get("arguments", {}).get("code") + ) + assert first_content_index < actual_tool_start_index + + assert calls == [ + ( + "render_html", + { + "code": "
red
", + "title": "Simple Red Square", + }, + ) + ] + assert any(e.get("type") == "tool_end" and e.get("tool_name") == "render_html" for e in events) + + # The second llama-server request should include the assistant preface + # plus the structured tool call, preserving OpenAI-compatible ordering. + assert len(payloads) == 2 + assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"] + assert assistant_messages[-1]["content"] == "Here is the artifact.\n\n" + assert assistant_messages[-1]["tool_calls"][0]["id"] == tool_call_id + assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" + + +def test_repeat_render_html_nudge_is_not_user_visible_error(monkeypatch): + """A repeated render_html call is an internal no-op, not a visible card.""" + + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "first", + "title": "First", + } + ), + }, + } + ] + } + ), + _done(), + ] + repeat_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_repeat", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "repeat", + "title": "Repeat", + } + ), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Short note."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, repeat_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: First." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + }, + {"type": "function", "function": {"name": "web_search"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 2, + ) + ) + + assert calls == [ + ( + "render_html", + {"code": "first", "title": "First"}, + ) + ] + assert _tool_names(payloads[1]) == ["web_search"] + + actual_tool_starts = [ + event + for event in events + if event.get("type") == "tool_start" and event.get("arguments", {}).get("code") + ] + tool_ends = [ + event + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") == "render_html" + ] + assert len(actual_tool_starts) == 1 + assert len(tool_ends) == 1 + + assert len(payloads) == 3 + render_tool_messages = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "tool" and message.get("name") == "render_html" + ] + assert len(render_tool_messages) == 1 + internal_nudges = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "user" + and "Do not call render_html again" in message.get("content", "") + ] + assert len(internal_nudges) == 1 + + +def test_render_html_success_drops_tool_schema_before_final_pass(monkeypatch): + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "ok"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Done."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + return "Rendered HTML artifact: Done." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Render this."}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + max_tool_iterations = 3, + ) + ) + + assert len(payloads) == 2 + assert "tools" not in payloads[1] + assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) + final_user_messages = [ + m.get("content", "") for m in payloads[1]["messages"] if m.get("role") == "user" + ] + assert not any("used all available tool calls" in message for message in final_user_messages) + + +def test_non_consecutive_duplicate_web_search_is_internal_noop(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + python_call = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print('ok')"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from gathered data."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, python_call, duplicate_search, final_stream], + payloads, + ) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return f"ok:{name}" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + tools = tools, + max_tool_iterations = 3, + ) + ) + + assert calls == [ + ("web_search", {"query": "gpu prices 2026"}), + ("python", {"code": "print('ok')"}), + ] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_start" and event.get("tool_name") + ] == ["web_search", "python"] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") + ] == ["web_search", "python"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 4 + assert _tool_names(payloads[3]) == ["web_search", "python"] + duplicate_nudges = [ + message + for message in payloads[3]["messages"] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + + +def test_duplicate_web_search_noop_allows_distinct_followup_tool(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + python_call = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print('ok')"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from gathered data."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, duplicate_search, python_call, final_stream], + payloads, + ) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return f"ok:{name}" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus in 2026 prices and use python"}], + tools = tools, + max_tool_iterations = 4, + ) + ) + + assert calls == [ + ("web_search", {"query": "gpu prices 2026"}), + ("python", {"code": "print('ok')"}), + ] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_start" and event.get("tool_name") + ] == ["web_search", "python"] + assert [ + event.get("tool_name") + for event in events + if event.get("type") == "tool_end" and event.get("tool_name") + ] == ["web_search", "python"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 4 + assert _tool_names(payloads[2]) == ["web_search", "python"] + duplicate_nudges = [ + message + for message in payloads[2]["messages"] + if message.get("role") == "user" + and "already completed successfully" in message.get("content", "") + ] + assert len(duplicate_nudges) == 1 + + +def test_repeated_duplicate_noop_transitions_to_final_pass(monkeypatch): + first_search = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_one = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + duplicate_two = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_3", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer from first search."}), _done()] + payloads: list[dict] = [] + backend = _make_backend( + monkeypatch, + [first_search, duplicate_one, duplicate_two, final_stream], + payloads, + ) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 10, + ) + ) + + assert calls == [("web_search", {"query": "gpu prices 2026"})] + assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ + "call_search_1" + ] + assert len(payloads) == 4 + assert "tools" not in payloads[-1] + assert any( + event.get("type") == "content" and event.get("text") == "Final answer from first search." + for event in events + ) + + +def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch): + same_turn_duplicates = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_search_1", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + }, + { + "index": 1, + "id": "call_search_2", + "type": "function", + "function": { + "name": "web_search", + "arguments": json.dumps({"query": "gpu prices 2026"}), + }, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [same_turn_duplicates, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "search-result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search gpus"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [("web_search", {"query": "gpu prices 2026"})] + assert [event.get("tool_call_id") for event in events if event.get("type") == "tool_end"] == [ + "call_search_1" + ] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_search_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + + +def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch): + same_turn_render_calls = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_html_1", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "one"}), + }, + }, + { + "index": 1, + "id": "call_html_2", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps({"code": "two"}), + }, + }, + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "Final answer."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [same_turn_render_calls, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: One." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "render html"}], + tools = [{"type": "function", "function": {"name": "render_html"}}], + max_tool_iterations = 2, + ) + ) + + assert calls == [("render_html", {"code": "one"})] + assert [ + event.get("tool_call_id") + for event in events + if event.get("type") == "tool_start" and not event.get("arguments") + ] == ["call_html_1"] + assert not [ + event + for event in events + if event.get("tool_call_id") == "call_html_2" + and event.get("type") in {"tool_start", "tool_end"} + ] + assert len(payloads) == 2 + assert "tools" not in payloads[1] + render_nudges = [ + message + for message in payloads[1]["messages"] + if message.get("role") == "user" + and "Do not call render_html again" in message.get("content", "") + ] + assert len(render_nudges) == 1 + + +def test_disabled_tool_call_is_internal_noop(monkeypatch): + disabled_python = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_python_disabled", + "type": "function", + "function": { + "name": "python", + "arguments": json.dumps({"code": "print(1)"}), + }, + } + ] + } + ), + _done(), + ] + final_stream = [_sse({"content": "I cannot run Python here."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [disabled_python, final_stream], payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "run python"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert not [event for event in events if event.get("type") in {"tool_start", "tool_end"}] + assert len(payloads) == 2 + disabled_nudges = [ + message + for message in payloads[1]["messages"] + if message.get("role") == "user" and "not enabled" in message.get("content", "") + ] + assert len(disabled_nudges) == 1 + + +def test_render_html_success_does_not_reprompt_render_html_intent(monkeypatch): + """After render_html succeeds, do not force another render_html call. + + The post-tool model pass can say it will use render_html again without + emitting a tool call. That should be accepted as a final model mistake, + not turned into repeated internal re-prompts after the artifact already + exists. + """ + + first_stream = [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_first", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "first", + "title": "First", + } + ), + }, + } + ] + } + ), + _done(), + ] + post_tool_stream = [ + _sse({"content": "I will now use render_html again."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, post_tool_stream], payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: First." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + assert len(payloads) == 2 + assert len(calls) == 1 + assert any( + event.get("type") == "content" and event.get("text") == "I will now use render_html again." + for event in events + ) + + +def test_internal_reprompt_attempts_do_not_duplicate_visible_text(monkeypatch): + """No-tool re-prompt attempts should not concatenate into the UI.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [_sse({"content": "Understood. I will use render_html now."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 2 + + +def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch): + """A hidden forced re-prompt may fall back to a plain final answer.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [ + _sse({"content": "No tool is needed. Final answer: use a red square."}), + _done(), + ], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ], + max_tool_iterations = 1, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == [ + "I will use render_html now.", + "No tool is needed. Final answer: use a red square.", + ] + assert len(payloads) == 2 + + +def test_internal_reprompt_disabled_when_auto_heal_disabled(monkeypatch): + streams = [[_sse({"content": "I will use render_html now."}), _done()]] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + def fake_execute_tool(name, arguments, **_kwargs): + raise AssertionError(f"unexpected tool execution: {name} {arguments}") + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + auto_heal_tool_calls = False, + ) + ) + + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now."] + assert len(payloads) == 1 + + +def test_auto_heal_disabled_parses_well_formed_xml_when_tools_enabled(monkeypatch): + streams = [ + [ + _sse( + { + "content": '{"name":"web_search","arguments":{"query":"x"}}' + } + ), + _done(), + ], + [_sse({"content": "done"}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "result" + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + auto_heal_tool_calls = False, + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "x"})] + assert not any( + event.get("type") == "content" and "" in event.get("text", "") + for event in events + ) + + +def test_reprompted_tool_call_still_streams_final_answer(monkeypatch): + """Suppression ends once a forced re-prompt actually calls a tool.""" + + streams = [ + [_sse({"content": "I will use render_html now."}), _done()], + [ + _sse( + { + "tool_calls": [ + { + "index": 0, + "id": "call_forced", + "type": "function", + "function": { + "name": "render_html", + "arguments": json.dumps( + { + "code": "forced", + "title": "Forced", + } + ), + }, + } + ] + } + ), + _done(), + ], + [_sse({"content": "Final note after tool."}), _done()], + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, streams, payloads) + + calls: list[tuple[str, dict]] = [] + + def fake_execute_tool(name, arguments, **_kwargs): + calls.append((name, arguments)) + return "Rendered HTML artifact: Forced." + + monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool) + + tools = [ + { + "type": "function", + "function": { + "name": "render_html", + "description": "Render HTML.", + "parameters": { + "type": "object", + "properties": {"code": {"type": "string"}}, + "required": ["code"], + }, + }, + } + ] + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "Make a red square."}], + tools = tools, + max_tool_iterations = 1, + ) + ) + + assert len(calls) == 1 + content_texts = [event.get("text", "") for event in events if event.get("type") == "content"] + assert content_texts == ["I will use render_html now.", "Final note after tool."] + assert len(payloads) == 3 diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 55d2a48df1..16ed1dc182 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -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 = '{"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 = "print('hi')" 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 = "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 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 {"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 {"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 '{"name":"render_html","arguments":{"code":"one"}}' + 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": "one"})] + 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( + [ ['{"name":"web_search","arguments":{"query":"x"}}'], ['{"name":"web_search","arguments":{"query":"x"}}'], ["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( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["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( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["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( - ['{"name":"terminal","arguments":{"command":"echo bypass"}}'] - ), - 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( - ['{"name":"python","arguments":{"code":"print(1)"}}'] - ), - 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 '{"name":"terminal","arguments":{"command":"echo bypass"}}' + 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( + [ + ['{"name":"python","arguments":{"code":"print(1)"}}'], + ["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( + [ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ['{"name":"web_search","arguments":{"query":"literal"}}'], + ] + ) + + 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 "" 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 = [ + ['{"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 "" 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 = [["{not valid json}"]], + 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 "" 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 = [ + [ + '{"name":"web_search","arguments":{"query":"A"}}' + '{"name":"web_search","arguments":{"query":"A"}}' + ], + ["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)"} diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py new file mode 100644 index 0000000000..8ff41342d7 --- /dev/null +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -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 +``...`` call even when the model appends prose +after the closing tag -- matching the JSON-style ``...`` 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 = ( + "weather london" + " 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 = "cats \n\n" + call = _only(text) + assert call == {"name": "web_search", "arguments": {"query": "cats"}} + + def test_closed_function_without_trailing_text_still_parses(self): + text = "cats" + call = _only(text) + assert call == {"name": "web_search", "arguments": {"query": "cats"}} + + def test_multi_param_with_trailing_prose(self): + text = ( + "ls -la" + "home 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 is the last one; the literal inside + # the code argument must survive (rfind, not the first match). + text = ( + "" + 'print("")' + " all done" + ) + call = _only(text) + assert call == {"name": "python", "arguments": {"code": 'print("")'}} + + def test_incomplete_function_without_close_is_still_rejected(self): + text = "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 present, but the single parameter never closes. + text = "weather london" + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + + +class TestParityWithJsonStyle: + def test_json_tool_call_with_trailing_prose_is_accepted(self): + text = ( + '{"name":"web_search","arguments":{"query":"weather london"}}' + " 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"{q} trailing", + allow_incomplete = False, + ) + js = parse_tool_calls_from_text( + f'{{"name":"web_search","arguments":{{"query":"{q}"}}}} 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 = "cats" + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "web_search" diff --git a/studio/backend/tests/test_tool_loop_controller.py b/studio/backend/tests/test_tool_loop_controller.py new file mode 100644 index 0000000000..dea5de6d6e --- /dev/null +++ b/studio/backend/tests/test_tool_loop_controller.py @@ -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": ""}, "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": ""}, "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" diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 6cc01f2595..2ba3310fbe 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -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 {"name":"web_search"} survives' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = False) == text + assert "" 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" diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 31f73c37bd..34cae518e3 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -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) - : null; - const argsGoogle = - argsObj && typeof argsObj.google === "object" && argsObj.google !== null - ? (argsObj.google as Record) - : null; - const hasNativePart = Boolean( - argsGoogle && - typeof argsGoogle.native_part === "object" && - argsGoogle.native_part !== null, - ); - const hasServerToolMarker = Boolean( - argsObj && (argsObj as Record)._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) - : null; - const argsGoogle = - argsObj && typeof argsObj.google === "object" && argsObj.google !== null - ? (argsObj.google as Record) - : 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[number]; +type SerializedToolResult = { + role: "tool"; + content: string; + tool_call_id: string; + name?: string; +}; + +type ToolPartReplayMetadata = { + argsObj: Record | null; + argsGoogle: Record | 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) + : null; + const argsGoogle = + argsObj && typeof argsObj.google === "object" && argsObj.google !== null + ? (argsObj.google as Record) + : null; + const hasNativePart = Boolean( + argsGoogle && + typeof argsGoogle.native_part === "object" && + argsGoogle.native_part !== null, + ); + const hasServerToolMarker = Boolean( + argsObj && (argsObj as Record)._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) + : {}; + const google = + extra.google && + typeof extra.google === "object" && + !Array.isArray(extra.google) + ? (extra.google as Record) + : {}; + 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 { // ... 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, - ) => { - 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[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) } 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 += ""; + 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) .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 += ""; - 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 tag so the reasoning - // panel parses cleanly. - if (reasoningContentOpen) { - cumulativeText += ""; - reasoningContentOpen = false; - } + // If the stream ended while we were still inside a + // delta.reasoning_content block (Kimi / DeepSeek path), close + // the open 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, ],