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/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(

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,111 @@ 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.
+                #
+                # ``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
+                    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
+                        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:
+                                _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", "")

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,92 @@ 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 ""

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,

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()

studio/backend/tests/test_validation_retry_loop.py

@@ -0,0 +1,250 @@
+# 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(
+            [
+                ['<tool_call>{"name":"missing_tool","arguments":{}}</tool_call>'],
+                [
+                    '<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
+                ],
+                ["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(
+            [
+                ['<tool_call>{"name":"missing_a","arguments":{}}</tool_call>'],
+                ['<tool_call>{"name":"missing_b","arguments":{}}</tool_call>'],
+                ['<tool_call>{"name":"missing_c","arguments":{}}</tool_call>'],
+                ["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(
+            [
+                ['<tool_call>{"name":"missing","arguments":{}}</tool_call>'],
+                ["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(
+            [
+                [
+                    '<tool_call>{"name":"web_search","arguments":"some text"}</tool_call>'
+                ],
+                ["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(
+            [
+                [
+                    '<tool_call>{"name":"web_search","arguments":"not a dict"}</tool_call>'
+                ],
+                [
+                    '<tool_call>{"name":"web_search","arguments":{"query":"sf"}}</tool_call>'
+                ],
+                ["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(
+            [
+                ['<tool_call>{"name":"anything","arguments":{}}</tool_call>'],
+                ["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(
+            [
+                [
+                    '<tool_call>{"name":"web_search","arguments":{"query":"x"}}</tool_call>'
+                ],
+                ["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"})]