studio: tool calling + healing parity for Llama-3, Mistral, Gemma 4 on safetensors + MLX

[danielhanchen]

DRAFT - holding open while we validate end to end with real models. Do not merge until safetensors + MLX behaviour is verified for every covered family.

[gemini-code-assist]

Code Review

This pull request significantly expands the tool-call parsing capabilities to support multiple model families, including Llama-3, Mistral, and Gemma 4, while ensuring parity between GGUF and safetensors backends. It introduces a re-prompting mechanism to nudge models that express intent without taking action and implements per-tool canonical heal keys. Feedback highlights potential issues with regex patterns for Mistral and Llama-3 that could lead to incorrect truncation or leaked markup, a bug in string unescaping that corrupts non-ASCII characters, and an order-dependency in sentinel stripping logic.

[gemini-code-assist]

The regex for Mistral v11+ tool calls uses a non-greedy .*? to match the JSON arguments, which will incorrectly truncate the block at the first closing brace }. If the model emits nested JSON (e.g., {"a": {"b": 1}}), the stripping logic will leave trailing garbage in the assistant stream. A similar issue affects the Mistral pre-v11 array pattern on line 55 if the array contains nested brackets.

[gemini-code-assist]

Using unicode_escape on a UTF-8 encoded byte string will corrupt literal non-ASCII characters (e.g., emojis or non-Latin scripts) emitted by the model. For example, a literal ✨ (U+2728) becomes â\x9c¨ after this transformation. Since the regex already ensures that internal quotes are escaped, a safer way to unescape the string while preserving literal UTF-8 is to use json.loads on the quoted value.

Suggested change

	args[k] = bytes(kv.group(2), "utf-8").decode("unicode_escape")
	args[k] = json.loads('"' + kv.group(2) + '"')

[gemini-code-assist]

The sentinel stripping logic is order-dependent and will fail to remove multiple sentinels if they appear in an order different from the tuple (e.g., <|eot_id|><|begin_of_text|> will leave <|begin_of_text|> behind because the loop already passed it). It should use a loop that repeatedly checks for any of the sentinels at the start of the string.

Suggested change

	for sentinel in (
	"<|begin_of_text|>",
	"<|eot_id|>",
	"<|start_header_id|>",
	"<|end_header_id|>",
	"<|eom_id|>",
	):
	stripped = stripped.lstrip()
	if stripped.startswith(sentinel):
	stripped = stripped[len(sentinel) :]
	while True:
	stripped = stripped.lstrip()
	found = False
	for sentinel in (
	"<|begin_of_text|>",
	"<|eot_id|>",
	"<|start_header_id|>",
	"<|end_header_id|>",
	"<|eom_id|>",
	):
	if stripped.startswith(sentinel):
	stripped = stripped[len(sentinel) :]
	found = True
	break
	if not found:
	break

References

1. Combine checks and transformations into a single loop to improve efficiency and handle multiple sentinels correctly.

[gemini-code-assist]

The regex for stripping Llama-3 tags is too restrictive. By stopping at the first newline or < character, it will leak the remainder of the tool call to the user if the call contains a < (common in queries or code) or spans multiple lines. This results in internal markup being visible in the chat UI during streaming.

[danielhanchen]

Looks safe to land. Pulled pr-5620 and verified against pr-5615.

What 5620 adds beyond 5615 (parser/inference scope):

File	LOC	What's new
tool_call_parser.py	+137	Balanced brace/bracket Mistral stripping, sentinel-loop, UTF-8-safe KV decoder
safetensors_agentic.py	+82	Intent re-prompt loop, _MAX_REPROMPTS=3, BUFFERING fall-through to STREAMING
llama_cpp.py	+39	Shared TOOL_XML_SIGNALS + per-tool canonical heal key (parity with safetensors)
routes/inference.py	+46	All 9 _TOOL_XML_RE.sub sites switched to _strip_tool_xml()
tests/test_safetensors_tool_loop.py	+253	TestLoopRePrompt (6) + TestLoopCanonicalHealKey (3)

Each of the five #5615 regressions has a fix + test in 5620 — details in #5615/#5619.

Tests: pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q → 110 passed in 3.86s (had to pip install structlog — missing dev dep, not a code issue).

