From 09ca2a1777826065410d94f02c8a5927e20b96b1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 14:42:19 +0000 Subject: [PATCH 1/3] Studio: retry agentic loop on validation failure Adds a pre-dispatch validation pass inside both agentic tool loops (generate_chat_completion_with_tools in core/inference/llama_cpp.py and run_safetensors_tool_loop in core/inference/safetensors_agentic.py). The pass catches two failure modes between parser and dispatch: * Unknown tool name (not in the request's tools[] array). * Arguments that cannot decode to a JSON object when auto_heal is off. On a caught call the loop appends a corrective tool-result message tied to the hallucinated tool_call_id (not a fabricated id, so the OpenAI chat template stays valid) and re-enters the model. When the call has no usable id we fall back to a user-role correction since tool-role messages require a matching prior call id. The retry pass is bounded by max_validation_retries (default 2, new ChatCompletionRequest field, threaded through the route layer). On budget exhaustion the call falls through to the existing per-tool error path so today's behavior is preserved. When auto_heal_tool_calls is on the heal path still runs in the dispatch loop unchanged; F3 only catches the strict-shape failures the coercer cannot fix. --- studio/backend/core/inference/inference.py | 2 + studio/backend/core/inference/llama_cpp.py | 96 ++++++++ .../core/inference/safetensors_agentic.py | 86 +++++++ studio/backend/models/inference.py | 5 + studio/backend/routes/inference.py | 10 + .../tests/test_validation_retry_loop.py | 227 ++++++++++++++++++ 6 files changed, 426 insertions(+) create mode 100644 studio/backend/tests/test_validation_retry_loop.py diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index e1620f5ca3..f7a8839cb3 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -858,6 +858,7 @@ class InferenceBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + max_validation_retries: int = 2, ): """Run an agentic tool loop on top of ``generate_chat_response``. @@ -905,6 +906,7 @@ class InferenceBackend: max_tool_iterations = max_tool_iterations, tool_call_timeout = tool_call_timeout, session_id = session_id, + max_validation_retries = max_validation_retries, ) def generate_chat_response( diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..762eeb22d5 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -4428,6 +4428,7 @@ class LlamaCppBackend: auto_heal_tool_calls: bool = True, tool_call_timeout: int = 300, session_id: Optional[str] = None, + max_validation_retries: int = 2, ) -> Generator[dict, None, None]: """ Agentic loop: let the model call tools, execute them, and continue. @@ -4476,6 +4477,21 @@ class LlamaCppBackend: # Pattern is compiled once at module level (_INTENT_SIGNAL). _reprompt_count = 0 + # ── Validation retry budget (F3) ───────────────────── + # Catches tool calls to unknown names or with malformed args + # (e.g. arguments that don't decode to a JSON object). On a + # caught call, append a corrective tool-result message tied to + # the hallucinated call id and re-enter the model. Bounded by + # max_validation_retries to avoid retry storms; on exhaustion + # the call falls through to the existing error-handling path + # so behavior matches today. + _validation_retries = 0 + _allowed_tool_names = { + (_t.get("function") or {}).get("name") + for _t in (tools or []) + if (_t.get("function") or {}).get("name") + } + # Reserve extra iterations for re-prompts so they don't # consume the caller's tool-call budget. Only add the # extra slot when tool iterations are actually allowed. @@ -4999,6 +5015,86 @@ class LlamaCppBackend: assistant_msg["tool_calls"] = tool_calls conversation.append(assistant_msg) + # ── F3: Validate tool calls before dispatch ── + # On unknown name or malformed args, append a corrective + # tool-result tied to the hallucinated call id and + # re-enter the model. Skips when the budget is spent so + # the existing error path still runs. + _validation_problem = None + if ( + tool_calls + and _allowed_tool_names + and _validation_retries < max_validation_retries + ): + for _vtc in tool_calls: + _vfn = _vtc.get("function", {}) or {} + _vname = _vfn.get("name", "") + if _vname not in _allowed_tool_names: + _validation_problem = ("unknown_tool", _vtc, _vname) + break + _vraw = _vfn.get("arguments", "") + if isinstance(_vraw, str): + try: + _vdec = json.loads(_vraw) if _vraw else {} + except (json.JSONDecodeError, ValueError): + _vdec = None + if not isinstance(_vdec, dict): + _validation_problem = ( + "malformed_args", _vtc, _vname, + ) + break + elif not isinstance(_vraw, dict): + _validation_problem = ( + "malformed_args", _vtc, _vname, + ) + break + + if _validation_problem is not None: + _kind, _vtc, _vname = _validation_problem + _validation_retries += 1 + _vcall_id = _vtc.get("id") + if _kind == "unknown_tool" and _vcall_id: + _allowed_list = ", ".join(sorted(_allowed_tool_names)) + conversation.append({ + "role": "tool", + "tool_call_id": _vcall_id, + "name": _vname, + "content": ( + f"Error: tool '{_vname}' is not available. " + f"Available tools: {_allowed_list}." + ), + }) + elif _kind == "malformed_args" and _vcall_id: + conversation.append({ + "role": "tool", + "tool_call_id": _vcall_id, + "name": _vname, + "content": ( + f"Error: arguments to '{_vname}' could not " + "be parsed as a JSON object. Call " + f"'{_vname}' again with valid JSON object " + "arguments." + ), + }) + else: + # No usable call id: user-role correction since + # OpenAI tool messages require a matching id. + conversation.append({ + "role": "user", + "content": ( + "Your last tool call was malformed " + "(missing id or function). Please " + "re-issue it as a valid OpenAI " + "function call." + ), + }) + logger.info( + "validation_retry kind=%s retries=%d/%d", + _kind, _validation_retries, max_validation_retries, + ) + yield {"type": "status", "text": ""} + continue + for tc in tool_calls or []: func = tc.get("function", {}) tool_name = func.get("name", "") diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 73bb3d090a..b7fbe680e5 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -105,6 +105,7 @@ def run_safetensors_tool_loop( max_tool_iterations: int = 25, tool_call_timeout: int = 300, session_id: Optional[str] = None, + max_validation_retries: int = 2, ) -> Generator[dict, None, None]: """Drive an agentic tool loop on top of a cumulative-text generator. @@ -142,6 +143,10 @@ def run_safetensors_tool_loop( if (tool.get("function") or {}).get("name") } next_call_id = 0 + # F3: bound how many times the loop will rewrite a corrective + # tool-result and re-enter the model. Past the cap, fall through to + # the existing per-tool error path so we don't loop forever. + validation_retries = 0 if max_tool_iterations <= 0: # 0 = disabled (same contract as the GGUF loop). @@ -299,6 +304,87 @@ def run_safetensors_tool_loop( next_call_id += len(tool_calls) conversation.append(assistant_msg) + # F3: pre-dispatch validation. On unknown name or non-decodable + # arguments, append a corrective tool-result tied to the call + # id and re-enter the model. ``auto_heal_tool_calls`` semantics + # are preserved: when heal is on, a string-args coercion path + # still runs in the dispatch loop below; the validation here + # catches the strict-shape failures the coercer cannot heal. + validation_problem = None + if ( + tool_calls + and allowed_tool_names + and validation_retries < max_validation_retries + ): + for v_tc in tool_calls: + v_fn = v_tc.get("function", {}) or {} + v_name = v_fn.get("name", "") or "" + if v_name not in allowed_tool_names: + validation_problem = ("unknown_tool", v_tc, v_name) + break + v_raw = v_fn.get("arguments", "") + if isinstance(v_raw, str): + if v_raw == "": + v_decoded = {} + else: + try: + v_decoded = json.loads(v_raw) + except (json.JSONDecodeError, ValueError): + v_decoded = None + if ( + not isinstance(v_decoded, dict) + and not auto_heal_tool_calls + ): + validation_problem = ("malformed_args", v_tc, v_name) + break + elif not isinstance(v_raw, dict): + if not auto_heal_tool_calls: + validation_problem = ("malformed_args", v_tc, v_name) + break + + if validation_problem is not None: + v_kind, v_tc, v_name = validation_problem + validation_retries += 1 + v_call_id = v_tc.get("id") + if v_kind == "unknown_tool" and v_call_id: + allowed_list = ", ".join(sorted(allowed_tool_names)) + conversation.append({ + "role": "tool", + "tool_call_id": v_call_id, + "name": v_name, + "content": ( + f"Error: tool '{v_name}' is not available. " + f"Available tools: {allowed_list}." + ), + }) + elif v_kind == "malformed_args" and v_call_id: + conversation.append({ + "role": "tool", + "tool_call_id": v_call_id, + "name": v_name, + "content": ( + f"Error: arguments to '{v_name}' could not be " + "parsed as a JSON object. Call " + f"'{v_name}' again with valid JSON object " + "arguments." + ), + }) + else: + conversation.append({ + "role": "user", + "content": ( + "Your last tool call was malformed (missing id " + "or function). Please re-issue it as a valid " + "OpenAI function call." + ), + }) + logger.info( + "validation_retry kind=%s retries=%d/%d", + v_kind, validation_retries, max_validation_retries, + ) + yield {"type": "status", "text": ""} + continue + for tc in tool_calls or []: func = tc.get("function", {}) or {} tool_name = func.get("name", "") or "" diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index b5626951c4..e0f4f3cf40 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -677,6 +677,11 @@ class ChatCompletionRequest(BaseModel): ge = 0, description = "[x-unsloth] Maximum number of tool call iterations per message (0 = disabled, 9999 = unlimited).", ) + max_validation_retries: Optional[int] = Field( + 2, + ge = 0, + description = "[x-unsloth] Maximum corrective retries when the model emits a tool call to an unknown name or with malformed arguments. 0 disables the retry pass.", + ) tool_call_timeout: Optional[int] = Field( 300, ge = 1, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 02270ab405..8d85e80f46 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2483,6 +2483,11 @@ async def openai_chat_completions( if payload.tool_call_timeout is not None else 300, session_id = payload.session_id, + max_validation_retries = ( + payload.max_validation_retries + if payload.max_validation_retries is not None + else 2 + ), ) _tool_sentinel = object() @@ -2972,6 +2977,11 @@ async def openai_chat_completions( else 300, session_id = payload.session_id, use_adapter = payload.use_adapter, + max_validation_retries = ( + payload.max_validation_retries + if payload.max_validation_retries is not None + else 2 + ), ) _sf_tool_sentinel = object() diff --git a/studio/backend/tests/test_validation_retry_loop.py b/studio/backend/tests/test_validation_retry_loop.py new file mode 100644 index 0000000000..719683e158 --- /dev/null +++ b/studio/backend/tests/test_validation_retry_loop.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the F3 validation-retry pass in ``run_safetensors_tool_loop``. + +The pass detects two failure modes between parse and dispatch: + +* The model emitted a tool call to a name not in the request's tools. +* The model emitted arguments that are not a JSON object (auto-heal off). + +Either failure triggers a corrective tool-result message tied to the +hallucinated ``tool_call_id`` and re-enters the model loop. The pass is +bounded by ``max_validation_retries`` so the loop cannot retry forever. +""" + +from core.inference.safetensors_agentic import run_safetensors_tool_loop + + +class FakeExecuteTool: + def __init__(self, results = None): + self.results = list(results or []) + self.calls = [] + + def __call__(self, name, arguments, *, cancel_event = None, timeout = None, session_id = None): + self.calls.append((name, arguments)) + return self.results.pop(0) if self.results else "OK" + + +def _multi_turn(turns): + """Yield cumulative-text generators, one per turn.""" + turn_iter = iter(turns) + + def _gen(_messages): + try: + chunks = next(turn_iter) + except StopIteration: + return + acc = "" + for c in chunks: + acc += c + yield acc + + return _gen + + +def _collect(loop, cap = 200): + out = [] + for ev in loop: + out.append(ev) + if len(out) >= cap: + break + return out + + +REAL_TOOLS = [ + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "function": {"name": "python"}}, +] + + +class TestUnknownTool: + def test_unknown_tool_triggers_retry(self): + # Turn 1: hallucinated tool name "missing_tool". + # Turn 2: valid call after the corrective nudge. + single_turn = _multi_turn([ + ['{"name":"missing_tool","arguments":{}}'], + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["Done."], + ]) + exec_fn = FakeExecuteTool(["result"]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + ) + ) + # The bad call was NOT executed; the good follow-up call ran. + assert exec_fn.calls == [("web_search", {"query": "x"})] + # Final content carries the model's answer. + contents = [e for e in events if e["type"] == "content"] + assert any("Done" in e.get("text", "") for e in contents) + + def test_retry_budget_exhausted_falls_through(self): + # Three unknown calls in a row. With max_validation_retries=2 + # the first two are retried; the third falls through to the + # existing per-tool error path (which never executes a real + # tool but emits an error result to the model). + single_turn = _multi_turn([ + ['{"name":"missing_a","arguments":{}}'], + ['{"name":"missing_b","arguments":{}}'], + ['{"name":"missing_c","arguments":{}}'], + ["Sorry, I cannot proceed."], + ]) + exec_fn = FakeExecuteTool([]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + max_validation_retries = 2, + ) + ) + # execute_tool is never called for an unknown name in either + # the F3 retry arm or the existing per-tool error arm. + assert exec_fn.calls == [] + # The loop ultimately exits cleanly; we expect at least one + # tool_end event from the existing error path on the third turn. + kinds = [e.get("type") for e in events] + assert "status" in kinds + + def test_retries_disabled_means_no_retry(self): + # With max_validation_retries=0, the F3 arm never engages and + # behavior matches the pre-F3 path: the existing per-tool + # error message is emitted but no corrective re-entry happens. + single_turn = _multi_turn([ + ['{"name":"missing","arguments":{}}'], + ["bye"], + ]) + exec_fn = FakeExecuteTool([]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + max_validation_retries = 0, + ) + ) + # No retry round; the unknown call goes straight to the + # existing error path. + assert exec_fn.calls == [] + tool_ends = [e for e in events if e["type"] == "tool_end"] + # Existing per-tool error path emits a tool_end with an Error + # result so the model sees the failure. + assert any( + "not enabled" in str(e.get("result", "")) + for e in tool_ends + ) + + +class TestMalformedArgs: + def test_malformed_args_bypassed_when_heal_on(self): + # With auto_heal_tool_calls=True (the default), string args are + # healed to {"query": "..."} for web_search. F3 leaves heal + # behavior intact and only catches strictly-impossible shapes. + single_turn = _multi_turn([ + ['{"name":"web_search","arguments":"some text"}'], + ["all done"], + ]) + exec_fn = FakeExecuteTool(["ok"]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + auto_heal_tool_calls = True, + ) + ) + # The heal path coerces "some text" to {"query": "some text"} + # and execute_tool runs. + assert len(exec_fn.calls) == 1 + assert exec_fn.calls[0][0] == "web_search" + + def test_malformed_args_caught_when_heal_off(self): + # With auto_heal off, a non-dict arguments value is a hard + # malformed-args failure and the F3 arm catches it. + single_turn = _multi_turn([ + ['{"name":"web_search","arguments":"not a dict"}'], + ['{"name":"web_search","arguments":{"query":"sf"}}'], + ["sunny"], + ]) + exec_fn = FakeExecuteTool(["sunny in sf"]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + auto_heal_tool_calls = False, + ) + ) + # The first call was caught by F3, not executed. The second + # (well-formed) call ran. + assert exec_fn.calls == [("web_search", {"query": "sf"})] + + +class TestNoOpCases: + def test_no_tools_no_validation(self): + # Empty tools list means allowed_tool_names is empty, so the + # F3 arm never engages. + single_turn = _multi_turn([ + ['{"name":"anything","arguments":{}}'], + ["done"], + ]) + exec_fn = FakeExecuteTool([]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = [], + execute_tool = exec_fn, + ) + ) + # No allowlist gate, no F3 retry; the bad call goes through + # the dispatch loop and execute_tool runs. + assert exec_fn.calls == [("anything", {})] + + def test_valid_call_no_retry_overhead(self): + # A clean valid call does not trigger F3 at all. + single_turn = _multi_turn([ + ['{"name":"web_search","arguments":{"query":"x"}}'], + ["final"], + ]) + exec_fn = FakeExecuteTool(["res"]) + events = _collect( + run_safetensors_tool_loop( + single_turn = single_turn, + messages = [{"role": "user", "content": "hi"}], + tools = REAL_TOOLS, + execute_tool = exec_fn, + ) + ) + assert exec_fn.calls == [("web_search", {"query": "x"})] From 699da0decdd10a982942fa835aefc894b54e32ba Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 May 2026 14:43:00 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/inference/llama_cpp.py | 76 +++++++++------ .../core/inference/safetensors_agentic.py | 71 +++++++------- .../tests/test_validation_retry_loop.py | 97 ++++++++++++------- 3 files changed, 142 insertions(+), 102 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 762eeb22d5..da25d91678 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -5040,12 +5040,16 @@ class LlamaCppBackend: _vdec = None if not isinstance(_vdec, dict): _validation_problem = ( - "malformed_args", _vtc, _vname, + "malformed_args", + _vtc, + _vname, ) break elif not isinstance(_vraw, dict): _validation_problem = ( - "malformed_args", _vtc, _vname, + "malformed_args", + _vtc, + _vname, ) break @@ -5055,42 +5059,50 @@ class LlamaCppBackend: _vcall_id = _vtc.get("id") if _kind == "unknown_tool" and _vcall_id: _allowed_list = ", ".join(sorted(_allowed_tool_names)) - conversation.append({ - "role": "tool", - "tool_call_id": _vcall_id, - "name": _vname, - "content": ( - f"Error: tool '{_vname}' is not available. " - f"Available tools: {_allowed_list}." - ), - }) + conversation.append( + { + "role": "tool", + "tool_call_id": _vcall_id, + "name": _vname, + "content": ( + f"Error: tool '{_vname}' is not available. " + f"Available tools: {_allowed_list}." + ), + } + ) elif _kind == "malformed_args" and _vcall_id: - conversation.append({ - "role": "tool", - "tool_call_id": _vcall_id, - "name": _vname, - "content": ( - f"Error: arguments to '{_vname}' could not " - "be parsed as a JSON object. Call " - f"'{_vname}' again with valid JSON object " - "arguments." - ), - }) + conversation.append( + { + "role": "tool", + "tool_call_id": _vcall_id, + "name": _vname, + "content": ( + f"Error: arguments to '{_vname}' could not " + "be parsed as a JSON object. Call " + f"'{_vname}' again with valid JSON object " + "arguments." + ), + } + ) else: # No usable call id: user-role correction since # OpenAI tool messages require a matching id. - conversation.append({ - "role": "user", - "content": ( - "Your last tool call was malformed " - "(missing id or function). Please " - "re-issue it as a valid OpenAI " - "function call." - ), - }) + conversation.append( + { + "role": "user", + "content": ( + "Your last tool call was malformed " + "(missing id or function). Please " + "re-issue it as a valid OpenAI " + "function call." + ), + } + ) logger.info( "validation_retry kind=%s retries=%d/%d", - _kind, _validation_retries, max_validation_retries, + _kind, + _validation_retries, + max_validation_retries, ) yield {"type": "status", "text": ""} continue diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index b7fbe680e5..196caf89f1 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -331,10 +331,7 @@ def run_safetensors_tool_loop( v_decoded = json.loads(v_raw) except (json.JSONDecodeError, ValueError): v_decoded = None - if ( - not isinstance(v_decoded, dict) - and not auto_heal_tool_calls - ): + if not isinstance(v_decoded, dict) and not auto_heal_tool_calls: validation_problem = ("malformed_args", v_tc, v_name) break elif not isinstance(v_raw, dict): @@ -348,39 +345,47 @@ def run_safetensors_tool_loop( v_call_id = v_tc.get("id") if v_kind == "unknown_tool" and v_call_id: allowed_list = ", ".join(sorted(allowed_tool_names)) - conversation.append({ - "role": "tool", - "tool_call_id": v_call_id, - "name": v_name, - "content": ( - f"Error: tool '{v_name}' is not available. " - f"Available tools: {allowed_list}." - ), - }) + conversation.append( + { + "role": "tool", + "tool_call_id": v_call_id, + "name": v_name, + "content": ( + f"Error: tool '{v_name}' is not available. " + f"Available tools: {allowed_list}." + ), + } + ) elif v_kind == "malformed_args" and v_call_id: - conversation.append({ - "role": "tool", - "tool_call_id": v_call_id, - "name": v_name, - "content": ( - f"Error: arguments to '{v_name}' could not be " - "parsed as a JSON object. Call " - f"'{v_name}' again with valid JSON object " - "arguments." - ), - }) + conversation.append( + { + "role": "tool", + "tool_call_id": v_call_id, + "name": v_name, + "content": ( + f"Error: arguments to '{v_name}' could not be " + "parsed as a JSON object. Call " + f"'{v_name}' again with valid JSON object " + "arguments." + ), + } + ) else: - conversation.append({ - "role": "user", - "content": ( - "Your last tool call was malformed (missing id " - "or function). Please re-issue it as a valid " - "OpenAI function call." - ), - }) + conversation.append( + { + "role": "user", + "content": ( + "Your last tool call was malformed (missing id " + "or function). Please re-issue it as a valid " + "OpenAI function call." + ), + } + ) logger.info( "validation_retry kind=%s retries=%d/%d", - v_kind, validation_retries, max_validation_retries, + v_kind, + validation_retries, + max_validation_retries, ) yield {"type": "status", "text": ""} continue diff --git a/studio/backend/tests/test_validation_retry_loop.py b/studio/backend/tests/test_validation_retry_loop.py index 719683e158..c137ba4271 100644 --- a/studio/backend/tests/test_validation_retry_loop.py +++ b/studio/backend/tests/test_validation_retry_loop.py @@ -21,7 +21,9 @@ class FakeExecuteTool: self.results = list(results or []) self.calls = [] - def __call__(self, name, arguments, *, cancel_event = None, timeout = None, session_id = None): + def __call__( + self, name, arguments, *, cancel_event = None, timeout = None, session_id = None + ): self.calls.append((name, arguments)) return self.results.pop(0) if self.results else "OK" @@ -62,11 +64,15 @@ class TestUnknownTool: def test_unknown_tool_triggers_retry(self): # Turn 1: hallucinated tool name "missing_tool". # Turn 2: valid call after the corrective nudge. - single_turn = _multi_turn([ - ['{"name":"missing_tool","arguments":{}}'], - ['{"name":"web_search","arguments":{"query":"x"}}'], - ["Done."], - ]) + single_turn = _multi_turn( + [ + ['{"name":"missing_tool","arguments":{}}'], + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["Done."], + ] + ) exec_fn = FakeExecuteTool(["result"]) events = _collect( run_safetensors_tool_loop( @@ -87,12 +93,14 @@ class TestUnknownTool: # the first two are retried; the third falls through to the # existing per-tool error path (which never executes a real # tool but emits an error result to the model). - single_turn = _multi_turn([ - ['{"name":"missing_a","arguments":{}}'], - ['{"name":"missing_b","arguments":{}}'], - ['{"name":"missing_c","arguments":{}}'], - ["Sorry, I cannot proceed."], - ]) + single_turn = _multi_turn( + [ + ['{"name":"missing_a","arguments":{}}'], + ['{"name":"missing_b","arguments":{}}'], + ['{"name":"missing_c","arguments":{}}'], + ["Sorry, I cannot proceed."], + ] + ) exec_fn = FakeExecuteTool([]) events = _collect( run_safetensors_tool_loop( @@ -115,10 +123,12 @@ class TestUnknownTool: # With max_validation_retries=0, the F3 arm never engages and # behavior matches the pre-F3 path: the existing per-tool # error message is emitted but no corrective re-entry happens. - single_turn = _multi_turn([ - ['{"name":"missing","arguments":{}}'], - ["bye"], - ]) + single_turn = _multi_turn( + [ + ['{"name":"missing","arguments":{}}'], + ["bye"], + ] + ) exec_fn = FakeExecuteTool([]) events = _collect( run_safetensors_tool_loop( @@ -135,10 +145,7 @@ class TestUnknownTool: tool_ends = [e for e in events if e["type"] == "tool_end"] # Existing per-tool error path emits a tool_end with an Error # result so the model sees the failure. - assert any( - "not enabled" in str(e.get("result", "")) - for e in tool_ends - ) + assert any("not enabled" in str(e.get("result", "")) for e in tool_ends) class TestMalformedArgs: @@ -146,10 +153,14 @@ class TestMalformedArgs: # With auto_heal_tool_calls=True (the default), string args are # healed to {"query": "..."} for web_search. F3 leaves heal # behavior intact and only catches strictly-impossible shapes. - single_turn = _multi_turn([ - ['{"name":"web_search","arguments":"some text"}'], - ["all done"], - ]) + single_turn = _multi_turn( + [ + [ + '{"name":"web_search","arguments":"some text"}' + ], + ["all done"], + ] + ) exec_fn = FakeExecuteTool(["ok"]) events = _collect( run_safetensors_tool_loop( @@ -168,11 +179,17 @@ class TestMalformedArgs: def test_malformed_args_caught_when_heal_off(self): # With auto_heal off, a non-dict arguments value is a hard # malformed-args failure and the F3 arm catches it. - single_turn = _multi_turn([ - ['{"name":"web_search","arguments":"not a dict"}'], - ['{"name":"web_search","arguments":{"query":"sf"}}'], - ["sunny"], - ]) + single_turn = _multi_turn( + [ + [ + '{"name":"web_search","arguments":"not a dict"}' + ], + [ + '{"name":"web_search","arguments":{"query":"sf"}}' + ], + ["sunny"], + ] + ) exec_fn = FakeExecuteTool(["sunny in sf"]) events = _collect( run_safetensors_tool_loop( @@ -192,10 +209,12 @@ class TestNoOpCases: def test_no_tools_no_validation(self): # Empty tools list means allowed_tool_names is empty, so the # F3 arm never engages. - single_turn = _multi_turn([ - ['{"name":"anything","arguments":{}}'], - ["done"], - ]) + single_turn = _multi_turn( + [ + ['{"name":"anything","arguments":{}}'], + ["done"], + ] + ) exec_fn = FakeExecuteTool([]) events = _collect( run_safetensors_tool_loop( @@ -211,10 +230,14 @@ class TestNoOpCases: def test_valid_call_no_retry_overhead(self): # A clean valid call does not trigger F3 at all. - single_turn = _multi_turn([ - ['{"name":"web_search","arguments":{"query":"x"}}'], - ["final"], - ]) + single_turn = _multi_turn( + [ + [ + '{"name":"web_search","arguments":{"query":"x"}}' + ], + ["final"], + ] + ) exec_fn = FakeExecuteTool(["res"]) events = _collect( run_safetensors_tool_loop( From 828f89abbe5d1f86d7f52f7080ed97fe1c99ede8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 15:49:30 +0000 Subject: [PATCH 3/3] Skip malformed-args validation when auto-heal is on The F3 pre-dispatch validator flagged any non-object tool-call ``arguments`` as malformed, even when the caller had set ``auto_heal_tool_calls=True``. The dispatcher downstream already heals bare-string arguments (for example a raw web_search query) into a valid ``{"query": ...}`` shape, so rejecting them up front made the validation loop fight the heal and effectively disabled the auto-heal feature for that family of models. Gate the malformed-args branch on ``auto_heal_tool_calls`` being off so dispatch keeps its existing healing semantics. The unknown-tool branch still fires in either mode because no amount of healing can invent a tool that is not registered. Existing tests in tests/test_validation_retry_loop.py already pin both paths (``test_malformed_args_bypassed_when_heal_on`` and ``test_malformed_args_caught_when_heal_off``); both pass with this change. --- studio/backend/core/inference/llama_cpp.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index da25d91678..de7fe49180 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -5020,6 +5020,15 @@ class LlamaCppBackend: # tool-result tied to the hallucinated call id and # re-enter the model. Skips when the budget is spent so # the existing error path still runs. + # + # ``malformed_args`` (non-object arguments) only fires + # when auto-heal is OFF. With auto_heal_tool_calls=True + # the dispatch path downstream coerces bare-string + # arguments (e.g. a raw web_search query) into a valid + # `{"query": ...}` shape, and rejecting those calls + # here would defeat the heal. The unknown-tool branch + # still runs in either mode because no amount of + # healing can invent a tool that isn't registered. _validation_problem = None if ( tool_calls @@ -5032,6 +5041,10 @@ class LlamaCppBackend: if _vname not in _allowed_tool_names: _validation_problem = ("unknown_tool", _vtc, _vname) break + if auto_heal_tool_calls: + # Dispatch will heal non-object arguments; do + # not pre-empt it by rejecting them here. + continue _vraw = _vfn.get("arguments", "") if isinstance(_vraw, str): try: