From b3e244d658a6edc90e802626a2ea350e5d92238e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 11:03:21 +0000 Subject: [PATCH] Harden Gemma tool-call parsing and stream-error detection Address three issues in the Gemma-native tool-call path: - _quote_gemma_object_keys stopped a bare (unquoted) string value at the first comma, so an argument like `location:New York, NY` was split mid-value and the synthesized JSON failed to parse, dropping the whole tool call. A bare value now ends only at `}` or a comma that begins the next `key:` pair. - parse_tool_calls_from_text scanned the entire response for Gemma markers even inside a tool call already parsed from a `{...}` JSON block, so a marker-like string inside an argument (data) was promoted to a second, unintended tool call. Matches inside an already-consumed call span are now skipped. - _openai_passthrough_stream relied on _monitor_openai_sse_line to flag a stream error, which returns early when monitor_id is None (skip_api_monitor), so an upstream error chunk left saw_stream_error unset and the synthetic-finish guard emitted a successful finish_reason after a failed stream. Error chunks are now detected independently of API monitoring. Adds tests/test_gemma_tool_parse_edge_cases.py covering the comma and marker-injection cases. --- studio/backend/core/tool_healing.py | 21 +++++- studio/backend/routes/inference.py | 7 ++ .../tests/test_gemma_tool_parse_edge_cases.py | 64 +++++++++++++++++++ 3 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 studio/backend/tests/test_gemma_tool_parse_edge_cases.py diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 4663512033..e91b541c0d 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -37,6 +37,10 @@ _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 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. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[\w-]+\s*:") def _balanced_brace_end( @@ -146,7 +150,14 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[ws:i]) if i < len(src) and src[i] not in '"{[': v_start = i - while i < len(src) and src[i] not in ",}": + # 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 + if src[i] == "," and _GEMMA_NEXT_KEY_RE.match(src, i + 1): + break i += 1 raw = src[v_start:i] try: @@ -196,6 +207,10 @@ def parse_tool_calls_from_text( ... """ tool_calls: list[dict] = [] + # Byte spans already claimed by a parsed tool call. A tool-call marker that + # appears INSIDE another call's argument string is data, not a real call, so + # it must not be re-parsed into a spurious second call. + consumed: list[tuple[int, int]] = [] for m in _TC_JSON_START_RE.finditer(content): brace_start = m.end() - 1 @@ -220,10 +235,13 @@ def parse_tool_calls_from_text( if isinstance(tc["function"]["arguments"], dict): tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) tool_calls.append(tc) + consumed.append((m.start(), i + 1)) except (json.JSONDecodeError, ValueError): pass for m in _TC_GEMMA_START_RE.finditer(content): + if any(start <= m.start() < end for start, end in consumed): + continue brace_start = m.end() - 1 i = _balanced_brace_end(content, brace_start, gemma_quotes = True) if i < 0: @@ -243,6 +261,7 @@ def parse_tool_calls_from_text( }, } ) + consumed.append((m.start(), i + 1)) except (json.JSONDecodeError, ValueError): pass diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3ee0319a7d..5c88387e9b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -9577,6 +9577,13 @@ async def _openai_passthrough_stream( delta = choice.get("delta") if isinstance(delta, dict) and delta.get("tool_calls"): saw_tool_call_delta = True + # Detect an upstream error chunk independently of API + # monitoring: when monitor_id is None (skip_api_monitor), + # _monitor_openai_sse_line returns before inspecting the + # error, so without this the synthetic-finish guard would + # emit a successful finish_reason after a failed stream. + if _monitor_openai_error_message(chunk_data): + saw_stream_error = True monitor_event = _monitor_openai_sse_line( monitor_id, raw_line, diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py new file mode 100644 index 0000000000..45e3277a4b --- /dev/null +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Edge cases in Gemma-native tool-call parsing. + +Covers two failure modes: + 1. A bare (unquoted) string argument that contains a comma, e.g. + ``location:New York, NY`` -- the comma must not be treated as the next + key boundary, or the whole call is dropped. + 2. A tool-call marker that appears INSIDE another call's argument string is + data, not a real call, so it must not be promoted to a second tool call. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +from core.inference.tool_call_parser import parse_tool_calls_from_text + + +def _args(call: dict) -> dict: + return json.loads(call["function"]["arguments"]) + + +def test_bare_string_argument_with_comma_is_kept(): + calls = parse_tool_calls_from_text( + "<|tool_call>call:get_weather{location:New York, NY,unit:celsius}" + ) + assert len(calls) == 1, calls + assert calls[0]["function"]["name"] == "get_weather" + assert _args(calls[0]) == {"location": "New York, NY", "unit": "celsius"} + + +def test_normal_multi_key_arguments_still_split(): + calls = parse_tool_calls_from_text('<|tool_call>call:f{a:1,b:hello,c:"x,y"}') + assert len(calls) == 1, calls + # Numbers stay numeric, bare strings get quoted, an explicit quoted comma + # stays inside its value. + assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} + + +def test_marker_inside_json_argument_is_not_a_second_call(): + # A python call whose `code` argument contains a Gemma marker string. The + # marker is data and must not execute as a second `terminal` call. + content = ( + '{"name":"python","arguments":{"code":' + '"x = 1 # <|tool_call>call:terminal{command:ls}"}}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_two_separate_gemma_calls_both_parse(): + content = "<|tool_call>call:a{x:1} and <|tool_call>call:b{y:2}" + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["a", "b"], calls + assert _args(calls[0]) == {"x": 1} + assert _args(calls[1]) == {"y": 2}