diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 5c99773b95..14d87f69e9 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -80,6 +80,80 @@ def _balanced_brace_end( return -1 +def _balanced_bracket_end(src: str, start: int) -> int: + """Index of the ``]`` matching the ``[`` at ``start``, or -1. Tracks nested + ``[]``/``{}`` and double-quoted strings.""" + depth = 0 + i = start + in_string = False + while i < len(src): + ch = src[i] + if in_string: + if ch == "\\" and i + 1 < len(src): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch in "[{": + depth += 1 + elif ch in "]}": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _split_top_level_commas(src: str) -> list: + """Split on commas that are not inside a nested ``[]``/``{}`` or a string.""" + parts: list[str] = [] + depth = 0 + in_string = False + start = 0 + i = 0 + while i < len(src): + ch = src[i] + if in_string: + if ch == "\\" and i + 1 < len(src): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch in "[{": + depth += 1 + elif ch in "]}": + depth -= 1 + elif ch == "," and depth == 0: + parts.append(src[start:i]) + start = i + 1 + i += 1 + parts.append(src[start:]) + return parts + + +def _quote_gemma_array_elements(body: str) -> str: + """Quote bare (unquoted) string elements in a Gemma array value. Gemma may + emit ``labels:[bug,ui]`` without per-element quotes; left as-is json.loads + fails and the whole call is dropped. 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() + if not stripped or stripped[0] in '"{[': + out.append(element) + continue + try: + json.loads(stripped) + out.append(element) + except (json.JSONDecodeError, ValueError): + out.append(json.dumps(stripped)) + return ",".join(out) + + def _normalise_gemma_quoted_strings(src: str) -> str: parts: list[str] = [] i = 0 @@ -148,7 +222,17 @@ def _quote_gemma_object_keys(src: str) -> str: while i < len(src) and src[i].isspace(): i += 1 parts.append(src[ws:i]) - if i < len(src) and src[i] not in '"{[': + 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:]) + i = len(src) + else: + parts.append("[" + _quote_gemma_array_elements(src[i + 1 : arr_end]) + "]") + i = arr_end + 1 + elif i < len(src) and src[i] not in '"{': v_start = i # Consume the bare value up to `}` or a comma that starts the # next key:value pair; a comma inside the value (e.g. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a7b04c7c53..01ed86207a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -9599,6 +9599,25 @@ async def _openai_passthrough_stream( ) if monitor_event == "error": saw_stream_error = True + # If a trailing usage-only chunk (include_usage) arrives before + # any finish chunk, emit the synthetic finish first so the order + # stays finish -> usage -> [DONE], matching the other streams. + if ( + isinstance(chunk_data, dict) + and chunk_data.get("usage") + and not ( + isinstance(chunk_data.get("choices"), list) and chunk_data["choices"] + ) + and not saw_finish_reason + and not saw_stream_error + and not cancel_event.is_set() + ): + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, finish_line, llama_backend.context_length + ) + yield finish_line + "\n\n" + saw_finish_reason = True # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" 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 e828b6dcea..9dd8607bd5 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -85,3 +85,18 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_bare_string_array_argument_is_quoted(): + # Gemma may emit an array of bare strings without per-element quotes; they + # must be quoted so the call is not dropped. + calls = parse_tool_calls_from_text("<|tool_call>call:label{labels:[bug,ui]}") + assert len(calls) == 1, calls + assert _args(calls[0]) == {"labels": ["bug", "ui"]} + + +def test_array_keeps_numbers_and_quoted_elements(): + calls = parse_tool_calls_from_text( + '<|tool_call>call:f{nums:[1,2],tags:[<|"|>a,b<|"|>,c]}' + ) + assert _args(calls[0]) == {"nums": [1, 2], "tags": ["a,b", "c"]}