Worth watching when the validation phase finishes:

• False-positive tool detection on prose JSON. _parse_llama3_bare_json (tool_call_parser.py:486-547) triggers on any leading { with name + parameters|arguments. The strict parameters/arguments dict guard mitigates this, but a user echoing back {"name":"foo","arguments":{}} in a code fence isn't stripped first — worth a quick fence-aware preamble check.
• python_tag multi-line code. The full parser (_parse_llama3_python_tag) handles multi-line via _balanced_brace_end, but the routes-level _strip_tool_xml() helper is single-line ([^\n]*). Multi-line python.call(code="…\n…") will have only the first line stripped at the route layer; remaining lines may leak in non-parser paths.
• Partial Mistral buffer. _strip_mistral_closed_calls correctly leaves unclosed [TOOL_CALLS]… in place, but the BUFFERING state machine caps at _MAX_BUFFER_CHARS=32 (safetensors_agentic.py:44). A long [TOOL_CALLS] very_long_name [ARGS] {… that exceeds 32 chars before reaching { may flush to user. Worth a test that names a tool with a long identifier and asserts no early flush.
• Re-prompt bound is _MAX_REPROMPTS=3 and _tool_iters_done = iteration + 1 - reprompt_count (safetensors_agentic.py:447), with test_max_reprompts_capped_at_three (test_safetensors_tool_loop.py:893) covering the cap. Looks correct.

Holding the [DRAFT] flag while the end-to-end real-model validation finishes is the right call. Once that's green, this is ready to merge.

[danielhanchen]

Pushed 7ef4b115 closing the routes-level multi-line python_tag gap flagged earlier. The full parser already handled this via balanced-brace scanning; the leak was only on the streaming strip path that runs while content is still arriving.

Earlier history of _TOOL_XML_RE's python_tag clause:

Revision	Pattern	Bug
#5615	<|python_tag|>[^\n<]*	Stopped at literal < so code="if x < 10" leaked < 10: pass)
#5620 round 1	<|python_tag|>[^\n]*	Single-line only — the second line of python.call(code="a\nb") leaked
this commit	<|python_tag|>(?:[^<]|<(?!|))*	Consumes everything that is NOT a Llama-3 <|sentinel|> start. Multi-line code, embedded JSON, and literal < in code all stay inside the strip.

pytest studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_safetensors_capability_advertise.py -q → 118 passed in 1.81s (was 110, +8 new tests in TestRoutesPythonTagStrip covering single-line, less-than-in-code, multi-line, multi-line-with-less-than, sentinel-stop, and JSON-form multi-line cases).

The _MAX_BUFFER_CHARS=32 concern I raised earlier turned out to be unfounded — the BUFFERING state machine treats a fully arrived [TOOL_CALLS] as a match (length 12, well under 32) and transitions to DRAINING, so a long-tool-name [TOOL_CALLS] very_long_name [ARGS] {… never accumulates 32 chars in BUFFERING. The 32-char cap only applies to partial signal prefixes (e.g. <|py waiting to complete to <|python_tag|>), and the longest signal is 14 chars (<|python_tag|>), so the cap is already more than 2× the worst-case prefix length.

The remaining false-positive concern (a model echoing back {"name":"foo","arguments":{}} as the very first non-whitespace characters of an assistant response triggering _parse_llama3_bare_json) is bounded by the strict parameters/arguments dict guard, which keeps plain prose unaffected — happy to add a fence-aware preamble check in a follow-up if you want belt-and-suspenders on that path too.

[danielhanchen]

Confirmed the Llama-3 side of this on a B200 against the matching unsloth UD-Q4_K_XL GGUF. Llama-3.1-8B safetensors fires tools on 18/18 runs (6 prompts, 3 seeds, same generation params as the GGUF), matching the GGUF 18/18 and passing 5/6 of the deterministic code checks vs 2/6 for the GGUF. The bare-JSON form {"name": ..., "parameters": ...} is the key case, since the skip-special-token stream drops <|python_tag|> and only the multi-format parser recovers it.

Mistral-Small-3.2 refuses to call tools in both backends for these prompts, so that one is model behaviour rather than the parser. Full per-family numbers and the cross-checks against llama.cpp, vLLM, and SGLang are on the stacked PR #5624.