diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 455d1d084c..aab3470193 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -44,7 +44,8 @@ from core.inference.llama_server_args import ( from core.inference.tool_call_parser import ( _GEMMA_BARE_TC_PREFIX_RE, _GEMMA_BARE_TC_RE, - _TOOL_ALL_PATS, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, _balanced_brace_end, _strip_function_xml_calls, _strip_gemma_wrapperless_calls, @@ -58,6 +59,16 @@ from core.inference.tool_call_parser import ( strip_llama3_leading_sentinels, strip_tool_markup as _shared_strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the GGUF streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + apply_tool_strip_patterns, + strip_outside_think, + strip_tool_call_markup, +) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( @@ -256,6 +267,72 @@ _FINAL_ANSWER_SIGNAL = re.compile( ) +def _gguf_active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in (active_tools or []) + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Rehearsal NAME chars (word + hyphen, matching the parser); the lookbehind excludes the +# Mistral [CALL_ID]...[ARGS] shape. +_GGUF_REHEARSAL_ARGS_RE = re.compile(r"(? int: + """Index of the first ``NAME[ARGS]`` whose NAME is an active tool, else -1. A + bare/inactive-name ``foo[ARGS]`` in prose is not a call; mirrors the safetensors + ``_earliest_tool_signal`` name-gating (no unrestricted GGUF mode).""" + active = set(_gguf_active_tool_names(active_tools)) + if not active: + return -1 + for m in _GGUF_REHEARSAL_ARGS_RE.finditer(text): + if m.group(1) in active: + return m.start() + return -1 + + +def _gguf_has_genuine_tool_signal(text: str, signals, active_tools: list[dict]) -> bool: + """True when ``text`` holds a genuine tool-call boundary for one of ``signals``. + + Unambiguous markers (````, ``[TOOL_CALLS]``, ``= 0: + return True + continue + if sig in text: + return True + return False + + +def _is_rehearsal_prefix(stripped: str, active_tools: list[dict]) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of ``NAME[ARGS]`` for an + active tool -- the bare tool name arriving in its own chunk before ``[ARGS]{...}``. + Mirrors the safetensors loop so the split rehearsal call is not streamed.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + for name in _gguf_active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len(text: str, active_tools: list[dict]) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose. Mirrors safetensors.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return len(tail) if tail and _is_rehearsal_prefix(tail, active_tools) else 0 + + 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 @@ -8418,6 +8495,13 @@ class LlamaCppBackend: _reasoning_started_at: Optional[float] = None _reasoning_summary_emitted = False + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = no gate. + _enabled_names_gate = set(_gguf_active_tool_names(tools)) if tools else None + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = list(tools or []) + def _reasoning_summary_event(started_at: float) -> dict: return { "type": "reasoning_summary", @@ -8436,25 +8520,42 @@ class LlamaCppBackend: ) -> str: if not (auto_heal_tool_calls or force): return text + # Delegate to the shared parser-side strip so the GGUF cleanup covers every family the + # parser promotes (Llama <|python_tag|>, Mistral [TOOL_CALLS], bare rehearsal, function + # XML, Gemma) and stays aligned with detection; tool_healing's strip omits the loop-only + # forms (python_tag / Mistral name) and would leak them into display. return _shared_strip_tool_markup( - text, final = final, enabled_tool_names = _enabled_tool_names + text, final = final, enabled_tool_names = _enabled_names_gate ) def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - # Shared parser patterns (not the legacy tool_healing set) so textual - # Mistral/python_tag calls entering DRAINING never leak. Balanced strips - # first (nested JSON removed whole); no final trim so length compares hold. - text = _strip_mistral_closed_calls(text) - text = _strip_gemma_wrapperless_calls(text, _enabled_tool_names) - # Parser-accurate scans close at each call's REAL terminator before - # the regex arms: literal markup inside a value is data. - text = _strip_function_xml_calls(text, final = True) - text = _strip_glm_calls(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first (nested JSON removed whole; literal markup inside a value is that + # call's data), then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal + markerless + # strips are name-gated on the ORIGINAL list (strip/detect aligned). + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = _enabled_names_gate) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, _enabled_names_gate) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = _enabled_names_gate + ) + return seg + + # Preserve think blocks verbatim (a rehearsed call inside one must not be deleted). + return strip_outside_think(text, _seg) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its @@ -8814,12 +8915,18 @@ class LlamaCppBackend: in_thinking = False cumulative_display += token cleaned = _strip_tool_markup_streaming(cumulative_display) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Hold a trailing bare active-tool-name (split rehearsal) + # until [ARGS] arrives; released by later prose or stream end. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] if _hold else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } elif detect_state == _S_BUFFERING: @@ -8828,7 +8935,8 @@ class LlamaCppBackend: if not stripped_buf: continue - # Check tool signal prefixes. + # Bracket tags arrive mid-buffer, so substring-check too; + # ``[ARGS]`` counts only as a regex-matched NAME[ARGS]. is_prefix = False is_match = False for sig in _tool_xml_signals: @@ -8838,6 +8946,31 @@ class LlamaCppBackend: if sig.startswith(stripped_buf): is_prefix = True break + if sig == "[ARGS]": + # Active NAME[ARGS] only; inactive-name prose + # is gated out, not drained/parsed. + if ( + _gguf_rehearsal_signal_pos( + stripped_buf, _detect_tools + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped_buf: + is_match = True + break + + # Split rehearsal: hold the bare name until + # its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and _is_rehearsal_prefix(stripped_buf, _detect_tools) + ): + is_prefix = True + is_rehearsal_prefix = True # Signal-less call shapes (mirror the safetensors # loop): Llama-3.2 bare {"name":..} and Gemma @@ -8884,9 +9017,14 @@ class LlamaCppBackend: # Tool signal -- flush any visible # prefix before DRAINING so the # route sends it before tool_start. + # Use the final strip (all families incl. Llama + # <|python_tag|> / Mistral name): the buffer holds + # the whole call, so a streaming closed-only strip + # would leak its open-ended markup as display text. _flush_reasoning_and_buffer() - cleaned = _strip_tool_markup_streaming( + cleaned = _strip_tool_markup( cumulative_display, + final = True, force = True, ) if len(cleaned) > len(_last_emitted): @@ -8898,8 +9036,14 @@ class LlamaCppBackend: } detect_state = _S_DRAINING elif _hold_buffer or ( - is_prefix and len(stripped_buf) < _MAX_BUFFER_CHARS + is_prefix + and ( + is_rehearsal_prefix + or len(stripped_buf) < _MAX_BUFFER_CHARS + ) ): + # A rehearsal prefix is self-bounded; the buffer + # cap must not cut long MCP names short. pass # keep buffering else: # Not a tool -- flush buffer @@ -8910,12 +9054,20 @@ class LlamaCppBackend: cleaned = _strip_tool_markup( cumulative_display, ) - if len(cleaned) > len(_last_emitted): - _last_emitted = cleaned + # Same trailing-name hold as STREAMING for this + # first flush out of BUFFERING. + _hold = _held_rehearsal_tail_len(cleaned, _detect_tools) + _emit = ( + cleaned[: len(cleaned) - _hold] + if _hold + else cleaned + ) + if len(_emit) > len(_last_emitted): + _last_emitted = _emit if not _suppress_visible_output: yield { "type": "content", - "text": cleaned, + "text": _emit, } except json.JSONDecodeError: @@ -8933,7 +9085,9 @@ class LlamaCppBackend: _is_bare_tc = bool(active_tools) and _looks_like_enabled_bare_json( _bare_eos, _enabled_tool_names ) - if stripped_buf and any(s in stripped_buf for s in _tool_xml_signals): + if stripped_buf and _gguf_has_genuine_tool_signal( + stripped_buf, _tool_xml_signals, _detect_tools + ): detect_state = _S_DRAINING elif _is_bare_tc: detect_state = _S_DRAINING @@ -9066,6 +9220,12 @@ class LlamaCppBackend: "type": "content", "text": forced_visible_text, } + elif not _suppress_visible_output: + # Turn ended as a plain answer (no [ARGS] followed): the held + # rehearsal tail is real prose, release it. + _final_clean = _strip_tool_markup_streaming(cumulative_display) + if len(_final_clean) > len(_last_emitted): + yield {"type": "content", "text": _final_clean} # Content was already streamed. Yield metadata. yield {"type": "status", "text": ""} diff --git a/studio/backend/core/inference/passthrough_healing.py b/studio/backend/core/inference/passthrough_healing.py index fe1aca0e4a..a444431f8d 100644 --- a/studio/backend/core/inference/passthrough_healing.py +++ b/studio/backend/core/inference/passthrough_healing.py @@ -32,16 +32,15 @@ from typing import Any, Optional from core.inference.tool_loop_controller import coerce_tool_arguments from core.tool_healing import parse_tool_calls_from_text -# Signals limited to the formats parse_tool_calls_from_text (core.tool_healing) -# actually promotes. The parser module's broader signal list also covers Llama -# <|python_tag|> and Mistral [TOOL_CALLS] for the streaming DRAIN buffers whose -# full parser handles them; buffering those here would hold a streamed -# client-tool call until finalization and then flush it as prose (this healer -# cannot promote them), so the passthrough keeps its own aligned list. +# Only the formats this healer's parser can promote -- narrower than the loops' +# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare +# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a +# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in. _HEAL_SIGNALS = ( "", "<|tool_call>", " list: diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 8e86d09754..aa732f47e4 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via ``core.inference.tools``. """ +import bisect import re import threading from typing import Callable, Generator, Optional @@ -23,7 +24,8 @@ from loggers import get_logger from core.inference.tool_call_parser import ( _GEMMA_BARE_TC_PREFIX_RE, _GEMMA_BARE_TC_RE, - _TOOL_ALL_PATS, + _TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS, + _TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS, _balanced_brace_end, _strip_function_xml_calls, _strip_gemma_wrapperless_calls, @@ -39,6 +41,16 @@ from core.inference.tool_call_parser import ( strip_llama3_leading_sentinels, strip_tool_markup, ) + +# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated +# pattern lists, so the safetensors streaming strip stays aligned with the parser. +from core.tool_healing import ( + _REHEARSAL_TAIL_STRIP_RE, + _strip_bracket_tag_calls, + _think_spans_outside_tool_markup, + apply_tool_strip_patterns, + strip_outside_think, +) from core.inference.tool_loop_controller import ( ToolLoopController, coerce_tool_arguments, @@ -93,6 +105,147 @@ def _active_tool_names(active_tools: list[dict]) -> list[str]: return [name for name in names if name] +def _active_tool_names(active_tools: list[dict]) -> list[str]: + names = [ + (tool.get("function") or {}).get("name") + for tool in active_tools + if isinstance(tool, dict) and isinstance(tool.get("function"), dict) + ] + return [name for name in names if name] + + +# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal; +# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held. +_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?") + + +def _is_rehearsal_prefix( + stripped: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]`` + rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space + means prose. Unrestricted mode accepts any identifier; else NAME must be active.""" + if not stripped or any(ch.isspace() for ch in stripped): + return False + if unrestricted: + return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None + for name in _active_tool_names(active_tools): + if stripped == name or f"{name}[ARGS]".startswith(stripped): + return True + return False + + +def _held_rehearsal_tail_len( + text: str, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Length of a trailing bare tool-name token that may be a split rehearsal call + (``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it + instead of leaking the name. Returns 0 for ordinary prose.""" + i = len(text) + while i > 0 and not text[i - 1].isspace(): + i -= 1 + tail = text[i:] + return ( + len(tail) + if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted) + else 0 + ) + + +def _rehearsal_name_start( + candidate: str, + signal_pos: int, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding + bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the + signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode).""" + if not candidate.startswith("[ARGS]", signal_pos): + return signal_pos + j = signal_pos + while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"): + j -= 1 + if j < signal_pos and ( + unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools) + ): + return j + return signal_pos + + +def _earliest_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> int: + """Index where the turn's first genuine tool-call boundary begins, or -1. + + Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal + only when an active tool name (any name in unrestricted mode) precedes it, so a + literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a + real ``NAME[ARGS]`` the boundary is pulled back to NAME.""" + best = -1 + for sig in signals: + if sig != "[ARGS]": + p = candidate.find(sig) + if p >= 0 and (best < 0 or p < best): + best = p + continue + from_idx = 0 + while True: + p = candidate.find("[ARGS]", from_idx) + if p < 0: + break + name_start = _rehearsal_name_start( + candidate, p, active_tools, unrestricted = unrestricted + ) + if name_start < p: + # Genuine ``NAME[ARGS]``: the boundary is the start of NAME. + if best < 0 or name_start < best: + best = name_start + break + # Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found. + from_idx = p + len("[ARGS]") + return best + + +def _has_genuine_tool_signal( + candidate: str, + signals, + active_tools: list[dict], + *, + unrestricted: bool = False, +) -> bool: + """True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``. + + Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only + when an active tool name (any in unrestricted mode) precedes it. Mirrors the + ``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not + drain inactive-name prose.""" + for sig in signals: + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted + ) + >= 0 + ): + return True + continue + if sig in candidate: + return True + return False + + def strip_tool_markup_streaming( text: str, *, @@ -101,25 +254,46 @@ def strip_tool_markup_streaming( enabled_tool_names: Optional[set] = None, ) -> str: """Strip open-ended tool XML from display text without trimming whitespace. - ``enabled_tool_names`` gates the markerless Gemma ``call:NAME{...}`` strip so a - disabled/example name in prose is kept (mirrors the parser gate).""" + + Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so + streaming and final display agree: balanced strips first (nested JSON removed whole), + then the guarded function-XML / GLM scans that close at each call's REAL terminator so + literal markup inside argument values is data and trailing prose survives. Reasoning + ```` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must + not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names`` + keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose, + not a call), matching the parse / detection active-tool gate.""" if not (auto_heal_tool_calls or tool_protocol_active): return text - # Mirror the final strip's scan order so streaming and final display agree: - # balanced strips first (nested JSON removed whole), then the guarded - # function-XML/GLM scans that close at each call's REAL terminator, so literal - # markup inside argument values is data and trailing prose survives. No final - # trim so streaming length comparisons hold. Leading Magistral [THINK]...[/THINK] - # is dropped (bracket form, not the reasoning channel's ); an unclosed - # [THINK] holds until [/THINK] so the cleaned text stays monotonic. + + # Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the + # ```` channel) so raw reasoning does not leak into streamed display; an unclosed + # leading block is held (dropped to EOF) until its closer streams in. text = _strip_mistral_reasoning(text) - text = _strip_mistral_closed_calls(text) - text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) - text = _strip_function_xml_calls(text, final = True) - text = _strip_glm_calls(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + + def _seg(segment: str, is_last: bool) -> str: + # Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced + # strips first, then the guarded function-XML / GLM scans, then the regex arms + # (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last + # segment (a bare ``foo[ARGS]`` before is prose). Rehearsal strips are name-gated. + seg = _strip_mistral_closed_calls(segment) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_function_xml_calls(seg, final = is_last) + seg = _strip_glm_calls(seg, final = is_last) + pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if is_last: + seg = apply_tool_strip_patterns( + seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names + ) + return seg + + # Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then + # regrows the cumulative text, corrupting append-by-length consumers. + return strip_outside_think(text, _seg) def _strip_tool_markup_final( @@ -149,23 +323,66 @@ def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) _FUNCTION_SIGNAL_RE = re.compile(r"") _TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"') +# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional +# render-html card fires for bracket-tag serializations too. +_MISTRAL_RENDER_NAME_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) +_REHEARSAL_RENDER_NAME_RE = re.compile(r"(? bool: - """Return True when the first drained tool call is clearly render_html.""" - function_match = _FUNCTION_SIGNAL_RE.search(content) - tool_call_index = content.find("") - if not function_match and tool_call_index < 0: + """Return True when the FIRST tool call in ``content`` is clearly render_html. + + Covers every serialization the loop executes (XML ```` / ````, + Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a + render_html marker inside another call's argument is treated as data. Markers inside + a ```` / ``[THINK]`` block are dropped since the parser skips them.""" + think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in think_spans] + + def _in_think(pos: int) -> bool: + if not think_spans: + return False + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1] + + def _first_outside(start: int, finder) -> int: + # First occurrence at/after ``start`` that is not inside a think span. + pos = finder(start) + while pos >= 0 and _in_think(pos): + pos = finder(pos + 1) + return pos + + candidates: list[tuple[int, str]] = [] + for fm in _FUNCTION_SIGNAL_RE.finditer(content): + if not _in_think(fm.start()): + candidates.append((fm.start(), fm.group(1))) + break + tc = _first_outside(0, lambda i: content.find("", i)) + if tc >= 0: + nm = _TOOL_CALL_NAME_RE.search(content[tc:]) + candidates.append((tc, nm.group(1) if nm else "")) + mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i)) + if mt >= 0: + mm = _MISTRAL_RENDER_NAME_RE.match(content, mt) + if mm: + candidates.append((mt, mm.group(1))) + else: + # Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the + # first call through the parser (it reads top-level names). + arr_calls = parse_tool_calls_from_text(content[mt:]) + if arr_calls: + candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or "")) + for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content): + if not _in_think(rm.start(1)): + candidates.append((rm.start(1), rm.group(1))) + break + + if not candidates: return False - - if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index): - return function_match.group(1) == "render_html" - - if tool_call_index >= 0: - name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:]) - return bool(name_match and name_match.group(1) == "render_html") - - return False + _pos, name = min(candidates, key = lambda c: c[0]) + return name == "render_html" def _coerce_arguments_with_provenance( @@ -256,6 +473,12 @@ def run_safetensors_tool_loop( conversation.extend(_auto["messages"]) unrestricted_tools = not tools + # Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the + # ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted. + _enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools)) + # Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent + # one-shot), else its repeat is stripped but never drained and the turn ends blank. + _detect_tools = [] if unrestricted_tools else list(tools or []) tool_controller = ToolLoopController( tools = None if unrestricted_tools else tools, auto_heal_tool_calls = auto_heal_tool_calls, @@ -381,18 +604,18 @@ def run_safetensors_tool_loop( if detect_state == _state_streaming: candidate = cumulative_display + delta - signal_pos = -1 - for sig in tool_xml_signals: - p = candidate.find(sig) - if p >= 0 and (signal_pos < 0 or p < signal_pos): - signal_pos = p + # Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is + # pulled back to NAME so the name is not flushed. + signal_pos = _earliest_tool_signal( + candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools + ) if signal_pos >= 0: before_tool = candidate[:signal_pos] cleaned_before = strip_tool_markup_streaming( before_tool, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned_before) > len(last_emitted): last_emitted = cleaned_before @@ -423,11 +646,20 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives; + # released by later prose or the end-of-stream flush. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} continue # BUFFERING: hold until we know it is not a tool call. @@ -445,6 +677,34 @@ def run_safetensors_tool_loop( if sig.startswith(stripped): is_prefix = True break + # Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS] + # counts only with an active NAME so prose is not drained into a no-op. + if sig == "[ARGS]": + if ( + _earliest_tool_signal( + stripped, + ("[ARGS]",), + _detect_tools, + unrestricted = unrestricted_tools, + ) + >= 0 + ): + is_match = True + break + elif sig.startswith("[") and sig in stripped: + is_match = True + break + + # Split rehearsal: hold the bare name until its [ARGS] arrives and matches above. + is_rehearsal_prefix = False + if ( + not is_match + and not is_prefix + and tool_protocol_active + and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools) + ): + is_prefix = True + is_rehearsal_prefix = True # Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML # signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses @@ -512,7 +772,7 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if len(cleaned) > len(last_emitted): last_emitted = cleaned @@ -536,7 +796,8 @@ def run_safetensors_tool_loop( "arguments": {}, "provenance": _tool_event_provenance(provisional = True), } - elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS: + elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS): + # A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short. continue else: detect_state = _state_streaming @@ -545,24 +806,38 @@ def run_safetensors_tool_loop( cumulative_display, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = tool_protocol_active, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) - if len(cleaned) > len(last_emitted): - last_emitted = cleaned - yield {"type": "content", "text": cleaned} + # Same trailing-name hold as STREAMING for this first flush out of BUFFERING. + if tool_protocol_active: + _hold = _held_rehearsal_tail_len( + cleaned, _detect_tools, unrestricted = unrestricted_tools + ) + emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned + else: + emit = cleaned + if len(emit) > len(last_emitted): + last_emitted = emit + yield {"type": "content", "text": emit} # Stream finished -- resolve what we collected. if cancel_event is not None and cancel_event.is_set(): return if detect_state == _state_buffering: - # Buffer never resolved -- tool XML or plain content? + # Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal + # ``foo[ARGS]{...}`` is not parsed. stripped = content_buffer.lstrip() _bare_eos = strip_llama3_leading_sentinels(stripped) if ( stripped and tool_protocol_active - and any(sig in stripped for sig in tool_xml_signals) + and _has_genuine_tool_signal( + stripped, + tool_xml_signals, + _detect_tools, + unrestricted = unrestricted_tools, + ) ): detect_state = _state_draining elif tool_protocol_active and _looks_like_enabled_bare_json( @@ -629,6 +904,17 @@ def run_safetensors_tool_loop( # in full; route-level cleanup still applies the Auto-Heal policy. if content_accum and any(sig in content_accum for sig in tool_xml_signals): yield {"type": "content", "text": content_accum} + else: + # Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real + # prose, release it. + final_clean = strip_tool_markup_streaming( + cumulative_display, + auto_heal_tool_calls = auto_heal_tool_calls, + tool_protocol_active = tool_protocol_active, + enabled_tool_names = _enabled_names_gate, + ) + if len(final_clean) > len(last_emitted): + yield {"type": "content", "text": final_clean} yield {"type": "status", "text": ""} return tool_calls = safety_tc @@ -636,19 +922,23 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) logger.info( "Safetensors safety net: parsed %d tool call(s) from streamed content", len(tool_calls), ) else: - # DRAINING: parse tool calls out of full content. + # DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the + # ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to + # drain here: a spent one-shot (render_html) is off the active list but its re-emitted + # ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of + # being dropped into a blank continuation. tool_calls = parse_tool_calls_from_text( content_accum, id_offset = next_call_id, allow_incomplete = auto_heal_tool_calls, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if not tool_calls: # Parser found nothing. Auto-Heal-enabled display cleanup @@ -682,7 +972,7 @@ def run_safetensors_tool_loop( content_accum, auto_heal_tool_calls = auto_heal_tool_calls, tool_protocol_active = True, - enabled_tool_names = _enabled_tool_names, + enabled_tool_names = _enabled_names_gate, ) if tool_calls: diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 08a6bf418a..70115e5744 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -41,6 +41,9 @@ TOOL_XML_SIGNALS = ( "<|python_tag|>", "[TOOL_CALLS]", "<|tool_call>", + # Bare reasoning-rehearsal marker (``name[ARGS]{...}``, no leading [TOOL_CALLS]); + # keeps a rehearsed call held in the stream so it is promoted, not leaked as prose. + "[ARGS]", # DeepSeek R1 / V3 / V3.1 -- 5 opener variants llama.cpp keeps. "<｜tool▁calls▁begin｜>", "<｜tool▁call▁begin｜>", @@ -90,7 +93,7 @@ _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ # follows; a prose mention (``See [TOOL_CALLS] docs...``) keeps its tail. Bare marker at EOF drops. re.compile(r"<\|tool_call>(?=\s*call\s*:|\s*$).*$", re.DOTALL), re.compile( - r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*[\[{])|\s*$).*$", + r"\[TOOL_CALLS\](?=\s*(?:[\[{]|[A-Za-z_][\w.\-]*(?:[\[{]|\s*$))|\s*$).*$", re.DOTALL, ), re.compile( @@ -572,26 +575,51 @@ def strip_tool_markup( enabled_tool_names: Optional[set] = None, ) -> str: """Strip tool-call markup. ``final=False`` keeps in-progress markup buffered; - ``final=True`` also drops trailing unclosed runs and trims. ``enabled_tool_names`` - gates the markerless Gemma ``call:NAME{...}`` strip so a disabled/example name in - prose is kept (mirrors the parser gate); ``None`` strips every closed call.""" + ``final=True`` also drops trailing unclosed runs and trims. + + ``enabled_tool_names`` gates the name-conditioned forms so a disabled/example name in + prose is kept (mirrors the parser gate): the bare reasoning-rehearsal ``name[ARGS]{...}`` + and the markerless Gemma ``call:NAME{...}`` strip. ``None`` strips every closed call. + """ if final: # Drop a leading Magistral ``[THINK]...[/THINK]`` at end-of-turn; its bracket # form is not the ```` the reasoning channel renders. text = _strip_mistral_reasoning(text) - text = _strip_mistral_closed_calls(text) - if final: - text = _strip_gemma_wrapperless_calls(text, enabled_tool_names) - # Scan-strip the function-XML form (a literal ```` inside a value is - # data). The regex arms below cover the other formats but no-op on function calls here. - text = _strip_function_xml_calls(text, final = final) - # GLM 4.x: scan to the call's real so a literal one inside a value is data, - # not a leak. Qwen {json} is left to the regex arms. - text = _strip_glm_calls(text, final = final) - pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in pats: - text = pat.sub("", text) - return text.strip() if final else text + + def _strip_segment(segment: str, is_last: bool) -> str: + seg_final = final and is_last + seg = _strip_mistral_closed_calls(segment) + # Bare reasoning-rehearsal ``name[ARGS]{json}`` and the Mistral name form promote through + # the shared balanced scan, so strip them the same way (any nesting depth removed whole). + # The rehearsal arm is name-gated: an inactive ``foo[ARGS]{..}`` is prose and is kept. + seg = _tool_healing._strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if seg_final: + # Markerless Gemma ``call:NAME{...}`` (name-gated, mirrors the parse gate); end-of-turn only. + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + # Scan-strip the function-XML form (parser-accurate: a literal ```` in a + # value is data, not a call); the regex arms below cover the other formats. + seg = _strip_function_xml_calls(seg, final = seg_final) + # GLM 4.x: scan to the call's real so a literal one inside a value is data, + # not a leak. Qwen {json} is left to the regex arms. + seg = _strip_glm_calls(seg, final = seg_final) + pats = _TOOL_ALL_PATS if seg_final else _TOOL_CLOSED_PATS + for pat in pats: + seg = pat.sub("", seg) + if seg_final: + # Drop a trailing partial bare rehearsal (``name[ARGS]`` with a truncated or absent + # body) the balanced scan cannot close; gated so prose ``foo[ARGS] ...`` survives. + seg = _tool_healing.apply_tool_strip_patterns( + seg, + [_tool_healing._REHEARSAL_TAIL_STRIP_RE], + enabled_tool_names = enabled_tool_names, + ) + return seg + + # ```` / ``[THINK]`` reasoning is preserved verbatim (a rehearsed call inside it is + # not executed, so it must not be stripped from display either); a literal think marker + # inside a real call's arguments is that call's data and is stripped with the call. + result = _tool_healing.strip_outside_think(text, _strip_segment) + return result.strip() if final else result def has_tool_signal(text: str) -> bool: @@ -711,6 +739,41 @@ def _xml_signal_inside_leading_mistral(content: str) -> bool: return _mistral_region_end(content, trig) is not None +def _parse_bare_rehearsals( + content: str, + *, + id_offset: int = 0, + enabled_tool_names: Optional[set] = None, +) -> list[dict]: + """Promote bare reasoning-rehearsal ``name[ARGS]{json}`` calls that a leading [TOOL_CALLS] + owns-the-turn parse would miss. Only the ``rehearsal`` kind is taken (a Mistral + ``[TOOL_CALLS]name[ARGS]{..}`` yields ``name`` and is not double-counted), and a rehearsal + inside a ```` / ``[THINK]`` block is reasoning, so it is skipped.""" + out: list[dict] = [] + think_spans = _tool_healing._think_spans_outside_tool_markup(content) + for start, end, kind, m in _tool_healing._iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if kind != "rehearsal": + continue + if any(s <= start < e for s, e in think_spans): + continue + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + out.append( + { + "id": f"call_{id_offset + len(out)}", + "type": "function", + "function": {"name": m.group(1), "arguments": json.dumps(payload)}, + } + ) + return out + + _ATTR_FUNC_OPEN_RE = re.compile(r'`` quoted-string handling the GGUF path relies on). + # Qwen/Hermes, Qwen3.5 XML, Gemma 4, plus Mistral [TOOL_CALLS] / bare rehearsal + # ``name[ARGS]{json}`` use the shared tool_healing parser (strict/Auto-Heal contract + + # nested-marker, trailing-prose, and ``<|"|>`` quoted-string handling the GGUF path + # relies on). ``enabled_tool_names`` gates the ambiguous bare-rehearsal form so an + # inactive ``foo[ARGS]{..}`` stays prose. calls = _tool_healing.parse_tool_calls_from_text( content, id_offset = id_offset, allow_incomplete = allow_incomplete, + enabled_tool_names = enabled_tool_names, ) if calls: return calls diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index b91403ed57..1b6b05768a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -1,38 +1,91 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# +# Bracket-tag, rehearsal, and thinking-block-strip logic adapted from forge +# (https://github.com/antoinezambelli/forge), Copyright (c) 2025-2026 +# Antoine Zambelli, used under the MIT License. -"""Lightweight tool-call XML parsing and stripping helpers. +"""Lightweight tool-call parsing and stripping helpers. External inference servers import this module without pulling in the inference -orchestrator, structlog, httpx, or the rest of the studio backend. +orchestrator, structlog, httpx, or the rest of the studio backend. Kept in +lockstep with ``core/inference/tool_call_parser.py`` so those servers +(llama-server wrappers, llama-swap, custom shims) reuse the same logic. Any +change here must also land there. + +Handles these serializations (see ``parse_tool_calls_from_text``): + +* ``{json}`` +* ``<|tool_call>call:name{...}`` (Gemma) +* ``v`` +* ``[TOOL_CALLS]name{json}`` (Mistral / Devstral fallback) +* ``name[ARGS]{json}`` (reasoning-model rehearsal) """ +# PEP 604 annotations must stay import-safe on Python 3.9 (requires-python >=3.9). +from __future__ import annotations + +import bisect import json import re -# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs -# strip first so a closed call goes as a unit before any to-EOF sweep reaches -# nested markup; only the final list adds the .*$ EOF sweeps. +# One nesting level in the strip regexes; deeper may leak markup (still parsed). +_BRACKETED_JSON_ONE_LEVEL = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}" + +# Rehearsal ``name[ARGS]{..}`` strips; group 1 = name for tool-list gating. Closed = +# complete body, tail = truncated; ``(?.*?`` rescans to EOF from every opener +# (quadratic on a stream of unclosed openers). Also reused by the quote-aware Gemma pre-pass. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) -_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ _TC_JSON_CLOSED_PAT, _TC_GEMMA_CLOSED_PAT, + re.compile(r""), _TC_FUNC_CLOSED_PAT, - _TC_GEMMA_END_PAT, + # Mirror the parser regexes: tolerate whitespace and v11 [CALL_ID]/[ARGS] metadata. + re.compile( + r"\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*" + + _BRACKETED_JSON_ONE_LEVEL, + re.DOTALL, + ), + _REHEARSAL_CLOSED_STRIP_RE, + # Drop the bare v11 [/TOOL_CALLS] closer the balanced scan leaves behind. + re.compile(r"\[/TOOL_CALLS\]"), ] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r"<\|tool_call>.*$", re.DOTALL), +# Bare open markers strip a partial call mid-stream; the rehearsal tail needs `{` or EOF +# so prose ``foo[ARGS]`` survives. The XML open-tail forms reach EOF and are reused by +# _tool_call_markup_spans (a think tag in an unclosed call's args stays argument data). +_TOOL_OPEN_XML_TAIL_PATS = [ re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] -# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in -# their argument data cannot make the helper truncate the block and its tail. +_TOOL_ALL_PATS = ( + _TOOL_CLOSED_PATS + + _TOOL_OPEN_XML_TAIL_PATS + + [ + re.compile(r"\[TOOL_CALLS\].*$", re.DOTALL), + _REHEARSAL_TAIL_STRIP_RE, + ] +) + +# Rehearsal strips (name in group 1); name-gated via ``enabled_tool_names``, strip-all when None. +_REHEARSAL_STRIP_PATS = frozenset({_REHEARSAL_CLOSED_STRIP_RE, _REHEARSAL_TAIL_STRIP_RE}) + +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in argument +# data cannot make the helper truncate the block and its tail. _TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] -# A lazy closed-pair pattern whose close token is absent rescans to EOF from -# every opener (quadratic, re-run per streamed token); skip that doomed pass. +# A lazy closed-pair pattern whose close token is absent would rescan to EOF from every +# opener; skip that doomed (quadratic) pass. Shared by both strip helpers. _PAT_REQUIRED_TOKEN = { _TC_JSON_CLOSED_PAT: "", _TC_GEMMA_CLOSED_PAT: "", @@ -50,26 +103,102 @@ def strip_tool_patterns(text: str, patterns) -> str: return text +def apply_tool_strip_patterns( + text: str, + patterns, + enabled_tool_names = None, +) -> str: + """Apply strip ``patterns`` to ``text``. A bare rehearsal ``name[ARGS]{..}`` pattern + strips only when ``name`` is an enabled tool (or when ``enabled_tool_names`` is + ``None``); every other pattern is removed unconditionally. A closed-pair pattern whose + close token is absent is skipped so an unclosed-marker stream stays linear.""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + if enabled_tool_names is not None and pat in _REHEARSAL_STRIP_PATS: + text = pat.sub(lambda m: "" if m.group(1) in enabled_tool_names else m.group(0), text) + else: + text = pat.sub("", text) + return text + + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") -# Name class allows dots/hyphens for dotted Gemma names; whitespace-tolerant around -# ``call`` / ``:`` since drift emits ``call: name{`` and ``call : name{``. _TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# Horizontal whitespace only so the newline + value indentation survive (_trim_param_value trims one newline). +# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it. _TC_PARAM_START_RE = re.compile(r"[^\S\n]*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next -# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`, -# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value. +# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next +# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. +# `location:New York, NY`), so it must not terminate the value. The key token +# must be identifier-shaped (start with a letter or underscore); a comma +# followed by digits-then-colon is value text such as a timestamp or ratio +# (`meet at 10:00, 11:00 tomorrow`), not a new key. _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") +# A candidate starting inside a think block is a rehearsal (block kept so literal tags in +# real args survive); ``$`` accepts an unclosed block mid-stream. +_THINK_TAG_RE = re.compile(r".*?(?:|$)|\[THINK\].*?(?:\[/THINK\]|$)", re.DOTALL) +# Bare open/close markers for prefilled-reasoning turns (template opens in the prompt). +_THINK_OPEN_RE = re.compile(r"|\[THINK\]") +_THINK_CLOSE_RE = re.compile(r"|\[/THINK\]") + +# Mistral canonical array: [TOOL_CALLS] + JSON list of {"name","arguments"} objects. +_MISTRAL_ARRAY_RE = re.compile(r"\[TOOL_CALLS\]\s*(?=\[)") + +# Mistral name form + v11 [ARGS]/[CALL_ID] shapes; [CALL_ID] is metadata, not the name, +# and hyphens keep dashed MCP names whole. +_MISTRAL_BRACKET_RE = re.compile( + r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)" +) + +# Rehearsal ``name[ARGS]{json}`` (no [TOOL_CALLS]); the lookbehind keeps the v11 call-id +# from being taken as the function name. +_REHEARSAL_RE = re.compile(r"(? int | None: + """Return the end index of a balanced JSON object opening at ``start``, + or ``None`` if the braces don't balance. Honors escapes and strings. + """ + if start >= len(text) or text[start] != "{": + return None + depth = 0 + in_string = False + escape = False + for j in range(start, len(text)): + ch = text[j] + if escape: + escape = False + continue + if ch == "\\": + escape = True + continue + if in_string: + if ch == '"': + in_string = False + continue + if ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return j + return None + def _balanced_brace_end( content: str, @@ -134,6 +263,94 @@ def _balanced_bracket_end(src: str, start: int) -> int: return -1 +def _decode_array_items(text: str, body_start: int, body_end: int): + """Return ``(objs, ends)`` for each top-level element of the JSON array between + ``body_start`` (at or before its ``[``) and ``body_end`` (exclusive): the decoded + object and its absolute exclusive end offset. + + Decoding element-by-element with ``raw_decode`` tolerates the comma-less object + separators the repo's own Mistral/Ollama multi-call templates emit + (``[{...}{...}]``; see ollama_template_mappers.py). A single ``json.loads`` of the + whole body rejects that form and would drop every call. The ends also tile the + region across the calls' spans so a with_spans consumer strips each exactly once.""" + decoder = json.JSONDecoder() + objs: list = [] + ends: list[int] = [] + i = text.find("[", body_start) + if i < 0: + return objs, ends + i += 1 + while i < body_end: + while i < body_end and text[i] in " \t\r\n,": + i += 1 + if i >= body_end or text[i] == "]": + break + try: + obj, rel = decoder.raw_decode(text[i:body_end]) + except (json.JSONDecodeError, ValueError): + break + i += rel + objs.append(obj) + ends.append(i) + return objs, ends + + +def _iter_bracket_spans( + text: str, + start: int = 0, + enabled_tool_names = None, +): + """Yield ``(span_start, span_end, kind, match)`` for each balanced bracket-tag + call from ``start`` on, in document order; ``span_end`` exclusive. ``kind`` is + ``"array"`` ([TOOL_CALLS] [..]), ``"name"`` ([TOOL_CALLS]name{..}, incl. v11 + [CALL_ID]/[ARGS]) or ``"rehearsal"`` (name[ARGS]{..}). + + ``enabled_tool_names`` (set, or None = unrestricted) gates only the ambiguous + bare rehearsal form: name[ARGS]{..} is a call ONLY when ``name`` is enabled, so a + prose ``foo[ARGS]{..}`` (foo disabled) is neither parsed nor stripped. Explicit + [TOOL_CALLS] markers stay unconditional, keeping parse/strip/detection symmetric. + + Balance-only (no JSON validation) so strip and parse share one scan. The cursor + jumps past each consumed span, so a marker inside consumed JSON is never + re-matched and each regex re-searches only once its match falls behind: linear.""" + n = len(text) + specs = ( + ("array", _MISTRAL_ARRAY_RE), + ("name", _MISTRAL_BRACKET_RE), + ("rehearsal", _REHEARSAL_RE), + ) + nexts = {kind: rx.search(text, start) for kind, rx in specs} + cursor = start + while cursor < n: + for kind, rx in specs: + m = nexts[kind] + if m is not None and m.start() < cursor: + nexts[kind] = rx.search(text, cursor) + live = [(kind, m) for kind, m in nexts.items() if m is not None] + if not live: + return + kind, m = min(live, key = lambda km: km[1].start()) + if kind == "array": + end = _balanced_bracket_end(text, m.end()) + end = None if end < 0 else end + else: + end = _balanced_json_span(text, m.end()) + if end is None: + # Truncated body: skip and keep scanning; the caller's catch-all strips the tail. + cursor = m.end() + continue + if ( + kind == "rehearsal" + and enabled_tool_names is not None + and m.group(1) not in enabled_tool_names + ): + # Inactive-name rehearsal is prose: advance past its body without yielding. + cursor = end + 1 + continue + yield (m.start(), end + 1, kind, m) + cursor = end + 1 + + def _split_top_level_commas(src: str) -> list: """Split on commas that are not inside a nested ``[]``/``{}`` or a string.""" parts: list[str] = [] @@ -164,8 +381,14 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds: - quote bare strings, recurse into objects/arrays, keep quoted/JSON literals.""" + """Normalise the elements of a Gemma array value so json.loads succeeds. + + Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of + objects (``items:[{path:a}]``) whose keys/values also lack quotes; left + as-is json.loads fails and the whole call is dropped. Bare string elements + are quoted, object and nested-array elements are normalised recursively, and + quoted strings (already normalised from ``<|"|>``), numbers, and JSON + literals are preserved.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -173,9 +396,11 @@ def _quote_gemma_array_elements(body: str) -> str: out.append(element) continue if stripped[0] == "{": + # Object element: quote its keys/bare values like a top-level object. out.append(_quote_gemma_object_keys(stripped)) continue if stripped[0] == "[": + # Nested array: normalise its elements too. inner_end = _balanced_bracket_end(stripped, 0) if inner_end == len(stripped) - 1: out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") @@ -240,8 +465,6 @@ def _quote_gemma_object_keys(src: str) -> str: while i < len(src) and src[i].isspace(): i += 1 key_name_start = i - # Dots match the parser's key/name charset: Gemma emits dotted argument keys - # (user.name:...) for namespaced schemas. while i < len(src) and (src[i].isalnum() or src[i] in "_-."): i += 1 key_name = src[key_name_start:i] @@ -254,12 +477,15 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 - # Quote bare string values ({unit:celsius}); JSON stays as-is. + # Gemma may emit bare string values ({unit:celsius}); quote them so + # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. ws = i while i < len(src) and src[i].isspace(): i += 1 parts.append(src[ws:i]) if i < len(src) and src[i] == "[": + # Array value: quote bare string elements (e.g. labels:[bug,ui]) + # so json.loads succeeds instead of dropping the call. arr_end = _balanced_bracket_end(src, i) if arr_end < 0: parts.append(src[i:]) @@ -269,7 +495,9 @@ def _quote_gemma_object_keys(src: str) -> str: i = arr_end + 1 elif i < len(src) and src[i] not in '"{': v_start = i - # Bare value: up to `}` or a comma that starts the next key:pair. + # Consume the bare value up to `}` or a comma that starts the + # next key:value pair; a comma inside the value (e.g. + # `New York, NY`) does not terminate it. while i < len(src): if src[i] == "}": break @@ -329,7 +557,9 @@ def _func_close_index(content: str, body_start: int, body: str) -> int: def _trim_param_value(val: str) -> str: - """Trim only the wrapping newline (not str.strip) so code/diff argument indentation survives.""" + """Trim the single wrapping newline the chat template adds around an XML + parameter value, preserving indentation inside VALUE (``str.strip()`` destroyed + code/diff argument indentation).""" if val.startswith("\n"): val = val[1:] if val.endswith("\n"): @@ -404,6 +634,7 @@ def parse_tool_calls_from_text( *, id_offset: int = 0, allow_incomplete: bool = True, + enabled_tool_names = None, with_spans: bool = False, ): """Parse OpenAI-format tool calls from model text. @@ -412,22 +643,36 @@ def parse_tool_calls_from_text( {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... + [TOOL_CALLS]web_search{"query":"..."} (Mistral / Devstral fallback) + web_search[ARGS]{"query":"..."} (reasoning-model rehearsal) + + A call rehearsed inside a ```` / ``[THINK]`` block is skipped, not + executed; the block is kept so a literal tag in a real argument is preserved. With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]`` is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup in ``content`` (including its close tag when present), so a caller can remove exactly the parsed markup and keep every other byte intact. """ + # Candidates starting inside a think block are rehearsals, skipped; blocks are kept, and a + # think marker opening inside a call is argument data (excluded from spans). + _think_spans = _think_spans_outside_tool_markup(content) + _think_starts = [s for s, _e in _think_spans] + + def _in_think(pos: int) -> bool: + # Spans are ordered and non-overlapping; bisect gives O(log M) per candidate. + i = bisect.bisect_right(_think_starts, pos) - 1 + return i >= 0 and _think_spans[i][0] <= pos < _think_spans[i][1] + tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside - # another call's coverage, or an open value, is data not executed. - markers = _build_markers(content) - coverage = _marker_coverage(content, markers) + # Collect JSON/Gemma markers; _marker_coverage decides nesting so a marker inside + # another call's coverage (even one that failed to parse) is data, not executed. A + # marker opening inside a think block is a rehearsal and is skipped. parsed_items = [] # (start, span_end, name, arguments) in document order + markers = [mk for mk in _build_markers(content) if not _in_think(mk[0])] + coverage = _marker_coverage(content, markers) for idx, (start, brace_end, kind, m) in enumerate(markers): - # A marker starting inside another's coverage is that call's data. The - # end is exclusive so a marker at a close's end is an adjacent sibling. if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue if brace_end < 0: @@ -441,7 +686,7 @@ def parse_tool_calls_from_text( if kind == "json": obj = json.loads(content[m.end() - 1 : brace_end + 1]) name = obj.get("name", "") - # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). + # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes). arguments = obj.get("arguments") if arguments is None: arguments = obj.get("parameters", {}) @@ -452,7 +697,6 @@ def parse_tool_calls_from_text( arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - # Span reaches through the close tag when present, else just the braces. span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) @@ -461,14 +705,11 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) - # Function-XML calls promote in document order alongside marker calls (the - # #6801 contract). A inside any marker's coverage is excluded -- - # even if that marker failed to parse -- so nested XML cannot escape; one - # after a balanced close-less marker is a sibling, not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) + and not _in_think(fm.start()) and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): @@ -545,11 +786,170 @@ def parse_tool_calls_from_text( ) call_spans.append((start, span_end)) + # Patterns 3+4: Mistral [TOOL_CALLS] and bare rehearsal via one balanced scan in document + # order, so a Mistral call and a rehearsal in one message both parse. + if not tool_calls: + for start, end, kind, m in _iter_bracket_spans( + content, enabled_tool_names = enabled_tool_names + ): + if _in_think(start): + continue + # Extend the region over an immediately-following v11 closer so with_spans consumers strip it too. + closer = re.match(r"\s*\[/TOOL_CALLS\]", content[end:]) + region_end = end + closer.end() if closer else end + if kind == "array": + # Decode elements individually (comma-tolerant): one json.loads of the whole + # body rejects the comma-less multi-call arrays Mistral/Ollama templates emit. + payload, item_ends = _decode_array_items(content, m.end(), end) + if not payload: + continue + # Tile the region so every byte belongs to exactly one span; a with_spans consumer + # keeps skipped bytes visible and strips promoted markup exactly once. + tile_start = start + last_span_idx = -1 + for item_idx, item in enumerate(payload): + if not isinstance(item, dict) or "name" not in item: + continue + args = item.get("arguments", {}) + if isinstance(args, str): + # ``arguments`` may itself be a JSON string (OpenAI spec). + try: + args = json.loads(args) + except (json.JSONDecodeError, ValueError): + pass + if not isinstance(args, (dict, str)): + # ``"arguments": null`` (or any non-object scalar) becomes {} like the + # path, not the string "null" auto-heal would mangle to + # a bogus {"query":"null"}. + args = {} + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": item.get("name", ""), + # A bare scalar string stays raw (like the path); + # json.dumps would double-encode it so the arg healer wraps + # "weather" with its literal quotes. + "arguments": args if isinstance(args, str) else json.dumps(args), + }, + } + ) + item_end = item_ends[item_idx] if item_idx < len(item_ends) else region_end + last_span_idx = len(call_spans) + call_spans.append((tile_start, item_end)) + tile_start = item_end + if last_span_idx >= 0: + tile_start, _tile_end = call_spans[last_span_idx] + call_spans[last_span_idx] = (tile_start, region_end) + else: + try: + payload = json.loads(content[m.end() : end]) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(payload, dict): + continue + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(payload), + }, + } + ) + call_spans.append((start, region_end)) + if with_spans: return tool_calls, call_spans return tool_calls +def _strip_bracket_tag_calls(text: str, enabled_tool_names = None) -> str: + """Strip complete [TOOL_CALLS] arrays / name / bare name[ARGS]{..} calls with one + balanced forward scan, so nested JSON args are removed whole (a fixed-depth regex + left two-level args behind). Truncated tails go to the caller's catch-all. Linear. + ``enabled_tool_names`` gates the rehearsal form (inactive-name prose kept; None + strips every span).""" + if len(text) > _MAX_BRACKET_SCAN_CHARS: + return text + out: list[str] = [] + cursor = 0 + for start, end, _kind, _m in _iter_bracket_spans(text, enabled_tool_names = enabled_tool_names): + out.append(text[cursor:start]) + cursor = end + out.append(text[cursor:]) + return "".join(out) + + +def _tool_call_markup_spans(text: str) -> list[tuple[int, int]]: + """Spans of tool-call markup, so a literal /[THINK] inside a call's args is + stripped WITH the call, not kept as a reasoning block. Covers closed XML/bracket + calls and an unclosed XML call (run via allow_incomplete); without the open-ended + span the unclosed call's markup would leak after execution.""" + # Skip a lazy closed-pair pattern whose close token is absent: its finditer would rescan + # to EOF from every opener (quadratic on a stream of unclosed openers). + spans = [ + m.span() + for pat in _TOOL_CLOSED_PATS + if (_PAT_REQUIRED_TOKEN.get(pat) is None or _PAT_REQUIRED_TOKEN[pat] in text) + for m in pat.finditer(text) + ] + spans.extend((start, end) for start, end, _kind, _m in _iter_bracket_spans(text)) + # An unclosed opener is a real incomplete call only outside closed/bracket spans. + for pat in _TOOL_OPEN_XML_TAIL_PATS: + for m in pat.finditer(text): + if not any(s <= m.start() < e for s, e in spans): + spans.append(m.span()) + return spans + + +def _think_spans_outside_tool_markup(text: str) -> list[tuple[int, int]]: + """/[THINK] block spans, minus any whose opening marker sits INSIDE a + tool-call span (that tag is argument data, not reasoning). Keeping it would drop a + real call after it as rehearsed and leak the call's markup. START tested only, so + a greedy unclosed past the call is still that call's argument data.""" + think_spans = [m.span() for m in _THINK_TAG_RE.finditer(text)] + call_spans = _tool_call_markup_spans(text) + # Prefilled reasoning: the template opens in the prompt, so add a leading span + # (0..close) to skip calls rehearsed there; guarded so a stray close in a normal answer is safe. + close = _THINK_CLOSE_RE.search(text) + if close is not None: + opener = _THINK_OPEN_RE.search(text) + if ( + (opener is None or close.start() < opener.start()) + and not any(cs <= close.start() < ce for cs, ce in call_spans) + and any(cs >= close.end() for cs, ce in call_spans) + ): + think_spans = [(0, close.end())] + think_spans + if not think_spans: + return think_spans + if not call_spans: + return think_spans + return [(s, e) for (s, e) in think_spans if not any(cs <= s < ce for cs, ce in call_spans)] + + +def strip_outside_think(text: str, strip_segment) -> str: + """Apply ``strip_segment(segment, is_last)`` to visible text around /[THINK] + blocks, preserving the blocks verbatim (tool-looking text inside is rehearsal). + ``is_last`` is True only after the final block, so trailing-tail patterns apply + only there. Shared by every strip path so they stay consistent.""" + # A think marker opening inside a complete call is argument text; excluding it lets the + # stripper see the whole call. START-tested, so an unclosed match stays argument data. + think_spans = _think_spans_outside_tool_markup(text) + if not think_spans: + return strip_segment(text, True) + pieces: list[str] = [] + prev = 0 + for s, e in think_spans: + pieces.append(strip_segment(text[prev:s], False)) + pieces.append(text[s:e]) + prev = e + pieces.append(strip_segment(text[prev:], True)) + return "".join(pieces) + + def _strip_gemma_native_spans(text: str, *, final: bool) -> str: """Remove complete Gemma-native spans, brace/quote-balanced so a literal ```` in a quoted argument cannot truncate the span. An incomplete @@ -635,26 +1035,43 @@ def _strip_closed_blocks_outside_gemma(text: str) -> str: return text -def strip_tool_markup_final(text: str) -> str: - """Final display strip, shared with the streaming wrappers so all paths order - the passes identically: Gemma-aware closed JSON/function blocks first, then - well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed - spans and drop any unclosed remainder to EOF. Whitespace is kept.""" +def _strip_markup_segment( + text: str, + *, + final: bool, + enabled_tool_names = None, +) -> str: + # Bracket-tag calls (Mistral/rehearsal) first via balanced scan (any nesting depth, + # rehearsal name-gated); then the quote-aware Gemma-native passes so a literal + # in an argument cannot truncate a block; finally the regex XML/tail sweeps. + text = _strip_bracket_tag_calls(text, enabled_tool_names = enabled_tool_names) text = _strip_closed_blocks_outside_gemma(text) - text = _strip_gemma_native_spans(text, final = True) - return strip_tool_patterns(text, _TOOL_ALL_PATS) + text = _strip_gemma_native_spans(text, final = final) + patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS + return apply_tool_strip_patterns(text, patterns, enabled_tool_names = enabled_tool_names) -def strip_tool_call_markup(text: str, *, final: bool = False) -> str: +def strip_tool_call_markup( + text: str, + *, + final: bool = False, + enabled_tool_names = None, +) -> str: """Strip tool-call XML markup from text. When ``final`` is False, only fully closed tool-call blocks are removed. When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. + + ```` / ``[THINK]`` reasoning is preserved verbatim (see + ``strip_outside_think``); the trailing-tail patterns apply only after the + last block. ``enabled_tool_names`` keeps an inactive-name ``foo[ARGS]{..}`` + example visible (it is prose, not a call) so display cleanup matches detection. """ - if final: - return strip_tool_markup_final(text).strip() - # Non-final: same ordering as the final path, but incomplete blocks are kept. - text = _strip_closed_blocks_outside_gemma(text) - text = _strip_gemma_native_spans(text, final = False) - return strip_tool_patterns(text, _TOOL_CLOSED_PATS) + result = strip_outside_think( + text, + lambda seg, is_last: _strip_markup_segment( + seg, final = final and is_last, enabled_tool_names = enabled_tool_names + ), + ) + return result.strip() if final else result diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3341a9c628..1042dda004 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1366,17 +1366,60 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: return flags +def _generation_prompt_opens_think(template: Optional[str]) -> bool: + """True when rendering the template's generation prompt ends INSIDE an unclosed ````. + + Distinguishes templates that PREFILL an open ```` in the assistant generation + prompt (DeepSeek-R1, QwQ, Qwen3-Thinking) -- where the model emits only the closing + ```` and the extractor must start in reasoning mode -- from templates that merely + render PAST assistant ``...`` history while leaving the generation prompt + open with no ```` (e.g. Kimi-K2-Thinking), where the model self-emits its own block + and the extractor must start in normal mode. Renders a single-user-message probe with the + same sandbox transformers uses; on any failure returns True, preserving the historical + always-on prefill for templates that cannot be rendered here. + """ + if not template: + return False + try: + from jinja2.sandbox import ImmutableSandboxedEnvironment + + def _raise_exception(message: str): + raise RuntimeError(message) + + env = ImmutableSandboxedEnvironment( + trim_blocks = True, + lstrip_blocks = True, + extensions = ["jinja2.ext.loopcontrols"], + ) + env.filters["tojson"] = lambda value, **kwargs: json.dumps(value, ensure_ascii = False) + env.globals["raise_exception"] = _raise_exception + rendered = env.from_string(template).render( + messages = [{"role": "user", "content": "hi"}], + add_generation_prompt = True, + bos_token = "", + eos_token = "", + ) + except Exception: + return True + # ```` is not a substring of ```` (the ``/`` breaks it), so the last open + # tag sitting after the last close tag means the prompt ends inside an open block. + return rendered.rfind("") > rendered.rfind("") + + def _sf_reasoning_prefill_mode( features: dict, enable_thinking: Optional[bool], template: Optional[str] = None, reasoning_effort: Optional[str] = None, ) -> bool: - """Whether this request begins INSIDE an unclosed ```` (Qwen3/Qwen3.5/GLM prefill it). + """Whether a safetensors/MLX generation begins INSIDE an unclosed ````. - Gated on the STANDARD ````/```` markers: a bespoke reasoning channel (e.g. gemma) - never emits ````, so prefilled mode would swallow the whole answer -- excluded, as are - gpt-oss and thinking-disabled requests. ``enable_thinking=None`` defaults ON, so plain requests prefill. + ``enable_thinking`` templates (Qwen3/GLM) prefill an open ```` so the model + emits only the closing ````, and the extractor must start in reasoning mode. + Gated on the STANDARD ````/```` markers: bespoke channels (gemma's + ``<|think|>``) never emit ```` and would swallow the answer, so they and + gpt-oss and thinking-disabled requests return False. ``enable_thinking`` None + defaults thinking ON, so a plain request still prefills. """ if features.get("reasoning_style") not in ("enable_thinking", "enable_thinking_effort"): return False @@ -1384,16 +1427,21 @@ def _sf_reasoning_prefill_mode( if "" not in tpl and "" not in tpl: return False if features.get("reasoning_always_on"): - return True + # enable_thinking_effort + always-on: the effort mechanism (not the prompt shape) keeps + # thinking on, so always-on wins over reasoning_effort and we prefill. + if features.get("reasoning_style") == "enable_thinking_effort": + return True + # ``reasoning_always_on`` fires on paired ``...`` anywhere in the + # template, including markup that only renders PAST assistant history (Kimi-K2-Thinking) + # while the generation prompt opens none. Prefill only when the generation prompt opens + # one, else the extractor captures a normal answer as reasoning_content and returns blank. + return _generation_prompt_opens_think(tpl) if not features.get("supports_reasoning"): return False if enable_thinking is False: return False - # A reasoning_effort="none" request disables thinking for enable_thinking_effort - # (GLM-5.2) models the same way enable_thinking=False does (see - # ``_request_reasoning_kwargs``). Without this, the model emits no ```` and - # a plain answer is swallowed whole into reasoning_content, leaving the visible - # response empty. + # Thinking-off arrives as reasoning_effort "none" on enable_thinking_effort models; honor it + # so we don't prefill and capture the answer. Plain enable_thinking models ignore effort. if features.get("reasoning_style") == "enable_thinking_effort" and reasoning_effort == "none": return False return True @@ -1669,11 +1717,17 @@ def _apply_rag_nudge(nudge: str, tools: list[dict], *, rag_scope) -> str: return nudge + " " + _RAG_GROUNDING_NUDGE -# Strip leaked tool-call markup: every shared-parser format plus the four leak -# shapes llama_cpp.py's speculative buffer splits across the visible/DRAIN -# boundary. Mistral [TOOL_CALLS] uses the parser's balanced-brace helper (a -# non-greedy regex would truncate nested JSON); the DeepSeek opener alternation -# is the parser's own, so a signal we parse is never left un-stripped. +# Strip leaked tool-call markup: every shared-parser format plus the leak shapes +# llama_cpp.py's speculative buffer splits across the visible/DRAIN boundary: +# 1. well-formed `...` / `...` +# 2. orphan opening to EOF (close was DRAINED) +# 3. bare orphan close (open was DRAINED) +# 4. tail-only `` (outer close truncated by EOS); anchored to +# `\Z` so mid-text `` in user code samples survives. +# 5. Mistral `[TOOL_CALLS]name{json}` / rehearsal `name[ARGS]{json}`: the balanced +# scan removes the whole call (a non-greedy regex would truncate nested JSON). +# DeepSeek/GLM/Kimi envelopes are covered by the parser's own arms/scans, so a signal +# we parse is never left un-stripped; the DeepSeek opener alternation is the parser's own. from core.inference.tool_call_parser import _DEEPSEEK_OPEN_RE_SRC as _DS_OPEN_SRC _TOOL_XML_RE = _re.compile( @@ -1694,6 +1748,17 @@ _TOOL_XML_RE = _re.compile( r"|" r"|" r"|<\|python_tag\|>(?:[^<]|<(?!\|(?:eot_id|eom_id|python_tag|start_header_id|end_header_id|begin_of_text|finetune_right_pad_id)\|))*" + r"|\[/TOOL_CALLS\]" + # Truncated canonical array (closing ``]`` lost to EOS): the balanced scan cannot remove + # it, so strip its tail here. + r"|\[TOOL_CALLS\]\s*\[.*\Z" + # Named / v11 forms and bare rehearsal; arms aligned with the parser regexes. + r"|\[TOOL_CALLS\]\s*[\w-]+(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|.*?\Z)" + # Rehearsal: balanced/truncated body or bare marker at EOS only (prose ``foo[ARGS]`` + # survives); NAME captured as ``reh`` for the inactive-name display gate. + r"|(?[\w-]+)\[ARGS\]\s*(?:\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}|\{.*\Z|\Z)" + # DeepSeek envelopes (all opener variants), Kimi section blocks, and bare Kimi calls; + # each arm carries a call-shaped lookahead so prose merely mentioning a marker survives. r"|" + _DS_OPEN_SRC + r"(?=\s*(?:<｜tool▁call▁begin｜>|function)|\s*$).*?(?:<｜tool▁calls▁end｜>|\Z)" @@ -1705,6 +1770,17 @@ _TOOL_XML_RE = _re.compile( _re.DOTALL, ) +# Closed-only variant for segments before the last think block: the ``\Z``-anchored arms +# would treat a segment boundary as EOS and strip prose ``foo[ARGS]``. +_TOOL_XML_CLOSED_RE = _re.compile( + r"<(?:tool_call|function=[\w-]+)>.*?" + r"|<\|tool_call>.*?" + r"|" + r"|" + r"|\[/TOOL_CALLS\]", + _re.DOTALL, +) + def _gemma_strip_gate(tools) -> set: """Enabled tool NAMES gating the wrapper-less Gemma strip (mirrors the @@ -1720,18 +1796,18 @@ def _gemma_strip_gate(tools) -> set: return names -def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: - """Combine the parser's scan-based strips (Mistral balanced-brace, gated - Gemma wrapper-less, GLM real-close, guarded function-XML) with - ``_TOOL_XML_RE`` -- the scan strips close at each call's REAL terminator so - literal markup inside argument values is data, not a leaked tail. - ``enabled_tool_names`` gates the Gemma strip; ``None`` strips every closed call.""" - cleaned = _strip_glm_calls( - _strip_gemma_wrapperless_calls(_strip_mistral_closed_calls(text), enabled_tool_names), - final = True, - ) - cleaned = _strip_function_xml_calls(cleaned, final = True) - return _TOOL_XML_RE.sub("", cleaned) +def _display_tool_name_gate(active_tools): + """Active tool NAMES for gating the rehearsal display strip, or None when no tools + are enabled. ``None`` keeps the legacy strip-all behavior, mirroring the loop gate: + a bare ``NAME[ARGS]`` is a call only when NAME is active; without a tool list every + identifier stays ambiguous, so strip.""" + names = { + (t.get("function") or {}).get("name") + for t in (active_tools or []) + if isinstance(t, dict) and isinstance(t.get("function"), dict) + } + names.discard(None) + return names or None def _strip_tool_xml_for_display( @@ -1740,12 +1816,56 @@ def _strip_tool_xml_for_display( auto_heal_tool_calls: bool, enabled_tool_names: Optional[set] = None, ) -> str: - """Route-level leak cleanup (Auto-Heal only). Delegates to ``_strip_tool_xml`` - so the Mistral balanced-brace pass runs too (``_TOOL_XML_RE`` alone has no - ``[TOOL_CALLS]`` arm). ``enabled_tool_names`` gates the Gemma strip.""" + """Apply route-level XML leak cleanup only when Auto-Heal is enabled. + + Mirrors the parser-side segment scan: balanced strips first (Mistral, gated Gemma + wrapper-less, GLM real-close, guarded function-XML close at each call's REAL terminator + so literal markup inside a value is data), then the ``_TOOL_XML_RE`` arms cover the + DeepSeek / Kimi / orphan forms. ```` blocks are preserved verbatim and the + ``\\Z``-anchored tail arms run only on the last segment (prose ``foo[ARGS]`` before a + block survives). ``enabled_tool_names`` (when not None) gates the ambiguous bare-rehearsal + ``NAME[ARGS]{...}`` and wrapper-less Gemma ``call:NAME{...}`` strips on the active tool + list; an inactive NAME is prose and is kept. The ``[TOOL_CALLS]`` control-token arms strip + unconditionally regardless of NAME.""" if not auto_heal_tool_calls: return text - return _strip_tool_xml(text, enabled_tool_names) + from core.tool_healing import _strip_bracket_tag_calls, strip_outside_think + + def _keep_inactive_rehearsal(m) -> str: + # Only the bare-rehearsal arm captures ``reh``; with a tool list an inactive + # NAME[ARGS]{...} is prose -- keep it. + if enabled_tool_names is not None: + name = m.groupdict().get("reh") + if name is not None and name not in enabled_tool_names: + return m.group(0) + return "" + + def _strip_segment(seg: str, is_last: bool) -> str: + # Scan strips close at each call's REAL terminator (a literal ```` or a + # nested marker quoted inside a value cannot truncate the strip); the regex arms below + # cover the attribute form and the DeepSeek / Kimi / orphan families. + seg = _strip_mistral_closed_calls(seg) + seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names) + if is_last: + seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names) + seg = _strip_glm_calls(seg, final = is_last) + seg = _strip_function_xml_calls(seg, final = is_last) + if is_last: + return _TOOL_XML_RE.sub(_keep_inactive_rehearsal, seg) + return _TOOL_XML_CLOSED_RE.sub("", seg) + + return strip_outside_think(text, _strip_segment) + + +def _strip_tool_xml(text: str, enabled_tool_names: Optional[set] = None) -> str: + # Mistral balanced-brace pre-strip (kept explicit so the regression guards see it), then + # the shared think-aware display strip -- the one raw _TOOL_XML_RE.sub lives inside + # _strip_tool_xml_for_display, so every route cleanup site shares it. ``enabled_tool_names`` + # gates the Gemma wrapper-less strip; ``None`` strips every closed call. + text = _strip_mistral_closed_calls(text) + return _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = enabled_tool_names + ) logger = get_logger(__name__) @@ -6010,14 +6130,18 @@ async def openai_chat_completions( _gguf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _gguf_display_tool_names = _display_tool_name_gate(tools_to_use) # ── Strip stale tool-call XML from conversation history ─ for _msg in gguf_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): + # Gate on enabled tool names, like the live strip, so a documented inactive + # ``foo[ARGS]{...}`` survives in the replayed prompt context. _msg["content"] = _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ).strip() def gguf_generate_with_tools(): @@ -6151,7 +6275,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6258,7 +6382,7 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _gguf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(tools_to_use), + enabled_tool_names = _gguf_display_tool_names, ) return full_text, usage, finish finally: @@ -6631,14 +6755,17 @@ async def openai_chat_completions( _sf_tpl = (_sf_model_info.get("chat_template_info") or {}).get("template") _sf_features = _detect_safetensors_features(backend, _sf_tpl) - # Split prefilled-```` output into reasoning_content deltas (GGUF parity) so the UI - # renders the thinking block for safetensors AND MLX. + # GGUF parity: enable_thinking templates prefill an unclosed ; split into + # reasoning_content deltas so the UI renders the block for safetensors and MLX. _sf_parse_think = bool( _sf_features.get("supports_reasoning") or _sf_features.get("reasoning_always_on") ) - # Prefilled-open only for prefill styles with thinking on this request; gpt-oss excluded. + # Prefilled-open only for prefill styles with thinking on; gpt-oss uses the normal mode. _sf_reasoning_prefilled = _sf_reasoning_prefill_mode( - _sf_features, payload.enable_thinking, _sf_tpl, payload.reasoning_effort + _sf_features, + payload.enable_thinking, + _sf_tpl, + reasoning_effort = payload.reasoning_effort, ) def _new_sf_reasoning_extractor(): @@ -6722,6 +6849,8 @@ async def openai_chat_completions( _sf_auto_heal_tool_calls = ( payload.auto_heal_tool_calls if payload.auto_heal_tool_calls is not None else True ) + # Active tool names gating the bare-rehearsal strip, matching the loop gate. + _sf_display_tool_names = _display_tool_name_gate(_sf_tools_to_use) # Strip stale tool-call XML from prior assistant turns. _sf_chat_messages = [] @@ -6733,7 +6862,7 @@ async def openai_chat_completions( "content": _strip_tool_xml_for_display( _msg["content"], auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ).strip(), } ) @@ -6792,7 +6921,7 @@ async def openai_chat_completions( reasoning_extractor = _new_sf_reasoning_extractor() def _sf_flush_reasoning(): - # Drain the extractor at a turn boundary / stream end (GGUF parity); only visible text reaches the monitor. + # Drain the extractor at turn/stream end (mirrors GGUF); only visible text hits the monitor. fr, fv = reasoning_extractor.finish() out = [] if fr: @@ -6818,7 +6947,7 @@ async def openai_chat_completions( if event["type"] == "status": if not event["text"]: - # Iteration boundary: flush reasoning, then start a fresh extractor for the next turn. + # Iteration boundary: flush reasoning, then a fresh prefilled extractor for the next turn. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6834,7 +6963,7 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": - # Flush reasoning before the tool_start line so the thinking block closes ahead of the tool card. + # Flush reasoning before tool_start so the thinking block closes ahead of the card. for _c in _sf_flush_reasoning(): yield _c prev_text = "" @@ -6847,7 +6976,7 @@ async def openai_chat_completions( clean_cumulative = _strip_tool_xml_for_display( raw_cumulative, auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ) new_text = clean_cumulative[len(prev_text) :] prev_text = clean_cumulative @@ -6939,12 +7068,12 @@ async def openai_chat_completions( full_text = _strip_tool_xml_for_display( event.get("text", ""), auto_heal_tool_calls = _sf_auto_heal_tool_calls, - enabled_tool_names = _gemma_strip_gate(_sf_tools_to_use), + enabled_tool_names = _sf_display_tool_names, ) return full_text content_text = await asyncio.to_thread(_drain_to_text) - # Split prefilled reasoning out of the visible answer (GGUF parity); monitor gets visible text only. + # Split prefilled out of the visible answer (GGUF parity); the monitor gets visible text only. _reasoning_text, _visible_text = _extract_responses_reasoning( content_text, parse_think_markers = _sf_parse_think, @@ -7043,7 +7172,7 @@ async def openai_chat_completions( yield _chat_role_chunk(completion_id, created, model_name) prev_text = "" - # Split prefilled into reasoning_content deltas (GGUF parity). Single turn (no per-turn reset); also serves MLX. + # Split prefilled into reasoning_content deltas (GGUF parity); single turn, serves MLX. reasoning_extractor = _new_sf_reasoning_extractor() # Run the sync generator in a thread pool to avoid blocking the # event loop. Critical for compare mode: two SSE requests arrive @@ -7150,7 +7279,7 @@ async def openai_chat_completions( for token in generate(): full_text = token - # Split prefilled reasoning from the visible answer (GGUF parity); also covers MLX. + # Split prefilled reasoning (GGUF parity); also covers MLX via the shared generate(). _reasoning_text, _visible_text = _extract_responses_reasoning( full_text, parse_think_markers = _sf_parse_think, @@ -8000,8 +8129,8 @@ class _ResponsesReasoningExtractor: reasoning_prefilled: bool = False, ) -> None: self._buffer = "" - # ``reasoning_prefilled``: output begins INSIDE an unclosed ```` (Qwen3/GLM prefill), - # so start in reasoning to capture leading text until the first ````. Callers default False. + # reasoning_prefilled: the template inserts an unclosed , so output begins inside + # the block; start in reasoning until the first close tag. Existing callers pass False. self._in_reasoning = reasoning_prefilled # Splitting requires marker parsing; a prefilled open implies it. self._parse_think_markers = parse_think_markers or reasoning_prefilled @@ -8033,8 +8162,8 @@ class _ResponsesReasoningExtractor: self._buffer = self._buffer[close_idx + len(_RESPONSES_THINK_CLOSE) :] self._in_reasoning = False continue - # Hold back a trailing partial of EITHER marker: the close (clean chunk-boundary split) - # and a stray open (so a re-emitted ```` isn't leaked into the reasoning drawer). + # Hold back a trailing partial of either marker: the close (clean split across chunks) + # and a stray open (a re-emitted is suppressed, not leaked). keep = _responses_marker_holdback( self._buffer, (_RESPONSES_THINK_CLOSE, _RESPONSES_THINK_OPEN) ) @@ -9919,11 +10048,15 @@ async def anthropic_messages( else: openai_messages.insert(0, {"role": "system", "content": _nudge}) - # Strip stale tool-call XML from conversation + # Strip stale tool-call XML via the protected display helper (think rehearsal and [TOOL_CALLS] + # prose survive), gated on enabled tool names so documented inactive examples are kept. + _anthropic_history_gate = _display_tool_name_gate(openai_tools) for _msg in openai_messages: if _msg.get("role") == "assistant" and isinstance(_msg.get("content"), str): - _msg["content"] = _strip_tool_xml( - _msg["content"], _gemma_strip_gate(openai_tools) + _msg["content"] = _strip_tool_xml_for_display( + _msg["content"], + auto_heal_tool_calls = True, + enabled_tool_names = _anthropic_history_gate, ).strip() def _run_tool_gen(): @@ -10023,6 +10156,10 @@ async def _anthropic_tool_stream( """Streaming response for the tool-calling path.""" _sentinel = object() + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) + # Prompt-token count for message_start.usage.input_tokens. count_chat_tokens # makes blocking HTTP calls to llama-server, so run it off the event loop. # Pass the tools so tool-schema tokens are counted (the generator renders @@ -10074,9 +10211,15 @@ async def _anthropic_tool_stream( captured_finish_reason = _fr # Strip leaked tool-call XML from content events first, so a # content event that was purely tool XML doesn't count as text. + # Protected helper preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose (raw _TOOL_XML_RE.sub corrupts both). if etype == "content": event = dict(event) - event["text"] = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) + event["text"] = _strip_tool_xml_for_display( + event["text"], + auto_heal_tool_calls = True, + enabled_tool_names = _display_names, + ) # disable_parallel_tool_use: keep only the first tool_use block, # dropping every later tool_start and its paired tool_end (robust # to empty tool-call ids — tracked by state, not id matching). @@ -10250,6 +10393,9 @@ async def _anthropic_tool_non_streaming( usage = {} prev_text = "" captured_finish_reason = None + # Gate the display strip on the declared tools: an inactive NAME[ARGS]{...} in a final + # answer is prose and must survive in the delivered text. + _display_names = _display_tool_name_gate(openai_tools) # Pending client tool_use; cleared by tool_end (server execution) or # trailing text. See the stop_reason mapping below. ends_on_tool_use = False @@ -10259,8 +10405,10 @@ async def _anthropic_tool_non_streaming( for event in events: etype = event.get("type", "") if etype == "content": - # Strip leaked tool-call XML - clean = _strip_tool_xml(event["text"], _gemma_strip_gate(openai_tools)) + # Strip leaked tool XML (protected helper keeps think rehearsal and trailing prose). + clean = _strip_tool_xml_for_display( + event["text"], auto_heal_tool_calls = True, enabled_tool_names = _display_names + ) new = clean[len(prev_text) :] prev_text = clean if new: @@ -10730,13 +10878,16 @@ async def _anthropic_passthrough_non_streaming( text = message.get("content") or "" if text: # Keep unpromoted bytes when healing is active; legacy stripping is - # only for opted-out or no-client-tool requests. Use the full - # _strip_tool_xml pass so Mistral [TOOL_CALLS] and guarded - # function-XML leaks are cleaned too, not just _TOOL_XML_RE forms, - # with the Gemma display gate so a disabled/example call:NAME{...} - # in prose survives. + # only for opted-out or no-client-tool requests. Protected helper (not + # raw _TOOL_XML_RE.sub): preserves rehearsal and balanced + # [TOOL_CALLS] trailing prose, gated on the declared tools so an + # inactive NAME[ARGS]{...} example in the final text is kept. if not healing_active: - text = _strip_tool_xml(text, _gemma_strip_gate(openai_tools)) + text = _strip_tool_xml_for_display( + text, + auto_heal_tool_calls = True, + enabled_tool_names = _display_tool_name_gate(openai_tools), + ) text = text.strip() if text: content_blocks.append(AnthropicResponseTextBlock(text = text)) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index a6c1fcda9c..170b456eac 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -889,6 +889,24 @@ class TestAnthropicToolNonStreaming: assert tool_blocks[0]["name"] == "render_html" assert tool_blocks[0]["input"] == {"code": ""} + def test_display_strip_gates_on_declared_tools(self): + # A final answer containing NAME[ARGS]{json} is gated on the declared tools: undeclared + # ``foo`` markup is prose and survives, the declared web_search rehearsal strips. + def _run_gen(): + yield { + "type": "content", + "text": 'Try foo[ARGS]{"x": 1} but not web_search[ARGS]{"q": "hi"} here.', + } + + tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}] + response = asyncio.run( + _anthropic_tool_non_streaming(_run_gen, "msg_1", "m", openai_tools = tools) + ) + body = json.loads(response.body) + text = "".join(b["text"] for b in body["content"] if b["type"] == "text") + assert 'foo[ARGS]{"x": 1}' in text # inactive name preserved as prose + assert "web_search[ARGS]" not in text # active name stripped from display + # ===================================================================== # Pass-through emitter tests (client-side tool execution path) diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 7b653f47aa..e3055d2127 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -160,6 +160,20 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_unclosed_think_literal_inside_tool_argument_does_not_hide_later_call(): + # A literal inside a completed call's arguments is argument data; both calls must parse. + text = '[TOOL_CALLS]a{"x":"literal marker"} b[ARGS]{"y":2}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + + +def test_real_think_block_with_rehearsal_inside_still_skips_only_the_rehearsal(): + # A genuine reasoning block still hides its rehearsal while a real call after it parses. + text = 'web_search[ARGS]{"q":"draft"}real[ARGS]{"q":"go"}' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["real"], calls + + def test_wrapperless_nested_object_argument_is_parsed(): # skip_special_tokens stream: wrapper and <|"|> markers stripped, so a nested object arrives bare. calls = parse_tool_calls_from_text("call:f{loc:{city:NYC},n:3}") diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index dcc759a210..9e16be2160 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -2130,6 +2130,312 @@ def test_metadata_event_omits_prompt_tokens_details_when_absent(monkeypatch): assert "prompt_tokens_details" not in metadata[-1]["usage"] +def test_gguf_rehearsal_name_split_before_args_is_not_leaked(monkeypatch): + """Finding 6: a rehearsal call whose name (``web_search``) and ``[ARGS]{...}`` + arrive in separate content deltas must hold the bare name in the buffer until + ``[ARGS]`` flips it to a drain. Without _is_rehearsal_prefix the GGUF path + streams the tool name as visible content before the call executes.""" + + first_stream = [ + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, 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 cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_initial_buffer_flush_holds_split_rehearsal_name(monkeypatch): + """The first flush out of BUFFERING (prose plus a trailing active-tool-name in + the first delta, ``[ARGS]{...}`` in the next) must apply the same trailing-name + hold the STREAMING branch uses. The first delta has spaces so it is not a + rehearsal prefix and falls to the initial flush, which previously emitted the + bare name before the call drained.""" + + first_stream = [ + _sse({"content": "I will use web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + assert all("[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_name_after_prose_in_streaming_is_not_leaked(monkeypatch): + """Finding 9: the BUFFERING guard only covers a rehearsal at the turn start. + When prose has already streamed (STREAMING state) and the model then emits the + tool name and ``[ARGS]{...}`` in later deltas, the bare name must still be held, + not flushed as visible content before the call drains.""" + + first_stream = [ + _sse({"content": "Let me think. "}), + _sse({"content": "I will search "}), + _sse({"content": "web_search"}), + _sse({"content": '[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert all("web_search" not in t for t in content_texts), content_texts + + +def test_gguf_plain_answer_ending_with_tool_name_word_is_preserved(monkeypatch): + """End-of-stream flush: a plain answer that ENDS on a tool-name word with no + ``[ARGS]`` following is real prose and must not be dropped by the streaming + rehearsal hold.""" + + first_stream = [ + _sse({"content": "I think "}), + _sse({"content": "you should "}), + _sse({"content": "web_search"}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "advise"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert any(t.rstrip().endswith("web_search") for t in content_texts), content_texts + + +def test_gguf_long_tool_name_split_rehearsal_is_not_capped_and_executes(monkeypatch): + """Finding 11: a realistic MCP name longer than the 32-char buffer cap split as + NAME then [ARGS]{...} must still be held (a rehearsal prefix is self-bounding), + so the name does not leak and the call executes.""" + name = "mcp__github__create_pull_request" + assert len(name) >= 32, len(name) + + first_stream = [ + _sse({"content": name}), + _sse({"content": '[ARGS]{"x":1}'}), + _done(), + ] + final_stream = [_sse({"content": "done"}), _done()] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda n, a, **_k: (calls.append((n, a)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [(name, {"x": 1})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert not any(name in t for t in content_texts), content_texts + + +def test_gguf_streaming_keeps_bare_args_before_think_block(monkeypatch): + """F4: the GGUF streaming strip must run its open-ended ``[ARGS]`` tail cleanup + only on the LAST segment. A bare ``foo[ARGS]`` (no JSON body, ``foo`` not a tool) + before a block is prose, not a truncated call, so the final visible text + must keep it verbatim instead of dropping ``foo[ARGS]`` and corrupting the + sentence.""" + + first_stream = [ + _sse({"content": "Please pass foo[ARGS] "}), + _sse({"content": "pause "}), + _sse({"content": "to the template."}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + assert calls == [], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + assert content_texts, events + assert content_texts[-1] == "Please pass foo[ARGS] pause to the template." + + +def test_gguf_inactive_name_args_in_prose_is_not_drained(monkeypatch): + """BUG A: an inactive-name ``foo[ARGS]{...}`` in a prose answer must not be treated + as a tool call. The BUFFERING and end-of-stream safety-net ``[ARGS]`` checks gate on + active tool names (like the safetensors loop and the mid-stream path), so ``foo`` + (``web_search`` is the only enabled tool) is neither drained/parsed into a disabled + no-op nor forced into another generation turn.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"x":1} is just syntax.'}), + _done(), + ] + backend = _make_backend(monkeypatch, [first_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "x"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 2, + ) + ) + + # No tool executed for the inactive name; a spurious no-op re-prompt would exhaust the + # single supplied stream and error. + assert calls == [], calls + assert not any(e.get("type") in ("tool_start", "tool_end") for e in events), events + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive ``foo[ARGS]{...}`` is prose: the name-gated strip keeps the whole sentence. + assert any('foo[ARGS]{"x":1} is just syntax.' in t for t in content_texts), content_texts + + +def test_gguf_inactive_rehearsal_before_active_call_executes_and_keeps_prose(monkeypatch): + """BUG X (#5704): an inactive ``foo[ARGS]{...}`` before a real ``web_search[ARGS]{...}`` + in one delta must NOT swallow the real call; web_search executes while the inactive + rehearsal stays visible as prose.""" + first_stream = [ + _sse({"content": 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}'}), + _done(), + ] + final_stream = [_sse({"content": "Found cats."}), _done()] + backend = _make_backend(monkeypatch, [first_stream, final_stream], []) + + calls: list[tuple[str, dict]] = [] + monkeypatch.setattr( + "core.inference.tools.execute_tool", + lambda name, arguments, **_k: (calls.append((name, arguments)) or "result"), + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "search cats"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + # The real call runs; ``foo`` is not executed as a phantom disabled call. + assert calls == [("web_search", {"query": "cats"})], calls + content_texts = [e.get("text", "") for e in events if e.get("type") == "content"] + # The inactive rehearsal is preserved as prose; the active one is stripped. + assert any('foo[ARGS]{"a":1}' in t for t in content_texts), content_texts + assert all("web_search[ARGS]" not in t for t in content_texts), content_texts + + +def test_gguf_rehearsal_detection_recognises_spent_one_shot_with_original_tools(): + # Rehearsal detection is fed the ORIGINAL tool list, so a spent one-shot's re-emitted + # repeat is still detected (matching the strip gate) instead of blanking the turn. + from core.inference.llama_cpp import _gguf_has_genuine_tool_signal + from core.inference.tool_call_parser import TOOL_XML_SIGNALS + + repeat = 'render_html[ARGS]{"code":"x"}' + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, active_only) + assert _gguf_has_genuine_tool_signal(repeat, TOOL_XML_SIGNALS, original) + + +def test_gguf_rehearsal_prefix_and_tail_hold_recognise_spent_one_shot(): + # The BUFFERING prefix check and STREAMING/flush tail-holds use the ORIGINAL tool list, + # so a spent one-shot's split repeat is held rather than leaked as visible text. + from core.inference.llama_cpp import _held_rehearsal_tail_len, _is_rehearsal_prefix + + active_only = [{"type": "function", "function": {"name": "web_search"}}] + original = active_only + [{"type": "function", "function": {"name": "render_html"}}] + assert not _is_rehearsal_prefix("render_html", active_only) + assert _is_rehearsal_prefix("render_html", original) + assert _held_rehearsal_tail_len("answer render_html", active_only) == 0 + assert _held_rehearsal_tail_len("answer render_html", original) == len("render_html") + + def test_gguf_oversized_bare_json_not_leaked_and_executes(monkeypatch): """An oversized bare-JSON call drains rather than streams, and still executes via the safety net.""" diff --git a/studio/backend/tests/test_passthrough_healing.py b/studio/backend/tests/test_passthrough_healing.py index 06316e2243..83bcc5864a 100644 --- a/studio/backend/tests/test_passthrough_healing.py +++ b/studio/backend/tests/test_passthrough_healing.py @@ -280,6 +280,56 @@ class TestStreamHealer: assert [c["id"] for c in calls] == ["call_0", "call_1"] assert _events_text(events).strip() == "then" + def test_mistral_array_multiple_calls_all_promoted_in_stream(self): + # A canonical Mistral [TOOL_CALLS] array carries several calls under a + # SINGLE signal. Draining only the first call would leave the residue + # starting at ",{...}]" (no signal), so later calls in the same array + # must be promoted in the same pass, not flushed as raw text. + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = healer.feed(array) + healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert [c["id"] for c in calls] == ["call_0", "call_1"] + assert _events_text(events) == "" + + def test_mistral_array_multiple_calls_promoted_char_by_char(self): + healer = StreamToolCallHealer({"get_weather", "get_time"}) + array = ( + '[TOOL_CALLS][{"name":"get_weather","arguments":{"city":"Paris"}},' + '{"name":"get_time","arguments":{"tz":"UTC"}}]' + ) + events = [] + for ch in array: + events += healer.feed(ch) + events += healer.finalize() + calls = _events_calls(events) + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_time"] + assert _events_text(events) == "" + + def test_mistral_array_undeclared_middle_kept_as_text_others_promoted(self): + # A mid-array element for a tool that is not declared must survive as + # text while the declared neighbours on either side still promote in + # document order. + healer = StreamToolCallHealer({"a", "c"}) + array = ( + '[TOOL_CALLS][{"name":"a","arguments":{}},' + '{"name":"b","arguments":{}},{"name":"c","arguments":{}}]' + ) + events = healer.feed(array) + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "c"] + assert '"b"' in _events_text(events) + + def test_mistral_array_then_trailing_prose(self): + healer = StreamToolCallHealer({"a", "b"}) + array = '[TOOL_CALLS][{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + events = healer.feed(f"{array} all done") + healer.finalize() + assert [c["function"]["name"] for c in _events_calls(events)] == ["a", "b"] + assert "all done" in _events_text(events) + def test_incomplete_call_healed_at_finalize(self): healer = StreamToolCallHealer({"Bash"}) events = healer.feed('{"name":"Bash","arguments":{"cmd":"ls"}}') @@ -1356,3 +1406,42 @@ class TestOpenaiStreamingRoute: assert chunks[0] == line + "\n\n" # byte-for-byte relay asyncio.run(_run()) + + +class TestHealerSignalAlignment: + """The passthrough healer buffers only formats its parser can promote. + The loops' bare [ARGS] rehearsal signal is gated on active tool names + there; ungated in the healer it would stall legitimate prose until + finalization without ever producing a promotable call.""" + + def test_heal_signals_are_promotable_formats_only(self): + from core.inference.passthrough_healing import _HEAL_SIGNALS + assert set(_HEAL_SIGNALS) == { + "", + "<|tool_call>", + ", so -# generation begins inside the think block and emits only the closing ; the extractor starts in reasoning. +# reasoning_prefilled: enable_thinking templates prefill an unclosed , so +# generation begins inside the block; the extractor must start in reasoning. class TestReasoningPrefilledExtractor: def test_prefilled_single_feed_splits_lone_close(self): # T1: reasoning...answer with a prefilled (unseen) open tag. @@ -2077,9 +2077,7 @@ class TestReasoningPrefilledExtractor: assert visible == "hi" def test_not_prefilled_lone_close_preserves_current_behavior(self): - # T9: GGUF-parity guard -- WITHOUT prefilled, a lone keeps the - # pre-fix behavior (reasoning stays visible, tag dropped). Ensures GGUF and - # every existing caller are byte-identical. + # T9: without prefilled, a lone close tag keeps the pre-fix behavior (parity guard). reasoning, visible = _extract_responses_reasoning( "reasoningans", parse_think_markers = True, @@ -2099,8 +2097,7 @@ class TestReasoningPrefilledExtractor: assert visible == "v" def test_prefilled_ignored_when_markers_not_parsed(self): - # T11: a non-reasoning model (parse_think_markers False) still passes text - # straight through even if reasoning_prefilled were mistakenly set False. + # T11: a non-reasoning model passes text through even with reasoning_prefilled False. reasoning, visible = _extract_responses_reasoning( "just an answer", parse_think_markers = False, diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 3701a00dd2..9fd1535f22 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -183,7 +183,9 @@ def test_detect_safetensors_features_llama3_template_keeps_tools_on(): def test_detect_safetensors_features_mistral_template_keeps_tools_on(): - """Mistral emits [TOOL_CALLS]; parser now supports it.""" + """Mistral emits [TOOL_CALLS]name{json}, which the safetensors loop now parses + (the shared bracket-tag parser). The gate must no longer suppress it, or the + PR's Mistral tool support is unreachable through normal capability detection.""" from routes.inference import _detect_safetensors_features backend = SimpleNamespace(active_model_name = "unsloth/mistral-7b-instruct-v0.3") @@ -706,13 +708,28 @@ def test_detect_safetensors_features_keeps_tools_for_function_alias_bare_json(): assert flags["supports_tools"] is True -# _sf_reasoning_prefill_mode gates the prefilled- extractor so safetensors/MLX reach -# GGUF reasoning-block parity for enable_thinking models. +# _sf_reasoning_prefill_mode gates the prefilled- extractor (GGUF reasoning parity). class TestSafetensorsReasoningPrefillGate: # A minimal Qwen3-style template with the standard / markers. _QWEN_TPL = "{% if enable_thinking %}{% endif %}......" # gemma-style bespoke reasoning channel -- no standard markers. _GEMMA_TPL = "{% if enable_thinking %}<|think|>{% endif %}<|channel>thought" + # always-on template whose GENERATION PROMPT opens an unclosed (DeepSeek-R1 / QwQ / + # Qwen3-Thinking shape): the model emits only the closing , so prefill. + _ALWAYS_ON_OPEN_TPL = ( + "{% for m in messages %}{{ m['content'] }}{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>\n{% endif %}" + ) + # always-on template that renders PAST assistant ... history but leaves the + # generation prompt open with no (Kimi-K2-Thinking shape): the model self-emits its + # own block, so prefill mode would blank a normal answer. + _ALWAYS_ON_HISTORY_TPL = ( + "{% for m in messages %}" + "{% if m['role'] == 'assistant' %}{{ m.get('reasoning_content', '') }}" + "{{ m['content'] }}{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}<|im_assistant|>assistant<|im_middle|>{% endif %}" + ) def _features(self, **over): base = { @@ -756,11 +773,19 @@ class TestSafetensorsReasoningPrefillGate: feats = self._features(supports_reasoning = False, reasoning_style = None) assert _sf_reasoning_prefill_mode(feats, True, self._QWEN_TPL) is False - def test_g7_reasoning_always_on(self): - # G7: hardcoded- template -> prefilled regardless of the flag. + def test_g7_reasoning_always_on_prompt_opens_think(self): + # G7: always-on template whose generation prompt opens -> prefilled regardless of the flag. from routes.inference import _sf_reasoning_prefill_mode feats = self._features(reasoning_always_on = True) - assert _sf_reasoning_prefill_mode(feats, False, self._QWEN_TPL) is True + assert _sf_reasoning_prefill_mode(feats, False, self._ALWAYS_ON_OPEN_TPL) is True + + def test_g7b_reasoning_always_on_history_only_not_prefilled(self): + # G7b (#5704): always-on classification from rendered assistant HISTORY + # (Kimi-K2-Thinking) whose generation prompt opens no . Prefill mode would capture a + # normal answer entirely as reasoning_content and blank the visible answer, so it must be off. + from routes.inference import _sf_reasoning_prefill_mode + feats = self._features(reasoning_always_on = True) + assert _sf_reasoning_prefill_mode(feats, None, self._ALWAYS_ON_HISTORY_TPL) is False def test_g8_gemma_bespoke_channel_excluded(self): # G8: gemma's <|think|>/<|channel> format has no -> NOT prefilled diff --git a/studio/backend/tests/test_safetensors_reasoning_stream.py b/studio/backend/tests/test_safetensors_reasoning_stream.py index 4a5423fa87..4e708139b7 100644 --- a/studio/backend/tests/test_safetensors_reasoning_stream.py +++ b/studio/backend/tests/test_safetensors_reasoning_stream.py @@ -3,9 +3,11 @@ """Safetensors/MLX reasoning-block parity with GGUF. -enable_thinking templates prefill an unclosed ````, so the stream must split the leading -text into ``reasoning_content`` deltas (per turn, monitor gets visible text only). Replays a copy -of ``sf_tool_stream``'s reasoning loop from routes/inference.py against synthetic events. +enable_thinking templates (Qwen3/GLM) prefill an unclosed ```` so the model +emits only the closing ```` then the answer; the safetensors stream must +split the leading text into ``reasoning_content`` deltas (plain stream and tool +loop), resetting per turn and appending only visible text to the monitor. Replays a +copy of ``sf_tool_stream``'s reasoning loop against synthetic events. """ from __future__ import annotations @@ -24,8 +26,40 @@ from routes.inference import ( ) +_THINK_TPL = "........." +_ETHINK = {"reasoning_style": "enable_thinking", "supports_reasoning": True} +_ETHINK_EFFORT = {"reasoning_style": "enable_thinking_effort", "supports_reasoning": True} + + +def test_prefill_mode_on_for_enable_thinking_default(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, _THINK_TPL) is True + + +def test_prefill_mode_off_when_thinking_disabled(): + assert _sf_reasoning_prefill_mode(_ETHINK, False, _THINK_TPL) is False + + +def test_prefill_mode_off_for_reasoning_effort_none(): + # enable_thinking_effort turns thinking off via reasoning_effort="none"; prefilled mode + # would capture the whole answer as reasoning_content. + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "none") + is False + ) + assert ( + _sf_reasoning_prefill_mode(_ETHINK_EFFORT, None, _THINK_TPL, reasoning_effort = "high") + is True + ) + + +def test_prefill_mode_off_without_think_markers(): + assert _sf_reasoning_prefill_mode(_ETHINK, None, "no markers here") is False + + def _replay_sf_reasoning_stream(events: list[dict], *, prefilled: bool) -> dict: - """Mirror sf_tool_stream's reasoning loop: diff cumulative snapshots, reset (flushing) on turn end.""" + """Mirror sf_tool_stream's reasoning loop: diff each cumulative ``content`` + snapshot, feed the delta through the extractor, and reset (flushing first) on + ``tool_start`` / empty ``status`` so each turn splits independently.""" prev_text = "" extractor = _ResponsesReasoningExtractor( parse_think_markers = True, reasoning_prefilled = prefilled @@ -151,9 +185,6 @@ def test_s5_thinking_off_no_reasoning_deltas(): assert out["monitor"] == "Just the plain answer, no thinking." -_THINK_TPL = "...{% if enable_thinking %}{% endif %}......" - - def test_s6_reasoning_effort_none_disables_prefill_for_enable_thinking_effort(): # GLM-5.2-style enable_thinking_effort: a request with reasoning_effort="none" (and # enable_thinking omitted) disables thinking exactly like enable_thinking=False, so diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 38b30fe8f6..f826f3cddf 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -115,6 +115,22 @@ class TestParser: assert result[0]["function"]["name"] == "python" assert "print('hi')" in result[0]["function"]["arguments"] + def test_xml_param_preserves_leading_indentation(self): + import json + + # Only the wrapping newline is trimmed; code-argument indentation survives. + text = ( + "\n" + " indented = 1\n" + " more\n" + "" + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "code": " indented = 1\n more" + } + def test_xml_unclosed(self): # Closing tags omitted; parser must still extract the value. text = "ls -la" @@ -183,6 +199,8 @@ class TestParser: assert has_tool_signal("blah x") assert has_tool_signal("blah <|tool_call>call:terminal") assert has_tool_signal("hi ...") + assert has_tool_signal("ok [TOOL_CALLS]web_search{...") + assert has_tool_signal("fine python[ARGS]{...") assert not has_tool_signal("hello world") def test_render_html_start_detector_uses_first_tool(self): @@ -197,6 +215,44 @@ class TestParser: '{"name":"python","arguments":{"code":""}}' ) + def test_render_html_start_detector_covers_mistral_and_rehearsal_forms(self): + # The provisional render-html card must fire for bracket-tag forms too, not only XML. + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html{"code":""}') + assert _detect_render_html_tool_start('[TOOL_CALLS]render_html[ARGS]{"code":"x"}') + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"name":"render_html","arguments":{}}]' + ) + assert _detect_render_html_tool_start('render_html[ARGS]{"code":""}') + # A different first tool (or a prose mention with no JSON body) must not fire. + assert not _detect_render_html_tool_start('[TOOL_CALLS]web_search{"q":"x"}') + assert not _detect_render_html_tool_start('web_search[ARGS]{"q":"x"}') + assert not _detect_render_html_tool_start('python[ARGS]{"code":"render_html[ARGS]{}"}') + assert not _detect_render_html_tool_start("use render_html[ARGS] to render") + + def test_render_html_start_detector_skips_think_block_rehearsal(self): + # A render_html rehearsed inside think must not fire the card; the outside-think call decides. + assert not _detect_render_html_tool_start( + 'draft render_html[ARGS]{"code":"x"}python[ARGS]{"code":"print(1)"}' + ) + assert not _detect_render_html_tool_start( + '[THINK]render_html[ARGS]{"code":"x"}[/THINK]web_search[ARGS]{"q":"y"}' + ) + # A real render_html AFTER a rehearsed non-render_html inside think still fires. + assert _detect_render_html_tool_start( + 'web_search[ARGS]{"q":"x"}render_html[ARGS]{"code":""}' + ) + # A render_html rehearsed inside think with no real call after does not fire. + assert not _detect_render_html_tool_start('render_html[ARGS]{"code":"x"}') + + def test_render_html_start_detector_reads_top_level_array_name(self): + # Array form: the name is the object's top-level ``"name"``, not an argument key. + assert not _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"render_html"},"name":"python"}]' + ) + assert _detect_render_html_tool_start( + '[TOOL_CALLS] [{"arguments":{"name":"python"},"name":"render_html"}]' + ) + def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" @@ -237,6 +293,376 @@ class TestParser: == "before " ) + # Mistral [TOOL_CALLS] bracket-tag. + + def test_mistral_bracket_basic(self): + # Devstral / Mistral-Small fallback when bypassing native FC. + text = '[TOOL_CALLS]web_search{"query":"weather"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert isinstance(result[0]["function"]["arguments"], str) + assert "weather" in result[0]["function"]["arguments"] + + def test_rehearsal_inside_unclosed_think_is_ignored(self): + """Rehearsal-shaped markup inside an unclosed block must + not be executed as a real tool call. Mid-stream the + tag has not arrived yet, so the strip regex has to accept + end-of-string as a terminator. Regression for the Gemini + high-severity flag on this PR.""" + text = ( + "I should call web_search[ARGS]" '{"query":"weather"} next to find the answer.' + ) + result = parse_tool_calls_from_text(text) + # Inside an unclosed think block no calls are yielded. + assert result == [] + + def test_rehearsal_inside_unclosed_bracket_think_is_ignored(self): + text = "[THINK]planning to use python[ARGS]" '{"code":"print(1)"} but not yet.' + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_rehearsal_after_closed_think_still_parsed(self): + text = "planning" 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_inside_prefilled_think_is_ignored(self): + """Reasoning models (Qwen3.5 enable_thinking) open in the PROMPT, + so generated content starts inside the thought and carries only a closing + . A call rehearsed in that leading thought must be skipped, while a + real call after the close still fires.""" + text = 'planning web_search[ARGS]{"query":"draft"}python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_literal_close_think_in_leading_argument_not_prefill(self): + """A literal inside a real leading call's arguments must not be + read as a prefilled-reasoning close (which would skip the call).""" + text = 'web_search[ARGS]{"query":"what is "}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_stray_close_after_real_call_not_treated_as_prefill(self): + """A real leading call followed by a stray and no further call is + a normal answer, not prefilled reasoning; the call must still fire (the + virtual span only applies when a real call follows the close).""" + text = 'Now web_search[ARGS]{"query":"x"} answer' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_with_whitespace(self): + # Optional whitespace (incl. newlines) between the name and the opening brace. + text = '[TOOL_CALLS]python \n {"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_mistral_bracket_nested_json(self): + # Brace-balance scan handles nested objects and braces inside string literals. + text = "[TOOL_CALLS]web_search" '{"query":"a {nested} brace","opts":{"limit":5}}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + import json as _json + + args = _json.loads(result[0]["function"]["arguments"]) + assert args["query"] == "a {nested} brace" + assert args["opts"] == {"limit": 5} + + def test_mistral_bracket_with_prose(self): + # Bracket-tag surrounded by prose is still recognised. + text = ( + "Sure, I will look that up.\n" + '[TOOL_CALLS]web_search{"query":"weather"}\n' + "Calling now." + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_mistral_bracket_bad_json_dropped(self): + text = "[TOOL_CALLS]web_search{not valid}" + result = parse_tool_calls_from_text(text) + # No usable tool call; callers fall back to text. + assert result == [] + + def test_mistral_bracket_object_with_array_value(self): + # Args must be a JSON object; a dict wrapping an array value is accepted. + text = '[TOOL_CALLS]web_search{"opts":[1,2,3]}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + # Rehearsal syntax name[ARGS]{json}. + + def test_rehearsal_basic(self): + text = 'python[ARGS]{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + assert "print(1)" in result[0]["function"]["arguments"] + + def test_rehearsal_with_prose(self): + text = "I should call the python tool. Like this: " 'python[ARGS]{"code":"x = 1"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_rehearsal_bad_json_dropped(self): + text = "python[ARGS]{not valid json}" + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_mistral_bracket_hyphenated_mcp_name(self): + # Dashed MCP names must be captured whole, not truncated at the first dash. + text = '[TOOL_CALLS]mcp__srv__list-issues{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_rehearsal_hyphenated_mcp_name(self): + text = 'mcp__srv__list-issues[ARGS]{"q":"x"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__list-issues" + + def test_streaming_strip_removes_partial_bracket_marker(self): + # A bracket tag streamed before its opening brace must strip on the final pass, not leak. + assert strip_tool_markup("answer [TOOL_CALLS]web_search", final = True) == "answer" + assert strip_tool_markup("text python[ARGS]", final = True) == "text" + # Non-final must keep the in-progress tag buffered (not yet stripped). + partial = "answer [TOOL_CALLS]web_search" + assert strip_tool_markup(partial, final = False) == partial + + def test_strip_removes_two_level_nested_bracket_call_keeps_prose(self): + # Two-level-nested args must be removed whole; the balanced scan handles any depth. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + assert strip_tool_markup(text, final = False) == "before after" + assert strip_tool_markup(text, final = True) == "before after" + + def test_strip_removes_call_with_literal_think_in_argument(self): + # A literal think block inside arguments strips with the call, not as a reasoning block. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + assert strip_tool_markup(text, final = True) == "" + + def test_strip_preserves_real_think_but_strips_call_with_literal_think(self): + text = ( + "planning ok " + '{"name":"w","arguments":{"t":"x"}} done' + ) + out = strip_tool_markup(text, final = True) + assert "planning" in out + assert "" not in out and '"name"' not in out + assert "ok" in out and "done" in out + + def test_prose_mentioning_args_marker_is_not_truncated(self): + # ``foo[ARGS] to the template`` is prose; the catch-all must not delete the sentence. + text = "Please pass foo[ARGS] to the template and continue reading." + assert strip_tool_markup(text, final = True) == text + + def test_streaming_strip_handles_mistral_v11_call_id_args(self): + # The streaming strip uses the regex patterns directly, so they must cover the v11 + # [CALL_ID]/[ARGS] metadata (aligned with the parser). + raw = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = strip_tool_markup_streaming(raw) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + # pre-strip. + + def test_think_block_stripped_before_xml(self): + # The think block is stripped before matching so the post-thinking call is recognised. + text = ( + "I will use web_search to find the weather." + '{"name":"web_search","arguments":{"query":"sf"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_think_block_stripped_before_bracket_tag(self): + text = ( + "Let me search for that.\n" '[TOOL_CALLS]web_search{"query":"weather"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + + def test_uppercase_think_tag_stripped(self): + # Some templates use [THINK]...[/THINK] instead of . + text = "[THINK]planning my next call[/THINK]" '[TOOL_CALLS]python{"code":"print(1)"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "python" + + def test_think_block_hides_inner_tool_call(self): + # A call mentioned inside think is a rehearsal; the wrapper strip removes the inner markup. + text = ( + "I might call " + '{"name":"web_search","arguments":{}} ' + "but I am not sure\n" + "Let me just answer directly." + ) + result = parse_tool_calls_from_text(text) + assert result == [] + + def test_think_literal_inside_real_tool_argument_is_preserved(self): + # A real call whose argument contains a literal think tag must not be corrupted. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["text"] == ( + "compare and tags" + ) + + def test_bracket_tag_argument_with_think_literal_is_preserved(self): + text = '[TOOL_CALLS]search{"q":"explain [THINK] blocks"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "explain [THINK] blocks" + + def test_real_call_after_think_with_rehearsal_inside(self): + # A rehearsal inside is skipped, but the real call after the close tag parses. + text = 'plan: search[ARGS]{"q":"x"}search[ARGS]{"q":"real"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"])["q"] == "real" + + # XML takes precedence over bracket-tag. + + def test_xml_wins_over_bracket(self): + # When a model emits both forms in one message, the XML form is canonical and wins. + text = ( + '{"name":"primary","arguments":{}}' + '[TOOL_CALLS]secondary{"k":"v"}' + ) + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "primary" + + # Strip patterns include bracket-tag and rehearsal. + + def test_strip_bracket_tag_closed(self): + text = 'before [TOOL_CALLS]web_search{"q":"hi"} after' + assert "[TOOL_CALLS]" not in strip_tool_markup(text) + assert "before" in strip_tool_markup(text) + assert "after" in strip_tool_markup(text) + + def test_strip_rehearsal_closed(self): + text = 'prose python[ARGS]{"code":"x"} more prose' + cleaned = strip_tool_markup(text) + assert "[ARGS]" not in cleaned + assert "prose" in cleaned + assert "more prose" in cleaned + + def test_strip_bracket_tag_unclosed_final(self): + text = 'before [TOOL_CALLS]web_search{"q":"part' + # Final-mode strip drops the trailing unclosed run. + cleaned = strip_tool_markup(text, final = True) + assert "TOOL_CALLS" not in cleaned + assert cleaned == "before" + + # Canonical Mistral array, v11 [CALL_ID], unified multi-call (PR review fixes). + + def test_mistral_canonical_array_is_parsed(self): + # Canonical multi-call array: every call must parse (was dropped then deleted to EOS). + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}},{"name":"b","arguments":{"y":2}}]' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(result[1]["function"]["arguments"]) == {"y": 2} + + def test_mistral_array_string_arguments_are_decoded(self): + # OpenAI-spec arguments arrive as a JSON string; decode to an object. + text = '[TOOL_CALLS] [{"name":"a","arguments":"{\\"x\\":1}"}]' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == {"x": 1} + + def test_mistral_array_scalar_string_argument_not_double_encoded(self): + # A bare scalar string argument in the Mistral array form must be kept + # raw, exactly like the path, so the downstream argument + # healer wraps ``weather`` into the single-string tool's key -- not + # ``"weather"`` with literal quotes from a redundant json.dumps. + array = parse_tool_calls_from_text( + '[TOOL_CALLS][{"name":"web_search","arguments":"weather"}]' + ) + xml = parse_tool_calls_from_text( + '{"name":"web_search","arguments":"weather"}' + ) + assert array[0]["function"]["arguments"] == xml[0]["function"]["arguments"] == "weather" + healed = _coerce_arguments( + array[0]["function"]["arguments"], heal = True, tool_name = "web_search" + ) + assert healed == {"query": "weather"} + + def test_mistral_array_strip_keeps_trailing_prose(self): + # The array form must be removed whole, not deleted to end-of-string. + text = 'answer [TOOL_CALLS] [{"name":"a","arguments":{}}] tail' + assert strip_tool_markup(text, final = True) == "answer tail" + + def test_mistral_and_rehearsal_in_one_message_both_parse(self): + # A Mistral call and a rehearsal call together: both must parse. + text = '[TOOL_CALLS]a{"x":1} then b[ARGS]{"y":2}' + result = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in result] == ["a", "b"] + + def test_mistral_v11_call_id_is_not_the_function_name(self): + # v11 shape: the function name is ``name``, never the opaque call-id token. + result = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[CALL_ID]abc123[ARGS]{"q":"x"}') + assert len(result) == 1 + assert result[0]["function"]["name"] == "get_weather" + assert json.loads(result[0]["function"]["arguments"]) == {"q": "x"} + # v11 without a call-id parses the same name. + r2 = parse_tool_calls_from_text('[TOOL_CALLS]get_weather[ARGS]{"q":"y"}') + assert r2[0]["function"]["name"] == "get_weather" + + def test_strip_preserves_rehearsal_inside_think(self): + # A rehearsal inside is reasoning; strip keeps it verbatim. + text = 'plan: search[ARGS]{"q":"x"} A' + out = strip_tool_markup(text, final = True) + assert out == text + assert "search[ARGS]" in out + + def test_streaming_strip_preserves_rehearsal_inside_think(self): + # The streaming strip must also preserve a think rehearsal: a mid-stream strip shrinks + # then regrows the cumulative text (corrupts append-by-length consumers). Matches GGUF. + text = 'plan: search[ARGS]{"q":"x"} A' + assert strip_tool_markup_streaming(text) == text + assert strip_tool_markup_streaming(text, tool_protocol_active = True) == text + # An unclosed block during streaming is preserved too (the parser keeps it). + partial = 'plan: search[ARGS]{"q":"x"}' + assert strip_tool_markup_streaming(partial, tool_protocol_active = True) == partial + + def test_streaming_strip_still_removes_real_call_outside_think(self): + # The think guard must not stop the streaming strip removing a call outside the block. + text = 'reason web_search[ARGS]{"q":"x"}' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "reason" in out + + def test_strip_bracket_calls_is_linear(self): + # Many complete bracket calls must strip in ~linear time (was O(n^2) per match). + import time + + text = '[TOOL_CALLS]f{"a":1}' * 4000 # ~80KB, 4000 complete calls + t0 = time.perf_counter() + out = strip_tool_markup(text, final = True) + elapsed = time.perf_counter() - t0 + assert "[TOOL_CALLS]" not in out + assert elapsed < 1.0, f"strip took {elapsed * 1000:.0f}ms on 4000 bracket calls" + def test_streaming_strip_handles_nested_mistral_json(self): # The non-greedy [TOOL_CALLS]name{...} pattern truncates nested JSON at the first }; the # balanced helper must remove the whole call so no trailing brace leaks to the streaming ... @@ -1390,6 +1816,254 @@ def test_active_tools_are_passed_to_single_turn_after_render_html_success(): assert any(event.get("type") == "content" and event.get("text") == "Done." for event in events) +def test_spent_one_shot_rehearsal_repeat_is_detected_not_blank_continuation(): + # A spent one-shot (render_html) stays in the ORIGINAL tool list; detection is gated on + # that list (matching the strip gate) so a re-emitted repeat is drained and routed to the + # repeat no-op instead of stripped into a blank continuation. + exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) + turns = iter( + [ + [ + '{"name":"render_html","arguments":{"code":"one"}}' + ], + ['render_html[ARGS]{"code":"two"}'], # spent one-shot rehearsal + ["The chart is above."], + ] + ) + + def gen(_messages, *, active_tools = None): + try: + chunks = next(turns) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = gen, + messages = [{"role": "user", "content": "make a chart"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "web_search"}}, + ], + execute_tool = exec_fn, + max_tool_iterations = 5, + ) + ) + contents = [e["text"] for e in events if e["type"] == "content"] + # render_html ran exactly once; the repeat was a no-op, not a second execution. + assert exec_fn.calls == [("render_html", {"code": "one"})], exec_fn.calls + # The loop continued past the repeat to the real answer (not a blank continuation). + assert any("The chart is above." in t for t in contents), contents + # The raw rehearsal markup never leaked as visible content. + assert not any("render_html[ARGS]" in t for t in contents), contents + + +def test_rehearsal_call_name_is_not_streamed_before_args(): + # A rehearsal whose name and [ARGS] arrive together must drain, not stream the bare name. + loop, exec_fn = _make_loop( + turns = [['web_search[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_call_name_split_before_args_is_not_streamed(): + # Finding 5: name and [ARGS] in separate chunks -- the bare name is held until [ARGS] arrives. + loop, exec_fn = _make_loop( + turns = [["web_search", '[ARGS]{"query":"cats"}'], ["Found."]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_plain_word_matching_no_tool_still_streams(): + # The prefix guard must not swallow prose: a non-tool bare word streams. + loop, _exec = _make_loop( + turns = [["weather", " is nice today."]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "weather is nice today." in contents, contents + + +def test_rehearsal_name_after_prose_in_streaming_is_not_streamed(): + # After prose has streamed (STREAMING state), a split rehearsal name must still be held. + loop, exec_fn = _make_loop( + turns = [ + # _make_loop accumulates these deltas into cumulative snapshots. + ["Let me think. ", "I will search ", "web_search", '[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_rehearsal_name_after_prose_same_chunk_in_streaming_is_not_streamed(): + # Prose then ``web_search[ARGS]{...}`` in one chunk: the boundary is pulled back over the name. + loop, exec_fn = _make_loop( + turns = [ + ["Sure. ", 'now web_search[ARGS]{"query":"cats"}'], + ["Found."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_initial_buffer_flush_holds_split_rehearsal_name(): + # First flush out of BUFFERING applies the same trailing-name hold as STREAMING. + loop, exec_fn = _make_loop( + turns = [["I will use python", '[ARGS]{"code":"print(1)"}'], ["done"]], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("python", {"code": "print(1)"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("python" in t for t in contents), contents + + +def test_think_rehearsal_streams_monotonically_and_keeps_reasoning(): + # A think rehearsal streams the same text the final strip keeps: cumulative content is + # monotonically non-decreasing and ends with the markup intact. + loop, exec_fn = _make_loop( + turns = [["plan ", 'search[ARGS]{"q":"x"}', " visible"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + contents = [e["text"] for e in events if e["type"] == "content"] + assert exec_fn.calls == [], exec_fn.calls + assert all(len(b) >= len(a) for a, b in zip(contents, contents[1:])), contents + final = contents[-1] if contents else "" + assert 'search[ARGS]{"q":"x"}' in final, contents + assert "visible" in final, contents + + +def test_plain_answer_ending_with_tool_name_word_is_preserved(): + # End-of-stream flush: a plain answer ending on a tool-name word is prose, not dropped. + loop, exec_fn = _make_loop( + turns = [["I think ", "you should ", "web_search"]], + max_tool_iterations = 1, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any(t.rstrip().endswith("web_search") for t in contents), contents + + +def test_long_tool_name_split_rehearsal_is_not_capped_and_executes(): + # Finding 10/11: an MCP name longer than the buffer cap, split before [ARGS], is still + # held (self-bounding prefix); no leak and the call executes. + from core.inference.safetensors_agentic import _MAX_BUFFER_CHARS + + name = "mcp__github__create_pull_request" + assert len(name) >= _MAX_BUFFER_CHARS, len(name) + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([[name, name + '[ARGS]{"x":1}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [{"type": "function", "function": {"name": name}}], + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [(name, {"x": 1})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any(name in t for t in contents), contents + + +def test_unrestricted_mode_split_rehearsal_name_is_not_streamed(): + # Finding 6: unrestricted mode treats any bare identifier as a possible rehearsal NAME. + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search" in t for t in contents), contents + + +def test_unrestricted_mode_split_after_bracket_is_not_streamed(): + # Unrestricted mode: a chunk split right after ``NAME[`` is still held (parity with the + # restricted-mode startswith hold). + exec_fn = FakeExecuteTool(["RESULT"]) + _turns = iter([["web_search[", 'web_search[ARGS]{"q":"x"}'], ["done"]]) + + def st(_messages, active_tools = None): + yield from next(_turns) + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "go"}], + tools = [], # unrestricted + execute_tool = exec_fn, + max_tool_iterations = 2, + ) + ) + assert exec_fn.calls == [("web_search", {"q": "x"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert not any("web_search[" in t for t in contents), contents + + +def test_unrestricted_mode_plain_prose_still_streams(): + # The unrestricted hold releases a held identifier once the rest of the sentence follows. + def st(_messages, active_tools = None): + for snap in ("Hello", "Hello there friend."): + yield snap + + events = _collect_events( + run_safetensors_tool_loop( + single_turn = st, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = FakeExecuteTool([]), + max_tool_iterations = 1, + ) + ) + contents = "".join(e["text"] for e in events if e["type"] == "content") + assert "Hello there friend." in contents, contents + + def test_safety_net_honors_disabled_auto_heal_for_late_incomplete_call(): # A late call caught by the safety net: an unclosed ```` heals only with Auto-Heal on; # off, the safety net must not pass ``allow_incomplete=True`` and execute a truncated call. @@ -1977,6 +2651,42 @@ class TestLoopBasic: assert tool_starts[0]["tool_name"] == "python" assert exec_fn.calls == [("python", {"code": "print('')"})] + def test_render_html_rehearsed_in_think_block_emits_no_provisional_start(self): + # BUG B: a render_html rehearsed inside think before a real python call must not emit a + # provisional render_html card; only the outside-think call fires. + exec_fn = FakeExecuteTool(["ok"]) + turn_iter = iter( + [ + [ + 'draft render_html[ARGS]{"code":"x"}', + 'python[ARGS]{"code":"print(1)"}', + ], + ["Done."], + ] + ) + + def _gen(_messages): + chunks = next(turn_iter) + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "run code"}], + tools = [ + {"type": "function", "function": {"name": "render_html"}}, + {"type": "function", "function": {"name": "python"}}, + ], + execute_tool = exec_fn, + ) + events = _collect_events(loop) + tool_starts = [e for e in events if e["type"] == "tool_start"] + + assert [e["tool_name"] for e in tool_starts] == ["python"], tool_starts + assert exec_fn.calls == [("python", {"code": "print(1)"})] + def test_render_html_success_blocks_second_canvas_call(self): exec_fn = FakeExecuteTool(["Rendered HTML canvas."]) turn_iter = iter( @@ -3653,6 +4363,114 @@ if __name__ == "__main__": pytest.main([__file__, "-v"]) +def test_streaming_strip_keeps_bare_args_before_think_block(): + # F3: a bare ``foo[ARGS]`` before a think block is prose; EOS-anchored tail arms run only + # on the last segment. + text = "Please pass foo[ARGS] pause to the template." + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert out == text + + +def test_streaming_strip_still_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips in the non-last segment. + text = 'go web_search[ARGS]{"q":"x"} z done' + out = strip_tool_markup_streaming(text, tool_protocol_active = True) + assert "web_search[ARGS]" not in out + assert "z" in out + assert "go" in out and "done" in out + + +def test_prose_args_marker_before_real_call_does_not_drain_the_prose(): + # F5: an inactive ``foo[ARGS]`` in prose is not a call boundary; the prose streams in + # full and the later real call still executes. + loop, exec_fn = _make_loop( + turns = [ + ["Intro ", "foo[ARGS] syntax. ", 'web_search[ARGS]{"query":"cats"}'], + ["Cats are great."], + ], + exec_results = ["RESULT"], + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [("web_search", {"query": "cats"})], exec_fn.calls + contents = [e["text"] for e in events if e["type"] == "content"] + # The prose between the bogus marker and the real call must survive. + assert any("foo[ARGS] syntax." in t for t in contents), contents + # The real call markup is never shown as content. + assert not any("web_search[ARGS]" in t for t in contents), contents + + +def test_inactive_name_args_with_body_is_not_parsed_into_disabled_noop(): + # BUG A: a prose answer with an inactive ``foo[ARGS]{...}`` is not drained into a + # disabled no-op extra turn; the [ARGS] checks are name-gated. + turns = [['foo[ARGS]{"x":1} is just syntax.']] + turn_calls: list[int] = [] + + def _gen(_messages): + turn_calls.append(1) + chunks = turns[len(turn_calls) - 1] if len(turn_calls) <= len(turns) else [] + acc = "" + for chunk in chunks: + acc += chunk + yield acc + + exec_fn = FakeExecuteTool([]) + loop = run_safetensors_tool_loop( + single_turn = _gen, + messages = [{"role": "user", "content": "explain"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + execute_tool = exec_fn, + max_tool_iterations = 3, + ) + events = _collect_events(loop) + assert exec_fn.calls == [], exec_fn.calls + assert not any(e["type"] in ("tool_start", "tool_end") for e in events), events + # Exactly one generation turn -- no disabled ``foo`` no-op re-prompt. + assert len(turn_calls) == 1, turn_calls + contents = [e["text"] for e in events if e["type"] == "content"] + assert any("is just syntax." in t for t in contents), contents + + +class TestEnabledToolNameGate: + """The safetensors loop passes the active tool names into parse/strip so the + ambiguous bare-rehearsal ``NAME[ARGS]{json}`` is treated as a call only when NAME + is an active tool (#5704). Without the gate an inactive ``foo[ARGS]{...}`` in prose + was parsed into a disabled no-op call and stripped from the visible text.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_parse_inactive_rehearsal_does_not_swallow_active_call(self): + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_parse_inactive_rehearsal_alone_is_prose(self): + assert ( + parse_tool_calls_from_text('foo[ARGS]{"a":1}', enabled_tool_names = {"web_search"}) == [] + ) + + def test_streaming_strip_keeps_inactive_rehearsal(self): + raw = 'answer foo[ARGS]{"x":1} tail' + assert strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) == raw + + def test_streaming_strip_removes_active_rehearsal(self): + raw = 'answer web_search[ARGS]{"q":1} tail' + out = strip_tool_markup_streaming(raw, enabled_tool_names = {"web_search"}) + assert "web_search[ARGS]" not in out + assert out == "answer tail" + + def test_final_strip_keeps_inactive_rehearsal(self): + text = 'foo[ARGS]{"x":1} is just syntax.' + assert strip_tool_markup(text, final = True, enabled_tool_names = {"web_search"}) == text + + def test_gate_none_preserves_legacy_strip_and_parse(self): + text = 'foo[ARGS]{"x":1} tail' + assert self._names(parse_tool_calls_from_text(text)) == ["foo"] + assert strip_tool_markup_streaming(text) == " tail" + + def test_drain_truncated_enabled_name_json_preserved_when_auto_heal_disabled(): # F3: with Auto-Heal OFF, a truncated ENABLED-name bare-JSON fragment that did # not parse must stay visible (disabled-Auto-Heal contract: malformed markup is diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 7f47140b8d..c6da1e90e7 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -266,6 +266,133 @@ class TestHealingPathUnaffected: assert healed[span[0] : span[1]] == "dogs" +class TestEnabledToolNameGate: + """``enabled_tool_names`` disambiguates the ambiguous bare-rehearsal + ``NAME[ARGS]{json}`` form (#5704): NAME is a call only when it is an active tool, + otherwise it is prose. ``None`` (the default) keeps the legacy unrestricted parse + so existing callers are unaffected.""" + + def _names(self, calls): + return [c["function"]["name"] for c in calls] + + def test_inactive_rehearsal_before_active_call_does_not_swallow_it(self): + # P1: an inactive ``foo[ARGS]{...}`` before a real call must not consume the real call. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "cats"} + + def test_inactive_rehearsal_alone_is_not_a_call(self): + text = 'foo[ARGS]{"a":1}' + assert parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) == [] + + def test_active_rehearsal_is_still_parsed(self): + text = 'web_search[ARGS]{"query":"cats"}' + calls = parse_tool_calls_from_text(text, enabled_tool_names = {"web_search"}) + assert self._names(calls) == ["web_search"] + + def test_unrestricted_gate_none_preserves_legacy_behavior(self): + # Without a gate every ``NAME[ARGS]{...}`` is parsed, as before the gate landed. + text = 'foo[ARGS]{"a":1} web_search[ARGS]{"query":"cats"}' + assert self._names(parse_tool_calls_from_text(text)) == ["foo", "web_search"] + assert self._names(parse_tool_calls_from_text(text, enabled_tool_names = None)) == [ + "foo", + "web_search", + ] + + +class TestBracketCallSpans: + """with_spans tiling for Mistral bracket calls: promoted markup strips + exactly once, filtered calls' bytes stay visible, closers strip too.""" + + def test_mixed_array_filtered_first_keeps_its_bytes_only(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"bad","arguments":{"x":1}},' + '{"name":"lookup","arguments":{"q":"cats"}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + kinds = [k for k, _v in events] + assert kinds == ["text", "tool_call"] + text = events[0][1] + assert '"bad"' in text + # The promoted call's markup must not survive in the text event. + assert '"lookup"' not in text + + def test_mixed_array_filtered_second_stays_visible(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"cats"}},' + '{"name":"bad","arguments":{"x":1}}]' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert events[0][0] == "tool_call" + trailing = "".join(v for k, v in events if k == "text") + assert '"bad"' in trailing + + def test_v11_closer_inside_span(self): + from core.tool_healing import parse_tool_calls_from_text as parse_with_spans + + text = '[TOOL_CALLS]web_search[ARGS]{"query":"cats"}[/TOOL_CALLS] after' + calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True) + (call,) = calls + assert call["function"]["name"] == "web_search" + (span,) = spans + assert text[span[0] : span[1]].endswith("[/TOOL_CALLS]") + assert text[span[1] :] == " after" + + def test_fully_promoted_array_strips_whole_region(self): + from core.inference.passthrough_healing import heal_openai_message_events + + tools = [{"type": "function", "function": {"name": "lookup", "parameters": {}}}] + content = ( + '[TOOL_CALLS][{"name":"lookup","arguments":{"q":"a"}},' + '{"name":"lookup","arguments":{"q":"b"}}] after' + ) + events = heal_openai_message_events( + {"role": "assistant", "content": content}, {"lookup"}, tools + ) + assert [k for k, _v in events] == ["tool_call", "tool_call", "text"] + assert events[2][1] == " after" + + +class TestMistralArrayHealing: + """Draining the whole [TOOL_CALLS] array for the shapes the repo's own + Mistral/Ollama templates emit.""" + + def test_comma_less_multi_call_array_parses_all_calls(self): + # ollama_template_mappers.py renders multi-call turns as [{...}{...}] with no + # comma separator; a single json.loads of the body rejects it and dropped every + # call. The element-by-element decode must recover all of them. + text = '[TOOL_CALLS] [{"name":"a","arguments":{"x":1}}{"name":"b","arguments":{"y":2}}]' + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["a", "b"] + assert json.loads(calls[0]["function"]["arguments"]) == {"x": 1} + assert json.loads(calls[1]["function"]["arguments"]) == {"y": 2} + + def test_comma_separated_and_single_arrays_still_parse(self): + both = parse_tool_calls_from_text( + '[TOOL_CALLS] [{"name":"a","arguments":{}},{"name":"b","arguments":{}}]' + ) + assert [c["function"]["name"] for c in both] == ["a", "b"] + one = parse_tool_calls_from_text('[TOOL_CALLS] [{"name":"a","arguments":{}}]') + assert [c["function"]["name"] for c in one] == ["a"] + + def test_mistral_array_null_arguments_normalized_to_empty_object(self): + # ``"arguments": null`` is a no-arg call; it must become {} (as the + # path does), not the string "null" that auto-heal turns into {"query":"null"}. + calls = parse_tool_calls_from_text('[TOOL_CALLS][{"name":"get_time","arguments":null}]') + assert calls[0]["function"]["arguments"] == "{}" + + class TestGlmStrict: def test_closed_glm_call_is_accepted(self): text = ( @@ -740,22 +867,26 @@ class TestMistralOuterOverXmlLiteral: class TestHealerSignalAlignment: - """The healer buffers only promotable formats; Mistral/Llama text calls stream through.""" + """The healer buffers only formats its shared parser can promote. Mistral's + ``[TOOL_CALLS]`` is promotable (rescued), so it is a heal signal; the loop-only + text-call markers (Llama ``<|python_tag|>``, bare ``[ARGS]``) are not, so they + stream through instead of stalling as prose that never yields a call.""" def test_heal_signals_subset_of_promotable_formats(self): from core.inference.passthrough_healing import _HEAL_SIGNALS - assert set(_HEAL_SIGNALS) == {"", "<|tool_call>", "", "<|tool_call>", " is not a healer-promotable format, so it streams through as text. + events = list(healer.feed('<|python_tag|>web_search.call(query="cats")')) text_out = "".join(v for k, v in events if k == "text") - assert "[TOOL_CALLS]" in text_out # streamed through, not buffered + assert "<|python_tag|>" in text_out # streamed through, not buffered assert not list(healer.finalize()) or all(k == "text" for k, _v in healer.finalize()) diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index d50c27130f..f7792a2a71 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -52,6 +52,11 @@ _ns = { } exec(f"_TOOL_XML_RE = _re.compile({_m.group(1)})", _ns) _TOOL_XML_RE = _ns["_TOOL_XML_RE"] +# The display helper uses the closed-only variant before the last think block; keep it in scope. +_mc = _re.search(r"_TOOL_XML_CLOSED_RE = _re\.compile\((.*?)\n\)", _src, _re.DOTALL) +assert _mc, "could not extract _TOOL_XML_CLOSED_RE source" +exec(f"_TOOL_XML_CLOSED_RE = _re.compile({_mc.group(1)})", _ns) +_TOOL_XML_CLOSED_RE = _ns["_TOOL_XML_CLOSED_RE"] # Signatures may span multiple lines and now carry the enabled_tool_names gate; match # the whole (possibly multi-line) signature up to ``-> str:`` then the indented body. @@ -66,16 +71,19 @@ assert "_strip_mistral_closed_calls" in _xml_helper.group( exec(_xml_helper.group(0), _ns) _strip_tool_xml = _ns["_strip_tool_xml"] +# Extract the gate helper and display strip up to the next top-level ``logger =``. _helper = _re.search( - r"def _strip_tool_xml_for_display\((?:.|\n)*?\) -> str:\n(?: .+\n)+", + r"def _display_tool_name_gate\(.*?(?=\nlogger = get_logger)", _src, + _re.DOTALL, ) -assert _helper, "could not extract _strip_tool_xml_for_display source" -# After the V1 fix the display helper delegates to _strip_tool_xml; confirm the -# extracted body actually reached that call rather than truncating early. +assert _helper, "could not extract display strip helper source" +# The extracted block spans _display_tool_name_gate through _strip_tool_xml (defined before +# ``logger =``); confirm the shared _strip_tool_xml delegate is present. assert "_strip_tool_xml(" in _helper.group(0), "display helper no longer delegates" exec(_helper.group(0), _ns) _strip_tool_xml_for_display = _ns["_strip_tool_xml_for_display"] +_display_tool_name_gate = _ns["_display_tool_name_gate"] _gate_src = _re.search( r"def _gemma_strip_gate\((?:.|\n)*?\) -> set:\n(?: .+\n)+", @@ -95,6 +103,56 @@ def test_route_display_strip_respects_disabled_auto_heal_contract(): assert "" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) +def test_route_display_strip_preserves_rehearsal_inside_think(): + # A rehearsed bracket call inside think is reasoning: the block is preserved while a real + # call outside it still strips. + text = 'plan: search[ARGS]{"q":"x"} answer [TOOL_CALLS]web_search{"q":"y"} tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert 'plan: search[ARGS]{"q":"x"}' in out + assert "[TOOL_CALLS]web_search" not in out + assert "answer" in out and "tail" in out + + +def test_route_display_strip_keeps_bare_args_before_think_block(): + # A bare ``foo[ARGS]`` before a think block is prose: EOS-anchored tail arms run only on + # the last segment (earlier segments use the closed-only regex). + text = "Please pass foo[ARGS] pause to the template." + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) == text + + +def test_route_display_strip_removes_complete_call_before_think_block(): + # A complete bracket call before a think block still strips (balanced scan runs on every segment). + text = 'before search[ARGS]{"q":"x"} pause after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "search[ARGS]" not in out + assert "pause" in out + assert "before" in out and "after" in out + + +def test_route_display_strip_removes_closed_xml_before_think_block(): + # A closed before a think block is removed in the non-last segment. + text = 'pre {"name":"x"} p tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out + assert "p" in out + assert "pre" in out and "tail" in out + + +def test_all_route_cleanup_sites_use_protected_display_helper(): + # Every route cleanup site must use _strip_tool_xml_for_display (think-preserving, + # balanced); raw _TOOL_XML_RE.sub corrupted think rehearsal and trailing prose. The only + # legitimate raw sub lives inside the helper itself. + raw_sub_lines = [ + (i, line) + for i, line in enumerate(_src.splitlines(), 1) + if "_TOOL_XML_RE.sub(" in line and not line.lstrip().startswith("#") + ] + assert len(raw_sub_lines) == 1, ( + "raw _TOOL_XML_RE.sub must appear only inside _strip_tool_xml_for_display; " + f"found extra call sites: {raw_sub_lines!r}" + ) + + def test_route_display_strip_removes_mistral_tool_calls_with_nested_json(): # _TOOL_XML_RE has no [TOOL_CALLS] arm, so the helper delegates to _strip_tool_xml for the Mistral # balanced-brace strip (a non-greedy \{.*?\} would truncate nested JSON). @@ -234,6 +292,32 @@ def test_strips_tail_only_parameter_orphan_no_trailing_ws(): assert "Final answer." in cleaned +def test_strips_complete_bracket_tag_keeps_trailing_prose(): + # A complete Mistral call strips only its balanced JSON, leaving following prose intact. + cleaned = _TOOL_XML_RE.sub("", '[TOOL_CALLS]web_search{"q":"x"} and then prose') + assert "[TOOL_CALLS]" not in cleaned + assert "and then prose" in cleaned + + +def test_strips_unclosed_bracket_tail(): + # Close brace lost to EOS: the truncated tail strips to the end instead of leaking. + cleaned = _TOOL_XML_RE.sub("", 'here [TOOL_CALLS]web_search{"query":"weather"') + assert "[TOOL_CALLS]" not in cleaned + assert cleaned.strip() == "here" + + +def test_strips_unclosed_rehearsal_tail(): + cleaned = _TOOL_XML_RE.sub("", 'text python[ARGS]{"code":"print(1)"') + assert "[ARGS]" not in cleaned + assert cleaned.strip() == "text" + + +def test_strips_hyphenated_mcp_bracket_name(): + cleaned = _TOOL_XML_RE.sub("", 'x [TOOL_CALLS]mcp__srv__list-issues{"q":"x"}') + assert "list-issues" not in cleaned + assert cleaned.strip() == "x" + + def test_preserves_mid_string_parameter_in_code_sample(): # Tail-anchor on `` so doc/example prose survives. text = ( @@ -362,6 +446,238 @@ def test_no_catastrophic_backtracking_on_orphan_opening_spam(): assert "" not in cleaned +# ── Two-level-nested bracket JSON (balanced-scan strip) ────────── + + +def test_route_strip_two_level_nested_bracket_keeps_trailing_prose(): + # Two-level-nested args must be removed whole so the trailing prose survives. + text = 'before [TOOL_CALLS]search{"f":{"g":{"h":1}}} after' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "before after" + assert "[TOOL_CALLS]" not in cleaned + + +def test_route_strip_two_level_nested_rehearsal_keeps_trailing_prose(): + text = 'note python[ARGS]{"a":{"b":{"c":1}}} done' + cleaned = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert cleaned == "note done" + assert "[ARGS]" not in cleaned + + +def test_route_strip_removes_call_with_literal_think_in_argument(): + # A literal inside a call argument strips with the call, not as reasoning. + text = ( + '{"name":"write","arguments":' + '{"text":"compare and tags"}}' + ) + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "" not in out and '"name"' not in out + + +def test_route_strip_removes_truncated_mistral_array(): + # A canonical array truncated by EOS is stripped by the route fallback like other orphans. + text = 'before [TOOL_CALLS] [{"name":"a","arguments":{"x":1}}' # missing ] + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "{" not in out + assert "before" in out + + +def test_route_strip_keeps_prose_mentioning_args_marker(): + # ``foo[ARGS] in a sentence`` is prose; the rehearsal arm must not truncate the line. + text = "Please pass foo[ARGS] to the template and continue reading." + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out == text + + +def test_route_strip_handles_mistral_v11_call_id_args_shape(): + # v11 [CALL_ID]/[ARGS] shape (Mistral Small 3.2) must strip whole. + text = 'before [TOOL_CALLS]web_search[CALL_ID]abc123[ARGS]{"q":"x"} after' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[TOOL_CALLS]" not in out and "[CALL_ID]" not in out and "[ARGS]" not in out + assert "before" in out and "after" in out + + +# ── Mistral [/TOOL_CALLS] closer + literal inside a call ─────────────── + +from core.tool_healing import strip_tool_call_markup as _strip_tool_call_markup + + +def test_core_strip_removes_orphan_tool_calls_closer_array_form(): + # The bare v11 [/TOOL_CALLS] closer left by the balanced scan must not leak as content. + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + assert _strip_tool_call_markup(text, final = True) == "" + + +def test_core_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + assert _strip_tool_call_markup(text, final = True) == "tail" + + +def test_core_strip_removes_call_with_literal_think_in_argument(): + # An unclosed literal inside call arguments strips with the call (argument data). + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert _strip_tool_call_markup(text, final = True) == "before after" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_array_form(): + text = '[TOOL_CALLS] [{"name":"x","arguments":{}}][/TOOL_CALLS]' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert out.strip() == "" + + +def test_route_display_strip_removes_orphan_tool_calls_closer_named_form_keeps_tail(): + text = '[TOOL_CALLS]web_search{"q":"x"}[/TOOL_CALLS] tail' + out = _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + assert "[/TOOL_CALLS]" not in out + assert out.strip() == "tail" + + +def test_incomplete_xml_call_with_literal_think_in_arg_is_stripped(): + # An incomplete holding a literal strips to EOS, not as a reasoning + # block (the unclosed tail _tool_call_markup_spans previously missed). + from core.tool_healing import parse_tool_calls_from_text as _parse + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'before {"name":"write","arguments":{"text":"literal marker"}} after' + assert [c["function"]["name"] for c in _parse(text)] == ["write"] + assert _strip(text, final = True) == "before" + + # A real reasoning block with no tool call is still preserved verbatim. + assert ( + _strip("answer real done", final = True) == "answer real done" + ) + + # A complete call followed by a real reasoning block: call stripped, block kept. + mixed = '{"name":"a","arguments":{}} mid r end' + assert _strip(mixed, final = True) == "mid r end" + + +# ── enabled-tool gate for the ambiguous bare-rehearsal strip (#5704) ── + + +def test_display_tool_name_gate_returns_active_names_or_none(): + # Empty / no tools -> None (unrestricted; keep the legacy strip-all behavior). + assert _display_tool_name_gate([]) is None + assert _display_tool_name_gate(None) is None + # OpenAI-shaped tool dicts -> set of function names, malformed entries dropped. + tools = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "run_python"}}, + {"type": "function"}, # no name + {"nope": 1}, # no function + ] + assert _display_tool_name_gate(tools) == {"web_search", "run_python"} + + +def test_route_display_strip_keeps_inactive_rehearsal_when_gated(): + # P1 #5704: an inactive ``foo[ARGS]{...}`` is prose; the gated strip leaves the sentence intact. + gate = {"web_search"} + text = 'foo[ARGS]{"x":1} is just syntax.' + assert ( + _strip_tool_xml_for_display(text, auto_heal_tool_calls = True, enabled_tool_names = gate) + == text + ) + # A bare marker with no JSON body is likewise prose when inactive. + assert ( + _strip_tool_xml_for_display( + "use foo[ARGS] here", auto_heal_tool_calls = True, enabled_tool_names = gate + ) + == "use foo[ARGS] here" + ) + + +def test_route_display_strip_removes_active_rehearsal_when_gated(): + # Mirror case: an active tool name is a real rehearsal and still strips. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + 'web_search[ARGS]{"query":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "web_search[ARGS]" not in out + assert out.strip() == "done" + + +def test_route_display_strip_ungated_strips_all_rehearsal_unchanged(): + # Backwards-compat: with no gate (None) the bare rehearsal strips as before. + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip_tool_xml_for_display(text, auto_heal_tool_calls = True).strip() == "is just syntax." + assert ( + _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = None + ).strip() + == "is just syntax." + ) + + +def test_route_display_strip_control_token_stripped_regardless_of_gate(): + # [TOOL_CALLS] is a control token: stripped even when its NAME is not in the gate. + gate = {"web_search"} + out = _strip_tool_xml_for_display( + '[TOOL_CALLS]foo[ARGS]{"x":1} keep', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + assert "[TOOL_CALLS]" not in out and "foo[ARGS]" not in out + assert out.strip() == "keep" + + +def test_core_strip_gates_bare_rehearsal_on_enabled_tools(): + # P1 (#5704): the shared strip gate mirrors the parse gate -- inactive names are prose + # and preserved, active names strip, ``None`` keeps legacy strip-all. + from core.tool_healing import strip_tool_call_markup as _strip + + text = 'foo[ARGS]{"x":1} is just syntax.' + assert _strip(text, final = True, enabled_tool_names = {"web_search"}) == text + assert ( + _strip('web_search[ARGS]{"q":1} done', final = True, enabled_tool_names = {"web_search"}) + == "done" + ) + assert _strip(text, final = True).strip() == "is just syntax." + assert _strip(text, final = True, enabled_tool_names = None).strip() == "is just syntax." + + +def test_route_display_strip_gate_preserves_inactive_history_rehearsal(): + # The GGUF history sanitiser passes the gate, so a documented inactive shape survives in + # the replayed prompt context. + gate = _display_tool_name_gate([{"function": {"name": "web_search"}}]) + text = 'To call it write foo[ARGS]{"x":1} in your reply.' + assert 'foo[ARGS]{"x":1}' in _strip_tool_xml_for_display( + text, auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # An ACTIVE name is still stripped as a real rehearsed call. + assert "web_search[ARGS]" not in _strip_tool_xml_for_display( + 'Result web_search[ARGS]{"q":"x"} done', auto_heal_tool_calls = True, enabled_tool_names = gate + ) + # No gate (legacy) strips every NAME[ARGS]{...}. + assert "foo[ARGS]" not in _strip_tool_xml_for_display(text, auto_heal_tool_calls = True) + + +def test_gguf_history_sanitizer_forwards_enabled_tool_names_gate(): + # Wiring guard: the GGUF history strip must forward the display gate like the live strip. + block = _re.search( + r"Strip stale tool-call XML from conversation history.*?\.strip\(\)", + _src, + _re.DOTALL, + ) + assert block, "could not locate GGUF history sanitizer block" + assert "enabled_tool_names" in block.group( + 0 + ), "GGUF history sanitizer must pass enabled_tool_names to _strip_tool_xml_for_display" + + +def test_route_history_and_passthrough_forward_the_display_gate(): + # The safetensors/Anthropic history sanitisers and the Anthropic non-stream passthrough + # must forward the gate so inactive examples survive in replayed prompt / final text. + blocks = { + "safetensors history": r"Strip stale tool-call XML from prior assistant turns.*?\.strip\(\)", + "anthropic history": r"Strip stale tool-call XML via the protected display helper.*?\.strip\(\)", + "anthropic passthrough": r"gated on the declared tools so an\n.*?\.strip\(\)", + } + for label, pat in blocks.items(): + m = _re.search(pat, _src, _re.DOTALL) + assert m, f"could not locate {label} strip block" + assert "enabled_tool_names" in m.group( + 0 + ), f"{label} must forward enabled_tool_names to _strip_tool_xml_for_display" + + # ── DeepSeek opener variants + bare Kimi (parse/strip symmetry) ──