From 165f4838e304dd0f650ba32f36cf42c8b63e0c59 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 16:27:33 +0200 Subject: [PATCH 01/30] Fix Gemma 4 GGUF OpenAI API streams --- .../core/inference/tool_call_parser.py | 168 +++++++++-- studio/backend/core/tool_healing.py | 156 +++++++++-- studio/backend/models/inference.py | 4 + studio/backend/routes/inference.py | 109 ++++++-- .../tests/test_llama_route_timeouts.py | 29 ++ studio/backend/tests/test_mcp_servers.py | 23 ++ .../tests/test_openai_tool_passthrough.py | 261 ++++++++++++++++++ .../tests/test_responses_tool_passthrough.py | 52 +++- .../tests/test_safetensors_tool_loop.py | 45 +++ .../tests/test_tool_call_parser_strict.py | 37 +++ 10 files changed, 803 insertions(+), 81 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 8d5d45269e..d44c1cea1d 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -16,16 +16,18 @@ import re # so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins. _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), + re.compile(r"<\|tool_call>.*?", re.DOTALL), re.compile(r".*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] # Prefixes the streaming buffer watches for to gate in-progress text. -TOOL_XML_SIGNALS = ("", "", "<|tool_call>", "\s*\{") +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") +_TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") # [\w-] so hyphenated MCP param names (issue-number) aren't dropped. _TC_PARAM_START_RE = re.compile(r"\s*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" +_GEMMA_QUOTE = '<|"|>' def _inside_open_parameter(content: str, pos: int) -> bool: @@ -111,6 +116,116 @@ def strip_tool_markup(text: str, *, final: bool = False) -> str: return text.strip() if final else text +def _balanced_brace_end(content: str, brace_start: int) -> int: + depth = 0 + i = brace_start + in_string = False + in_gemma_string = False + while i < len(content): + if content.startswith(_GEMMA_QUOTE, i): + in_gemma_string = not in_gemma_string + i += len(_GEMMA_QUOTE) + continue + ch = content[i] + if in_gemma_string: + i += 1 + continue + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _normalise_gemma_quoted_strings(src: str) -> str: + parts: list[str] = [] + i = 0 + while i < len(src): + if not src.startswith(_GEMMA_QUOTE, i): + parts.append(src[i]) + i += 1 + continue + end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE)) + if end < 0: + parts.append(src[i:]) + break + raw_value = src[i + len(_GEMMA_QUOTE) : end] + parts.append(json.dumps(raw_value)) + i = end + len(_GEMMA_QUOTE) + return "".join(parts) + + +def _quote_gemma_object_keys(src: str) -> str: + parts: list[str] = [] + i = 0 + in_string = False + while i < len(src): + ch = src[i] + if in_string: + parts.append(ch) + if ch == "\\" and i + 1 < len(src): + parts.append(src[i + 1]) + i += 2 + continue + if ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + parts.append(ch) + i += 1 + continue + if ch not in "{,": + parts.append(ch) + i += 1 + continue + + parts.append(ch) + i += 1 + key_start = i + while i < len(src) and src[i].isspace(): + i += 1 + key_name_start = i + while i < len(src) and (src[i].isalnum() or src[i] in "_-"): + i += 1 + key_name = src[key_name_start:i] + colon_pos = i + while colon_pos < len(src) and src[colon_pos].isspace(): + colon_pos += 1 + if key_name and colon_pos < len(src) and src[colon_pos] == ":": + parts.append(src[key_start:key_name_start]) + parts.append(json.dumps(key_name)) + parts.append(src[i:colon_pos]) + parts.append(":") + i = colon_pos + 1 + else: + parts.append(src[key_start:i]) + return "".join(parts) + + +def _gemma_arguments_to_json(args_src: str) -> dict: + """Parse Gemma 4's native call:name{key:value} argument object.""" + args_src = args_src.strip() + if not args_src: + return {} + src = _normalise_gemma_quoted_strings(args_src) + src = "{" + src + "}" + src = _quote_gemma_object_keys(src) + return json.loads(src) + + def parse_tool_calls_from_text( content: str, *, @@ -123,10 +238,12 @@ def parse_tool_calls_from_text( dicts. ``arguments`` is always a JSON string so callers can hand it straight back into an OpenAI-style response. - Handles two shapes: + Handles three shapes: - JSON inside ```` tags: ``{"name":"web_search","arguments":{"query":"..."}}`` + - Gemma 4 native call blocks: + ``<|tool_call>call:web_search{query:"..." }`` - XML-style function blocks: ``v`` @@ -141,26 +258,8 @@ def parse_tool_calls_from_text( # JSON strings. for m in _TC_JSON_START_RE.finditer(content): brace_start = m.end() - 1 # opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 - continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - break - i += 1 - if depth != 0: + i = _balanced_brace_end(content, brace_start) + if i < 0: continue if not allow_incomplete: tail_after_json = content[i + 1 :].lstrip() @@ -183,6 +282,31 @@ def parse_tool_calls_from_text( except (json.JSONDecodeError, ValueError): pass + # Pattern 1b: Gemma 4 native call block: + # <|tool_call>call:terminal{command:"ls"} + for m in _TC_GEMMA_START_RE.finditer(content): + brace_start = m.end() - 1 + i = _balanced_brace_end(content, brace_start) + if i < 0: + continue + if not allow_incomplete: + tail_after_json = content[i + 1 :].lstrip() + if _TC_GEMMA_END_TAG_RE.match(tail_after_json) is None: + continue + try: + tool_calls.append( + { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])), + }, + } + ) + except (json.JSONDecodeError, ValueError): + pass + # Pattern 2: v... -- closing tags optional; # isn't a body boundary since code values can contain it. if not tool_calls: diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 973520d5cd..f3e52534e8 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -19,20 +19,133 @@ import re # issue-number) parse alongside the built-ins. _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), + re.compile(r"<\|tool_call>.*?", re.DOTALL), re.compile(r".*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") +_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") _TC_PARAM_START_RE = re.compile(r"\s*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") +_GEMMA_QUOTE = '<|"|>' + + +def _balanced_brace_end(content: str, brace_start: int) -> int: + depth = 0 + i = brace_start + in_string = False + in_gemma_string = False + while i < len(content): + if content.startswith(_GEMMA_QUOTE, i): + in_gemma_string = not in_gemma_string + i += len(_GEMMA_QUOTE) + continue + ch = content[i] + if in_gemma_string: + i += 1 + continue + if in_string: + if ch == "\\" and i + 1 < len(content): + i += 2 + continue + if ch == '"': + in_string = False + elif ch == '"': + in_string = True + elif ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if depth == 0: + return i + i += 1 + return -1 + + +def _normalise_gemma_quoted_strings(src: str) -> str: + parts: list[str] = [] + i = 0 + while i < len(src): + if not src.startswith(_GEMMA_QUOTE, i): + parts.append(src[i]) + i += 1 + continue + end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE)) + if end < 0: + parts.append(src[i:]) + break + raw_value = src[i + len(_GEMMA_QUOTE) : end] + parts.append(json.dumps(raw_value)) + i = end + len(_GEMMA_QUOTE) + return "".join(parts) + + +def _quote_gemma_object_keys(src: str) -> str: + parts: list[str] = [] + i = 0 + in_string = False + while i < len(src): + ch = src[i] + if in_string: + parts.append(ch) + if ch == "\\" and i + 1 < len(src): + parts.append(src[i + 1]) + i += 2 + continue + if ch == '"': + in_string = False + i += 1 + continue + if ch == '"': + in_string = True + parts.append(ch) + i += 1 + continue + if ch not in "{,": + parts.append(ch) + i += 1 + continue + + parts.append(ch) + i += 1 + key_start = i + while i < len(src) and src[i].isspace(): + i += 1 + key_name_start = i + while i < len(src) and (src[i].isalnum() or src[i] in "_-"): + i += 1 + key_name = src[key_name_start:i] + colon_pos = i + while colon_pos < len(src) and src[colon_pos].isspace(): + colon_pos += 1 + if key_name and colon_pos < len(src) and src[colon_pos] == ":": + parts.append(src[key_start:key_name_start]) + parts.append(json.dumps(key_name)) + parts.append(src[i:colon_pos]) + parts.append(":") + i = colon_pos + 1 + else: + parts.append(src[key_start:i]) + return "".join(parts) + + +def _gemma_arguments_to_json(args_src: str) -> dict: + args_src = args_src.strip() + if not args_src: + return {} + src = _normalise_gemma_quoted_strings(args_src) + src = "{" + src + "}" + src = _quote_gemma_object_keys(src) + return json.loads(src) def parse_tool_calls_from_text(content: str) -> list[dict]: @@ -41,6 +154,7 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: Handles formats like: {"name":"web_search","arguments":{"query":"..."}} + <|tool_call>call:web_search{query:"..."} ... Closing tags (, , ) are all optional since models frequently omit them. @@ -51,26 +165,8 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: # skips braces inside JSON strings. for m in _TC_JSON_START_RE.finditer(content): brace_start = m.end() - 1 # position of the opening { - depth, i = 0, brace_start - in_string = False - while i < len(content): - ch = content[i] - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 # skip escaped character - continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - break - i += 1 - if depth == 0: + i = _balanced_brace_end(content, brace_start) + if i >= 0: json_str = content[brace_start : i + 1] try: obj = json.loads(json_str) @@ -88,6 +184,26 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: except (json.JSONDecodeError, ValueError): pass + # Pattern 1b: Gemma 4 native <|tool_call>call:name{key:value}. + for m in _TC_GEMMA_START_RE.finditer(content): + brace_start = m.end() - 1 + i = _balanced_brace_end(content, brace_start) + if i < 0: + continue + try: + tool_calls.append( + { + "id": f"call_{len(tool_calls)}", + "type": "function", + "function": { + "name": m.group(1), + "arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])), + }, + } + ) + except (json.JSONDecodeError, ValueError): + pass + # Pattern 2: XML-style value # All closing tags optional; models frequently omit them. if not tool_calls: diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0d33cfa976..d520896798 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1101,6 +1101,8 @@ class ChoiceDelta(BaseModel): role: Optional[str] = None content: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None OpenAIFinishReason = Literal["stop", "length", "tool_calls", "content_filter", "function_call"] @@ -1136,6 +1138,8 @@ class CompletionMessage(BaseModel): role: Literal["assistant"] = "assistant" content: str refusal: Optional[str] = None + reasoning_content: Optional[str] = None + tool_calls: Optional[list[dict]] = None class CompletionChoice(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 4b068941b8..e658182dfd 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -787,7 +787,10 @@ async def _aiter_llama_stream_items( raise httpx.ReadTimeout("The model did not produce a first token in time.") if response is not None: _set_stream_response_read_timeout(response, remaining_s) - item = await asyncio.wait_for(async_iter.__anext__(), timeout = remaining_s) + # Keep httpx/httpcore's AnyIO cancel scope in this task. + # asyncio.wait_for would drive __anext__ in a child task. + async with asyncio.timeout(remaining_s): + item = await async_iter.__anext__() else: item = await async_iter.__anext__() except asyncio.TimeoutError as exc: @@ -1286,6 +1289,7 @@ _TOOL_XML_RE = _re.compile( # Hyphen in the name char-class matches MCP tool names with dashes # (mcp__srv__list-issues) that would otherwise leak past this strip. r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" + r"|<\|tool_call>.*?(?:|\Z)" r"|" r"|\s*\Z", _re.DOTALL, @@ -4842,6 +4846,28 @@ async def openai_chat_completions( completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) + def _new_chat_reasoning_extractor(): + return _ResponsesReasoningExtractor( + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ) + ) + + def _gguf_chat_delta_line(delta: ChoiceDelta, finish_reason = None) -> str: + chunk = ChatCompletionChunk( + id = completion_id, + created = created, + model = model_name, + choices = [ + ChunkChoice( + delta = delta, + finish_reason = finish_reason, + ) + ], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + # ── Tool-calling path (agentic loop) ────────────────── # `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools` # hard-override the per-request value, else falls back to @@ -5002,6 +5028,7 @@ async def openai_chat_completions( # stays free for disconnect detection. gen = gguf_generate_with_tools() prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() _stream_usage = None _stream_timings = None _stream_finish = None @@ -5024,6 +5051,7 @@ async def openai_chat_completions( # streams cleanly. if not event["text"]: prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() # Emit tool status as a custom SSE event (including # empty ones to clear UI badges) status_data = json.dumps( @@ -5038,6 +5066,7 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" continue @@ -5059,19 +5088,23 @@ async def openai_chat_completions( prev_text = clean_cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [ - ChunkChoice( - delta = ChoiceDelta(content = new_text), - finish_reason = None, - ) - ], + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = reasoning_delta) + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) + + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = final_reasoning) ) - yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) final_chunk = ChatCompletionChunk( id = completion_id, @@ -5186,6 +5219,7 @@ async def openai_chat_completions( # stays free for disconnect detection. gen = gguf_generate() prev_text = "" + reasoning_extractor = _new_chat_reasoning_extractor() _stream_usage = None _stream_timings = None _stream_finish = None @@ -5219,19 +5253,23 @@ async def openai_chat_completions( prev_text = cumulative if not new_text: continue - api_monitor.append_reply(monitor_id, new_text) - chunk = ChatCompletionChunk( - id = completion_id, - created = created, - model = model_name, - choices = [ - ChunkChoice( - delta = ChoiceDelta(content = new_text), - finish_reason = None, - ) - ], + reasoning_delta, visible_delta = reasoning_extractor.feed(new_text) + if reasoning_delta: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = reasoning_delta) + ) + if visible_delta: + api_monitor.append_reply(monitor_id, visible_delta) + yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) + + final_reasoning, final_visible = reasoning_extractor.finish() + if final_reasoning: + yield _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = final_reasoning) ) - yield f"data: {chunk.model_dump_json(exclude_none = True)}\n\n" + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) # Final chunk final_chunk = ChatCompletionChunk( @@ -5309,14 +5347,24 @@ async def openai_chat_completions( continue full_text = token + reasoning_text, visible_text = _extract_responses_reasoning( + full_text, + parse_think_markers = _responses_should_parse_think_markers( + payload, + llama_backend, + ), + ) + message_kwargs = {"content": visible_text} + if reasoning_text: + message_kwargs["reasoning_content"] = reasoning_text _choices.append( CompletionChoice( index = _idx, - message = CompletionMessage(content = full_text), + message = CompletionMessage(**message_kwargs), finish_reason = _clamp_finish_reason(completion_finish), ) ) - _monitor_replies.append(full_text) + _monitor_replies.append(visible_text) if completion_usage: # The prompt is shared across all n choices, so count its # tokens ONCE (OpenAI bills only generated tokens for each @@ -5338,7 +5386,7 @@ async def openai_chat_completions( prompt_tokens_details = _prompt_tokens_details(_prompt_details), ), ) - monitor_reply = full_text + monitor_reply = _monitor_replies[-1] if _monitor_replies else "" if _n > 1: monitor_reply = "\n\n".join( f"Choice {_idx + 1}:\n{text}" for _idx, text in enumerate(_monitor_replies) @@ -6686,8 +6734,9 @@ def _responses_should_parse_think_markers( if llama_backend is not None and getattr(llama_backend, "is_loaded", False): if getattr(llama_backend, "reasoning_always_on", False): return True - if not getattr(llama_backend, "supports_reasoning", False): - return False + if getattr(llama_backend, "supports_reasoning", False): + return True + return False if chat_req.enable_thinking is True: return True return chat_req.enable_thinking is None and chat_req.reasoning_effort not in (None, "none") diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index 5aee6198ba..dc6e13a3de 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -40,6 +40,35 @@ def test_stream_first_item_deadline_after_headers(): asyncio.run(_run()) +def test_stream_first_item_deadline_does_not_hop_tasks(): + async def _run(): + outer_task = asyncio.current_task() + seen_tasks = [] + + class _One: + def __init__(self): + self.done = False + + async def __anext__(self): + seen_tasks.append(asyncio.current_task()) + if self.done: + raise StopAsyncIteration + self.done = True + return "data: {}" + + out = [] + async for item in inf_mod._aiter_llama_stream_items( + _One(), + first_token_deadline = time.monotonic() + 1, + ): + out.append(item) + + assert out == ["data: {}"] + assert seen_tasks == [outer_task, outer_task] + + asyncio.run(_run()) + + def test_preheader_send_cleanup_on_disconnect_and_cancel(): async def _run(cancel_parent): state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 90b1ade03c..a61368f973 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -423,6 +423,29 @@ def test_tool_healing_strip_handles_hyphenated_function_names(): assert out == "before after" +def test_tool_healing_strip_handles_gemma_native_tool_call(): + from core.tool_healing import strip_tool_call_markup + + out = strip_tool_call_markup( + 'before <|tool_call>call:mcp__srv__list-issues{repo:"octocat/hello"} after' + ) + assert out == "before after" + + +def test_tool_healing_parser_handles_gemma_native_windows_path(): + from core.tool_healing import parse_tool_calls_from_text + import json as _json + + calls = parse_tool_calls_from_text( + r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + ) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "ls" + assert _json.loads(calls[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } + + def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch): """A tool call not in the per-request list must be refused by the GGUF agentic loop (mirroring the safetensors path).""" diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 2586076321..e725d4a523 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1245,6 +1245,24 @@ class TestGgufVisionToolRouting: return TestGgufVisionToolRouting._drive(_consume()) + @staticmethod + def _sse_payloads(chunks): + payloads = [] + for chunk in chunks: + if isinstance(chunk, bytes): + chunk = chunk.decode() + for line in str(chunk).splitlines(): + if not line.startswith("data: "): + continue + data = line.removeprefix("data: ") + if data == "[DONE]": + continue + try: + payloads.append(json.loads(data)) + except json.JSONDecodeError: + pass + return payloads + def test_image_request_with_enabled_tools_enters_gguf_tool_loop(self, monkeypatch): import routes.inference as inf_mod @@ -1390,6 +1408,249 @@ class TestGgufVisionToolRouting: assert "confirm_tool_calls requires stream=true" in entry["error"] assert monitor.active_count() == 0 + def test_standard_gguf_stream_splits_reasoning_content(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + + def _generate(**_kwargs): + yield "plan" + yield "planvis" + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + stream = True, + messages = [{"role": "user", "content": "hi"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + payloads = self._sse_payloads(self._consume_response(response)) + deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + assert "".join(d.get("content", "") for d in deltas) == "visible" + assert all("" not in d.get("content", "") for d in deltas) + [entry] = monitor.snapshot() + assert entry["reply"] == "visible" + + def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + + def _generate(**_kwargs): + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + stream = True, + messages = [{"role": "user", "content": "hi"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + payloads = self._sse_payloads(self._consume_response(response)) + deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + assert "".join(d.get("content", "") for d in deltas) == "visible" + [entry] = monitor.snapshot() + assert entry["reply"] == "visible" + + def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled( + self, monkeypatch + ): + import routes.inference as inf_mod + + reset_tool_policy() + + def _generate(**_kwargs): + yield "leakedvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = False, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + stream = True, + enable_thinking = False, + messages = [{"role": "user", "content": "hi"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + payloads = self._sse_payloads(self._consume_response(response)) + deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "leaked" + assert "".join(d.get("content", "") for d in deltas) == "visible" + assert all("" not in d.get("content", "") for d in deltas) + [entry] = monitor.snapshot() + assert entry["reply"] == "visible" + + def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + def _tools(**_kwargs): + yield { + "type": "content", + "text": 'planvisible <|tool_call>call:terminal{command:"ls"}', + } + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = True, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _plain, + generate_chat_completion_with_tools = _tools, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + stream = True, + enable_tools = True, + enabled_tools = ["terminal"], + messages = [{"role": "user", "content": "list files"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + payloads = self._sse_payloads(self._consume_response(response)) + deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + + assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" + combined_content = "".join(d.get("content", "") for d in deltas) + assert combined_content == "visible " + assert "<|tool_call>" not in combined_content + [entry] = monitor.snapshot() + assert entry["reply"] == "visible " + + def test_non_streaming_gguf_splits_reasoning_content(self, monkeypatch): + import routes.inference as inf_mod + + reset_tool_policy() + + def _generate(**_kwargs): + yield "planvisible" + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + backend = SimpleNamespace( + is_loaded = True, + is_vision = False, + supports_tools = False, + supports_reasoning = True, + reasoning_always_on = True, + _is_audio = False, + model_identifier = "test-gguf", + context_length = 4096, + generate_chat_completion = _generate, + ) + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + payload = ChatCompletionRequest( + model = "default", + messages = [{"role": "user", "content": "hi"}], + ) + + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + body = json.loads(response.body) + message = body["choices"][0]["message"] + + assert message["content"] == "visible" + assert message["reasoning_content"] == "plan" + [entry] = monitor.snapshot() + assert entry["reply"] == "visible" + def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 0bea355668..f9d029babe 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -927,6 +927,38 @@ class TestResponsesNonStreamingAdapter: assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "show x tags" + def test_reasoning_capable_gguf_parses_think_tags_by_default(self, monkeypatch): + body = self._run_with_message( + monkeypatch, + {"content": "plananswer"}, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = True, + ), + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "plan"}] + assert body["output"][1]["content"][0]["text"] == "answer" + + def test_reasoning_capable_gguf_sanitizes_think_tags_when_disabled(self, monkeypatch): + payload = ResponsesRequest(input = "hi", reasoning = {"effort": "none"}) + body = self._run_with_message( + monkeypatch, + {"content": "leakedanswer"}, + payload = payload, + llama_backend = SimpleNamespace( + is_loaded = True, + reasoning_always_on = False, + supports_reasoning = True, + ), + ) + + assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert body["output"][0]["content"] == [{"type": "reasoning_text", "text": "leaked"}] + assert body["output"][1]["content"][0]["text"] == "answer" + def test_structured_reasoning_content_extracts_text_parts(self, monkeypatch): body = self._run_with_message( monkeypatch, @@ -1332,10 +1364,10 @@ class TestResponsesStreamAdapter: assert entry["status"] == "completed" assert entry["reply"] == "plan" - def test_literal_think_tags_stream_as_visible_text_without_reasoning_request(self, monkeypatch): + def test_reasoning_capable_gguf_stream_parses_think_tags_by_default(self, monkeypatch): chunks = [ - {"choices": [{"delta": {"content": "show x tags"}}]}, + {"choices": [{"delta": {"content": "plananswer"}}]}, {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, ] self._install_stream_mock(monkeypatch, chunks) @@ -1350,13 +1382,15 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") - assert reasoning_deltas == [] - assert "".join(event["delta"] for event in text_deltas) == "show x tags" + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert "".join(event["delta"] for event in text_deltas) == "answer" completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == ["message"] - assert completed["response"]["output"][0]["content"][0]["text"] == ( - "show x tags" - ) + assert [item["type"] for item in completed["response"]["output"]] == [ + "reasoning", + "message", + ] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + assert completed["response"]["output"][1]["content"][0]["text"] == "answer" def test_non_reasoning_gguf_stream_keeps_literal_think_tags_visible(self, monkeypatch): chunks = [ diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 4098c87f4b..406764fba1 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -10,6 +10,7 @@ calls, tool-result feedback, bad-JSON heal, duplicate-call short-circuit, ``__IMAGES__`` sentinel stripping, executor errors, cancel, and the iteration cap. """ +import json import threading from typing import cast @@ -62,6 +63,46 @@ class TestParser: assert parse_tool_calls_from_text(text)[0]["function"]["name"] == "python" assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + def test_gemma_native_tool_call(self): + text = '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "terminal" + args = json.loads(result[0]["function"]["arguments"]) + assert args == {"command": "ls -la", "workdir": "."} + + def test_gemma_native_tool_call_template_quotes(self): + text = '<|tool_call>call:web_search{query:<|"|>openai news<|"|>}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "web_search" + assert json.loads(result[0]["function"]["arguments"]) == {"query": "openai news"} + + def test_gemma_native_tool_call_template_quotes_escape_backslashes(self): + text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "ls" + assert json.loads(result[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } + + def test_gemma_native_tool_call_hyphenated_argument_name(self): + text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "mcp__srv__create-issue" + assert json.loads(result[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + + def test_gemma_native_tool_call_keeps_braces_inside_string_value(self): + text = '<|tool_call>call:terminal{command:"echo {foo:bar}"}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert result[0]["function"]["name"] == "terminal" + assert json.loads(result[0]["function"]["arguments"]) == { + "command": "echo {foo:bar}" + } + def test_xml_function_call(self): text = "print('hi')" result = parse_tool_calls_from_text(text) @@ -121,6 +162,7 @@ class TestParser: def test_has_tool_signal(self): assert has_tool_signal("blah x") + assert has_tool_signal("blah <|tool_call>call:terminal") assert has_tool_signal("hi ...") assert not has_tool_signal("hello world") @@ -139,6 +181,8 @@ class TestParser: def test_strip_markup_closed(self): text = "before {} after" assert strip_tool_markup(text) == "before after" + text = 'before <|tool_call>call:terminal{command:"ls"} after' + assert strip_tool_markup(text) == "before after" def test_strip_markup_unclosed_final(self): text = "before {partial" @@ -146,6 +190,7 @@ class TestParser: assert strip_tool_markup(text, final = True) == "before" # Without final=True the unclosed run is preserved. assert "partial" in strip_tool_markup(text) + assert strip_tool_markup("before <|tool_call>call:terminal{", final = True) == "before" def test_streaming_strip_respects_disabled_healing(self): raw = 'before {"name":"web_search"' diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 8ff41342d7..8ba0b5b08d 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -106,6 +106,43 @@ class TestParityWithJsonStyle: assert json.loads(js[0]["function"]["arguments"]) == {"query": q} +class TestGemmaNativeStyle: + def test_closed_native_call_with_trailing_prose_is_accepted(self): + text = ( + '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' + " running it now" + ) + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "terminal" + assert json.loads(calls[0]["function"]["arguments"]) == { + "command": "ls -la", + "workdir": ".", + } + + def test_unclosed_native_call_requires_healing(self): + text = '<|tool_call>call:terminal{command:"ls"}' + assert parse_tool_calls_from_text(text, allow_incomplete = False) == [] + calls = parse_tool_calls_from_text(text, allow_incomplete = True) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "terminal" + + def test_hyphenated_native_argument_name_is_accepted(self): + text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert calls[0]["function"]["name"] == "mcp__srv__create-issue" + assert json.loads(calls[0]["function"]["arguments"]) == {"issue-title": "Bug report"} + + def test_native_template_quotes_preserve_windows_path(self): + text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "path": r"C:\Users\wasim\repo" + } + + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): text = "cats" From 05518f2769b307dc21fe5c0066e07eb46447c602 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jun 2026 14:29:06 +0000 Subject: [PATCH 02/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 8 ++------ studio/backend/tests/test_mcp_servers.py | 5 +---- studio/backend/tests/test_openai_tool_passthrough.py | 4 +--- studio/backend/tests/test_safetensors_tool_loop.py | 8 ++------ studio/backend/tests/test_tool_call_parser_strict.py | 7 ++----- 5 files changed, 8 insertions(+), 24 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index e658182dfd..07f265607a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5099,9 +5099,7 @@ async def openai_chat_completions( final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: - yield _gguf_chat_delta_line( - ChoiceDelta(reasoning_content = final_reasoning) - ) + yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning)) if final_visible: api_monitor.append_reply(monitor_id, final_visible) yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) @@ -5264,9 +5262,7 @@ async def openai_chat_completions( final_reasoning, final_visible = reasoning_extractor.finish() if final_reasoning: - yield _gguf_chat_delta_line( - ChoiceDelta(reasoning_content = final_reasoning) - ) + yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning)) if final_visible: api_monitor.append_reply(monitor_id, final_visible) yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index a61368f973..d1a9662b1d 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -425,7 +425,6 @@ def test_tool_healing_strip_handles_hyphenated_function_names(): def test_tool_healing_strip_handles_gemma_native_tool_call(): from core.tool_healing import strip_tool_call_markup - out = strip_tool_call_markup( 'before <|tool_call>call:mcp__srv__list-issues{repo:"octocat/hello"} after' ) @@ -441,9 +440,7 @@ def test_tool_healing_parser_handles_gemma_native_windows_path(): ) assert len(calls) == 1 assert calls[0]["function"]["name"] == "ls" - assert _json.loads(calls[0]["function"]["arguments"]) == { - "path": r"C:\Users\wasim\repo" - } + assert _json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch): diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index e725d4a523..6d2341b716 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1502,9 +1502,7 @@ class TestGgufVisionToolRouting: [entry] = monitor.snapshot() assert entry["reply"] == "visible" - def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled( - self, monkeypatch - ): + def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled(self, monkeypatch): import routes.inference as inf_mod reset_tool_policy() diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 406764fba1..6b470d87a7 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -83,9 +83,7 @@ class TestParser: result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "ls" - assert json.loads(result[0]["function"]["arguments"]) == { - "path": r"C:\Users\wasim\repo" - } + assert json.loads(result[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} def test_gemma_native_tool_call_hyphenated_argument_name(self): text = '<|tool_call>call:mcp__srv__create-issue{issue-title:"Bug report"}' @@ -99,9 +97,7 @@ class TestParser: result = parse_tool_calls_from_text(text) assert len(result) == 1 assert result[0]["function"]["name"] == "terminal" - assert json.loads(result[0]["function"]["arguments"]) == { - "command": "echo {foo:bar}" - } + assert json.loads(result[0]["function"]["arguments"]) == {"command": "echo {foo:bar}"} def test_xml_function_call(self): text = "print('hi')" diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 8ba0b5b08d..799f38710f 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -109,8 +109,7 @@ class TestParityWithJsonStyle: class TestGemmaNativeStyle: def test_closed_native_call_with_trailing_prose_is_accepted(self): text = ( - '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' - " running it now" + '<|tool_call>call:terminal{command:"ls -la",workdir:"."}' " running it now" ) calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 @@ -138,9 +137,7 @@ class TestGemmaNativeStyle: text = r'<|tool_call>call:ls{path:<|"|>C:\Users\wasim\repo<|"|>}' calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 - assert json.loads(calls[0]["function"]["arguments"]) == { - "path": r"C:\Users\wasim\repo" - } + assert json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} class TestHealingPathUnaffected: From 6705053e9ed990bc0fa0b37f214a36f411b10fa3 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 16:33:32 +0200 Subject: [PATCH 03/30] Avoid duplicate Responses stream disconnect watcher --- studio/backend/routes/inference.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 07f265607a..b680cd5a4a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7337,7 +7337,6 @@ async def _responses_stream( resp = None lines_iter = None disconnect_event = threading.Event() - disconnect_watcher = None try: req = client.build_request( "POST", target_url, json = body, headers = {"Connection": "close"} @@ -7397,14 +7396,10 @@ async def _responses_stream( ) return - disconnect_watcher = asyncio.create_task( - _await_disconnect_then_close(request, resp, disconnect_event) - ) lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( lines_iter, cancel_event = disconnect_event, - request = request, first_token_deadline = first_token_deadline, response = resp, ): @@ -7536,6 +7531,7 @@ async def _responses_stream( llama_backend.context_length, ) except asyncio.CancelledError: + disconnect_event.set() api_monitor.finish(monitor_id, "cancelled") raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: @@ -7561,12 +7557,6 @@ async def _responses_stream( ) return finally: - if disconnect_watcher is not None: - disconnect_watcher.cancel() - try: - await disconnect_watcher - except (asyncio.CancelledError, Exception): - pass if lines_iter is not None: try: await lines_iter.aclose() From 4c877cc7d886ac670fec6b754a48ac092eab9abb Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 16:51:40 +0200 Subject: [PATCH 04/30] Keep reasoning-only Responses output hidden --- .../core/inference/tool_call_parser.py | 147 ++---------------- studio/backend/core/tool_healing.py | 8 +- studio/backend/routes/inference.py | 15 -- .../tests/test_responses_tool_passthrough.py | 51 ++++-- 4 files changed, 56 insertions(+), 165 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index d44c1cea1d..78e8f14103 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -10,20 +10,19 @@ or ``v...`` shape. import json import re - -# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing unclosed -# runs so truncated tails don't leak markup. The [\w-] name set matches OpenAI's -# so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins. -_TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), - re.compile(r".*?", re.DOTALL), -] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), - re.compile(r"<\|tool_call>.*$", re.DOTALL), - re.compile(r".*$", re.DOTALL), -] +from core.tool_healing import ( + _TC_END_TAG_RE, + _TC_FUNC_CLOSE_RE, + _TC_FUNC_START_RE, + _TC_GEMMA_START_RE, + _TC_JSON_START_RE, + _TC_PARAM_CLOSE_RE, + _TC_PARAM_START_RE, + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + _balanced_brace_end, + _gemma_arguments_to_json, +) # Prefixes the streaming buffer watches for to gate in-progress text. @@ -76,19 +75,9 @@ RAG_SEARCH_CAP_NUDGE = ( ) -# Pre-compiled patterns reused by ``parse_tool_calls_from_text``. -_TC_JSON_START_RE = re.compile(r"\s*\{") -_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") -_TC_FUNC_START_RE = re.compile(r"\s*") -_TC_END_TAG_RE = re.compile(r"") _TC_GEMMA_END_TAG_RE = re.compile(r"") -_TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") -# [\w-] so hyphenated MCP param names (issue-number) aren't dropped. -_TC_PARAM_START_RE = re.compile(r"\s*") -_TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -_GEMMA_QUOTE = '<|"|>' def _inside_open_parameter(content: str, pos: int) -> bool: @@ -116,116 +105,6 @@ def strip_tool_markup(text: str, *, final: bool = False) -> str: return text.strip() if final else text -def _balanced_brace_end(content: str, brace_start: int) -> int: - depth = 0 - i = brace_start - in_string = False - in_gemma_string = False - while i < len(content): - if content.startswith(_GEMMA_QUOTE, i): - in_gemma_string = not in_gemma_string - i += len(_GEMMA_QUOTE) - continue - ch = content[i] - if in_gemma_string: - i += 1 - continue - if in_string: - if ch == "\\" and i + 1 < len(content): - i += 2 - continue - if ch == '"': - in_string = False - elif ch == '"': - in_string = True - elif ch == "{": - depth += 1 - elif ch == "}": - depth -= 1 - if depth == 0: - return i - i += 1 - return -1 - - -def _normalise_gemma_quoted_strings(src: str) -> str: - parts: list[str] = [] - i = 0 - while i < len(src): - if not src.startswith(_GEMMA_QUOTE, i): - parts.append(src[i]) - i += 1 - continue - end = src.find(_GEMMA_QUOTE, i + len(_GEMMA_QUOTE)) - if end < 0: - parts.append(src[i:]) - break - raw_value = src[i + len(_GEMMA_QUOTE) : end] - parts.append(json.dumps(raw_value)) - i = end + len(_GEMMA_QUOTE) - return "".join(parts) - - -def _quote_gemma_object_keys(src: str) -> str: - parts: list[str] = [] - i = 0 - in_string = False - while i < len(src): - ch = src[i] - if in_string: - parts.append(ch) - if ch == "\\" and i + 1 < len(src): - parts.append(src[i + 1]) - i += 2 - continue - if ch == '"': - in_string = False - i += 1 - continue - if ch == '"': - in_string = True - parts.append(ch) - i += 1 - continue - if ch not in "{,": - parts.append(ch) - i += 1 - continue - - parts.append(ch) - i += 1 - key_start = i - while i < len(src) and src[i].isspace(): - i += 1 - key_name_start = i - while i < len(src) and (src[i].isalnum() or src[i] in "_-"): - i += 1 - key_name = src[key_name_start:i] - colon_pos = i - while colon_pos < len(src) and src[colon_pos].isspace(): - colon_pos += 1 - if key_name and colon_pos < len(src) and src[colon_pos] == ":": - parts.append(src[key_start:key_name_start]) - parts.append(json.dumps(key_name)) - parts.append(src[i:colon_pos]) - parts.append(":") - i = colon_pos + 1 - else: - parts.append(src[key_start:i]) - return "".join(parts) - - -def _gemma_arguments_to_json(args_src: str) -> dict: - """Parse Gemma 4's native call:name{key:value} argument object.""" - args_src = args_src.strip() - if not args_src: - return {} - src = _normalise_gemma_quoted_strings(args_src) - src = "{" + src + "}" - src = _quote_gemma_object_keys(src) - return json.loads(src) - - def parse_tool_calls_from_text( content: str, *, diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index f3e52534e8..c419d8e21c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -1,14 +1,10 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Tool-call XML parsing and stripping helpers. +"""Lightweight tool-call XML parsing and stripping helpers. -Extracted verbatim from studio/backend/core/inference/llama_cpp.py so external -inference servers can reuse the logic without importing the inference +External inference servers import this module without pulling in the inference orchestrator, structlog, httpx, or the rest of the studio backend. - -Regexes and bodies are byte-for-byte identical to the original; any change must -preserve that. test_tool_healing_extraction_is_exact.py verifies via AST. """ import json diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b680cd5a4a..f40348a9ae 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -7606,21 +7606,6 @@ async def _responses_stream( "delta": final_visible, }, ) - if full_reasoning and not full_text and not tool_call_state: - for event in _ensure_message_open(): - yield event - full_text = full_reasoning - api_monitor.set_reply(monitor_id, full_text) - yield _sse( - "response.output_text.delta", - { - "type": "response.output_text.delta", - "item_id": message_state["item_id"], - "output_index": message_state["output_index"], - "content_index": 0, - "delta": full_text, - }, - ) close_items: list[tuple[int, str, dict[str, Any]]] = [] if reasoning_state["opened"]: diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index f9d029babe..582fcfbe87 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -61,6 +61,7 @@ from models.inference import ( from routes.inference import ( _build_chat_request, _chat_tool_calls_to_responses_output, + _extract_responses_reasoning, _normalise_responses_input, _responses_tool_output_content, _responses_non_streaming, @@ -782,6 +783,15 @@ class TestResponsesNonStreamingAdapter: assert "" not in body["output"][1]["content"][0]["text"] assert "" not in body["output"][1]["content"][0]["text"] + def test_unclosed_think_block_extracts_as_reasoning(self): + reasoning, visible = _extract_responses_reasoning( + "partial plan", + parse_think_markers = True, + ) + + assert reasoning == "partial plan" + assert visible == "" + def test_monitor_records_translated_visible_text(self, monkeypatch): import routes.inference as inf_mod @@ -1318,7 +1328,7 @@ class TestResponsesStreamAdapter: assert entry["status"] == "completed" assert entry["reply"] == "tail" - def test_reasoning_only_fallback_updates_monitor(self, monkeypatch): + def test_reasoning_only_stream_does_not_update_visible_monitor_reply(self, monkeypatch): import routes.inference as inf_mod class FakeExtractor: @@ -1359,10 +1369,11 @@ class TestResponsesStreamAdapter: lines = asyncio.run(run()) - assert self._payloads(lines, "response.output_text.delta")[-1]["delta"] == "plan" + assert self._payloads(lines, "response.output_text.delta") == [] + assert self._payloads(lines, "response.reasoning_text.delta")[-1]["delta"] == "plan" [entry] = monitor.snapshot() assert entry["status"] == "completed" - assert entry["reply"] == "plan" + assert entry["reply"] == "" def test_reasoning_capable_gguf_stream_parses_think_tags_by_default(self, monkeypatch): chunks = [ @@ -1418,7 +1429,7 @@ class TestResponsesStreamAdapter: "show x tags" ) - def test_reasoning_only_streams_as_visible_message_text(self, monkeypatch): + def test_reasoning_only_stream_stays_out_of_visible_message_text(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": "plan"}}]}, {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, @@ -1436,14 +1447,34 @@ class TestResponsesStreamAdapter: reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") text_deltas = self._payloads(lines, "response.output_text.delta") assert "".join(event["delta"] for event in reasoning_deltas) == "plan" - assert "".join(event["delta"] for event in text_deltas) == "plan" + assert text_deltas == [] completed = self._payloads(lines, "response.completed")[0] - assert [item["type"] for item in completed["response"]["output"]] == [ - "reasoning", - "message", - ] + assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] + assert completed["response"]["output"][0]["content"][0]["text"] == "plan" + + def test_unclosed_think_stream_stays_out_of_visible_message_text(self, monkeypatch): + chunks = [ + {"choices": [{"delta": {"content": "plan"}}]}, + {"choices": [], "usage": {"prompt_tokens": 2, "completion_tokens": 3}}, + ] + self._install_stream_mock(monkeypatch, chunks) + payload = ResponsesRequest(input = "hi", stream = True, reasoning = {"effort": "high"}) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + return await self._collect(response) + + lines = asyncio.run(run()) + + reasoning_deltas = self._payloads(lines, "response.reasoning_text.delta") + text_deltas = self._payloads(lines, "response.output_text.delta") + assert "".join(event["delta"] for event in reasoning_deltas) == "plan" + assert text_deltas == [] + completed = self._payloads(lines, "response.completed")[0] + assert [item["type"] for item in completed["response"]["output"]] == ["reasoning"] assert completed["response"]["output"][0]["content"][0]["text"] == "plan" - assert completed["response"]["output"][1]["content"][0]["text"] == "plan" def test_structured_reasoning_content_streams_as_reasoning(self, monkeypatch): chunks = [ From 475ff786d8ec1df3f7ec8416ca667891fe254fc0 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 17:07:47 +0200 Subject: [PATCH 05/30] Address Gemma stream review comments --- studio/backend/core/tool_healing.py | 1 + studio/backend/routes/inference.py | 77 ++++++++++++++++++- .../tests/test_llama_route_timeouts.py | 68 ++++++++++++++++ studio/backend/tests/test_tool_xml_strip.py | 8 ++ 4 files changed, 152 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index c419d8e21c..e13c56a6c3 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -135,6 +135,7 @@ def _quote_gemma_object_keys(src: str) -> str: def _gemma_arguments_to_json(args_src: str) -> dict: + """Parse Gemma 4's native call:name{key:value} argument object.""" args_src = args_src.strip() if not args_src: return {} diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f40348a9ae..9c0937907d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -697,6 +697,53 @@ def _set_stream_response_read_timeout( pass +_STREAM_DISCONNECT_POLL_TIMEOUT_S = 0.25 + + +class _CompatSameTaskTimeout: + """Same-task timeout fallback for Python versions before asyncio.timeout.""" + + def __init__(self, timeout_s: float): + self.timeout_s = timeout_s + self._task = None + self._handle = None + self._timed_out = False + self._cancelling = 0 + + async def __aenter__(self): + self._task = asyncio.current_task() + if self._task is None: + return self + if hasattr(self._task, "cancelling"): + self._cancelling = self._task.cancelling() + loop = asyncio.get_running_loop() + self._handle = loop.call_later(max(self.timeout_s, 0), self._cancel_task) + return self + + async def __aexit__(self, exc_type, exc, tb): + if self._handle is not None: + self._handle.cancel() + if exc_type is not None and issubclass(exc_type, asyncio.CancelledError): + if self._timed_out: + if self._task is not None and hasattr(self._task, "uncancel"): + if self._task.uncancel() > self._cancelling: + return None + raise asyncio.TimeoutError from exc + return None + + def _cancel_task(self) -> None: + self._timed_out = True + if self._task is not None: + self._task.cancel() + + +def _same_task_timeout(timeout_s: float): + timeout_ctx = getattr(asyncio, "timeout", None) + if timeout_ctx is not None: + return timeout_ctx(timeout_s) + return _CompatSameTaskTimeout(timeout_s) + + async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: if cancel_event is not None and cancel_event.is_set(): return True @@ -785,13 +832,31 @@ async def _aiter_llama_stream_items( remaining_s = first_token_deadline - time.monotonic() if remaining_s <= 0: raise httpx.ReadTimeout("The model did not produce a first token in time.") + read_timeout_s = remaining_s + if request is not None: + read_timeout_s = min(read_timeout_s, _STREAM_DISCONNECT_POLL_TIMEOUT_S) if response is not None: - _set_stream_response_read_timeout(response, remaining_s) + _set_stream_response_read_timeout(response, read_timeout_s) # Keep httpx/httpcore's AnyIO cancel scope in this task. # asyncio.wait_for would drive __anext__ in a child task. - async with asyncio.timeout(remaining_s): + async with _same_task_timeout(remaining_s): item = await async_iter.__anext__() else: + if ( + request is not None + and response is not None + and post_first_item_read_timeout_s is not None + and last_item_at is not None + ): + stall_remaining_s = post_first_item_read_timeout_s - ( + time.monotonic() - last_item_at + ) + if stall_remaining_s <= 0: + raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") + _set_stream_response_read_timeout( + response, + min(stall_remaining_s, _STREAM_DISCONNECT_POLL_TIMEOUT_S), + ) item = await async_iter.__anext__() except asyncio.TimeoutError as exc: if waiting_first_item: @@ -805,6 +870,12 @@ async def _aiter_llama_stream_items( if now >= first_token_deadline: raise continue + if ( + request is not None + and post_first_item_read_timeout_s is not None + and now - last_item_at < post_first_item_read_timeout_s + ): + continue raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") if ( last_item_at is None @@ -1291,6 +1362,7 @@ _TOOL_XML_RE = _re.compile( r"<(?:tool_call|function=[\w-]+)>.*?(?:|\Z)" r"|<\|tool_call>.*?(?:|\Z)" r"|" + r"|" r"|\s*\Z", _re.DOTALL, ) @@ -7400,6 +7472,7 @@ async def _responses_stream( async for raw_line in _aiter_llama_stream_items( lines_iter, cancel_event = disconnect_event, + request = request, first_token_deadline = first_token_deadline, response = resp, ): diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index dc6e13a3de..b0954619dc 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -5,6 +5,7 @@ import asyncio import os import sys import time +import threading from types import SimpleNamespace _backend = os.path.join(os.path.dirname(__file__), "..") @@ -69,6 +70,73 @@ def test_stream_first_item_deadline_does_not_hop_tasks(): asyncio.run(_run()) +def test_stream_first_item_deadline_uses_compat_timeout_without_task_hop(monkeypatch): + monkeypatch.setattr(inf_mod.asyncio, "timeout", None, raising = False) + + async def _run(): + outer_task = asyncio.current_task() + seen_tasks = [] + + class _One: + def __init__(self): + self.done = False + + async def __anext__(self): + seen_tasks.append(asyncio.current_task()) + if self.done: + raise StopAsyncIteration + self.done = True + return "data: {}" + + out = [] + async for item in inf_mod._aiter_llama_stream_items( + _One(), + first_token_deadline = time.monotonic() + 1, + ): + out.append(item) + + assert out == ["data: {}"] + assert seen_tasks == [outer_task, outer_task] + + asyncio.run(_run()) + + +def test_stream_wait_polls_disconnect_without_background_watcher(): + async def _run(): + state = SimpleNamespace(disconnect_checks = 0) + cancel_event = threading.Event() + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + + class _Request: + async def is_disconnected(self): + state.disconnect_checks += 1 + return state.disconnect_checks >= 2 + + class _SlowFirstItem: + async def __anext__(self): + await asyncio.sleep(0.02) + raise inf_mod.httpx.ReadTimeout("poll") + + started = time.monotonic() + async for _ in inf_mod._aiter_llama_stream_items( + _SlowFirstItem(), + cancel_event = cancel_event, + request = _Request(), + response = response, + first_token_deadline = started + 1, + ): + raise AssertionError("stream should stop after disconnect") + + assert cancel_event.is_set() + assert state.disconnect_checks >= 2 + assert time.monotonic() - started < 0.5 + assert response.request.extensions["timeout"]["read"] <= ( + inf_mod._STREAM_DISCONNECT_POLL_TIMEOUT_S + ) + + asyncio.run(_run()) + + def test_preheader_send_cleanup_on_disconnect_and_cancel(): async def _run(cancel_parent): state = SimpleNamespace(disconnected = False, closed = False, cancelled = False) diff --git a/studio/backend/tests/test_tool_xml_strip.py b/studio/backend/tests/test_tool_xml_strip.py index 2ba3310fbe..c2dc1fe8db 100644 --- a/studio/backend/tests/test_tool_xml_strip.py +++ b/studio/backend/tests/test_tool_xml_strip.py @@ -125,6 +125,14 @@ def test_strips_orphan_closing_tag(): # Mid-string intentionally preserved (see preserve test). +def test_strips_gemma_native_orphan_closing_tag(): + cleaned = _TOOL_XML_RE.sub("", "Tool call drained.Visible tail.") + + assert "" not in cleaned + assert "Tool call drained." in cleaned + assert "Visible tail." in cleaned + + # ── Tail-only (PR #5735 follow-up) ─────────────────── From ede6a2bcee2a1ba26a19169b7d6bc65e38fdd74d Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 17:23:02 +0200 Subject: [PATCH 06/30] Avoid Responses stream task-group cleanup --- studio/backend/routes/inference.py | 15 ++++++++- .../tests/test_responses_tool_passthrough.py | 31 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 9c0937907d..f58f3aeaba 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -12,6 +12,7 @@ import uuid from pathlib import Path from fastapi import APIRouter, Depends, HTTPException, Request, status from fastapi.responses import StreamingResponse, JSONResponse, Response +from starlette.requests import ClientDisconnect from typing import Any, List, Optional, Union import json import httpx @@ -744,6 +745,18 @@ def _same_task_timeout(timeout_s: float): return _CompatSameTaskTimeout(timeout_s) +class _SameTaskStreamingResponse(StreamingResponse): + """StreamingResponse without Starlette's legacy AnyIO task-group wrapper.""" + + async def __call__(self, scope, receive, send) -> None: + try: + await self.stream_response(send) + except OSError: + raise ClientDisconnect() + if self.background is not None: + await self.background() + + async def _preheader_cancelled(cancel_event = None, request: Optional[Request] = None) -> bool: if cancel_event is not None and cancel_event.is_set(): return True @@ -7839,7 +7852,7 @@ async def _responses_stream( api_monitor.finish(monitor_id) yield _sse("response.completed", completed_response) - return StreamingResponse( + return _SameTaskStreamingResponse( event_generator(), media_type = "text/event-stream", headers = { diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 582fcfbe87..9a9cb024d6 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -59,6 +59,7 @@ from models.inference import ( ResponsesUsage, ) from routes.inference import ( + _SameTaskStreamingResponse, _build_chat_request, _chat_tool_calls_to_responses_output, _extract_responses_reasoning, @@ -1075,6 +1076,36 @@ class TestResponsesStreamAdapter: ), ) + def test_stream_response_avoids_legacy_receive_watcher(self, monkeypatch): + self._install_stream_mock( + monkeypatch, + [{"choices": [{"delta": {"content": "33"}}]}], + ) + payload = ResponsesRequest(input = "hi", stream = True) + messages = [ChatMessage(role = "user", content = "hi")] + + async def run(): + response = await _responses_stream(payload, messages, self._Request()) + assert isinstance(response, _SameTaskStreamingResponse) + + sent = [] + + async def receive(): + raise AssertionError("Responses streams poll disconnects in the generator") + + async def send(message): + sent.append(message) + + await response({"type": "http", "asgi": {"spec_version": "2.3"}}, receive, send) + return sent + + sent = asyncio.run(run()) + + assert sent[0]["type"] == "http.response.start" + body = b"".join(message.get("body", b"") for message in sent).decode() + assert "response.output_text.delta" in body + assert '"delta":"33"' in body.replace(" ", "") + def test_split_think_markers_stream_as_reasoning_and_visible_text(self, monkeypatch): chunks = [ {"choices": [{"delta": {"content": " Date: Fri, 19 Jun 2026 19:03:09 +0200 Subject: [PATCH 07/30] Harden OpenAI chat completion streams --- studio/backend/routes/inference.py | 77 +++++++++++++++++-- .../tests/test_openai_tool_passthrough.py | 63 +++++++++++++++ .../test_stream_cancel_registration_timing.py | 46 +++++++---- 3 files changed, 164 insertions(+), 22 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f58f3aeaba..ee2aeaebf0 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4748,7 +4748,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( audio_input_stream(), media_type = "text/event-stream", headers = { @@ -5239,7 +5239,7 @@ async def openai_chat_completions( pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( gguf_tool_stream(), media_type = "text/event-stream", headers = { @@ -5393,7 +5393,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( gguf_stream_chunks(), media_type = "text/event-stream", headers = { @@ -5842,7 +5842,7 @@ async def openai_chat_completions( _sf_tracker.__exit__(None, None, None) if payload.stream: - return StreamingResponse( + return _SameTaskStreamingResponse( sf_tool_stream(), media_type = "text/event-stream", headers = { @@ -6062,7 +6062,7 @@ async def openai_chat_completions( finally: _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( stream_chunks(), media_type = "text/event-stream", headers = { @@ -9547,7 +9547,7 @@ async def _openai_passthrough_stream( except Exception: pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( iter(()), media_type = "text/event-stream", headers = { @@ -9598,6 +9598,26 @@ async def _openai_passthrough_stream( _await_disconnect_then_close(request, resp, cancel_event) ) monitor_done = False + saw_finish_reason = False + saw_done = False + last_chunk_id = completion_id + last_chunk_model = model_name + last_chunk_created = int(time.time()) + + def _synthetic_finish_line() -> str: + chunk = ChatCompletionChunk( + id = last_chunk_id, + created = last_chunk_created, + model = last_chunk_model, + choices = [ + ChunkChoice( + delta = ChoiceDelta(), + finish_reason = "stop", + ) + ], + ) + return f"data: {chunk.model_dump_json(exclude_none = True)}" + try: lines_iter = resp.aiter_lines() async for raw_line in _aiter_llama_stream_items( @@ -9611,6 +9631,37 @@ async def _openai_passthrough_stream( continue if not raw_line.startswith("data: "): continue + data_text = raw_line[6:].strip() + if data_text == "[DONE]": + saw_done = True + if not saw_finish_reason 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 + yield raw_line + "\n\n" + monitor_done = True + break + try: + chunk_data = json.loads(data_text) + except json.JSONDecodeError: + chunk_data = None + if isinstance(chunk_data, dict): + if isinstance(chunk_data.get("id"), str): + last_chunk_id = chunk_data["id"] + if isinstance(chunk_data.get("model"), str): + last_chunk_model = chunk_data["model"] + if isinstance(chunk_data.get("created"), int): + last_chunk_created = chunk_data["created"] + choices = chunk_data.get("choices") + if isinstance(choices, list) and choices: + choice = choices[0] + if isinstance(choice, dict) and choice.get("finish_reason"): + saw_finish_reason = True # Honor parallel_tool_calls=false (best-effort): drop tool_call # deltas with index>=1 so only the first call streams. Only # lines carrying tool_calls are reparsed; everything else is @@ -9625,9 +9676,19 @@ async def _openai_passthrough_stream( # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" - if monitor_event == "done" or raw_line[6:].strip() == "[DONE]": + if monitor_event == "done": monitor_done = True break + if not saw_done and not saw_finish_reason 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" + yield "data: [DONE]\n\n" + monitor_done = True if not monitor_done: api_monitor.finish( monitor_id, @@ -9680,7 +9741,7 @@ async def _openai_passthrough_stream( pass _tracker.__exit__(None, None, None) - return StreamingResponse( + return _SameTaskStreamingResponse( _stream(), media_type = "text/event-stream", headers = { diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 6d2341b716..03f353a063 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -48,6 +48,7 @@ from routes.inference import ( _openai_passthrough_stream, _openai_stream_usage_chunk, _proxy_to_external_provider, + _SameTaskStreamingResponse, _set_or_prepend_system_message, openai_completions, openai_embeddings, @@ -2239,6 +2240,7 @@ class TestApiMonitorProviderAndCompletionStreams: "chatcmpl-test", monitor_id = monitor_id, ) + assert isinstance(response, _SameTaskStreamingResponse) iterator = response.body_iterator first = await anext(iterator) assert "hello" in first @@ -2256,6 +2258,67 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_synthesizes_missing_finish_reason(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield 'data: {"id":"upstream","created":123,"model":"gguf","choices":[{"index":0,"delta":{"content":"hello"}}]}' + yield "data: [DONE]" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + body = "".join(chunks) + + assert '"finish_reason":"stop"' in body.replace(" ", "") + assert "data: [DONE]" in body + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 2e6b5f14da..35a7258a65 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -146,24 +146,42 @@ def test_async_generators_cleanup_tracker_in_finally(): ) -def test_streaming_responses_have_no_background_task(): - top = None - for n in ast.walk(_TREE): - if isinstance(n, ast.AsyncFunctionDef) and n.name == "openai_chat_completions": - top = n - break - assert top is not None +def test_chat_completions_streams_avoid_starlette_task_group(): + top = _async_function("openai_chat_completions") + legacy_calls = [] + same_task_calls = 0 for sub in ast.walk(top): if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): continue - if sub.func.id != "StreamingResponse": + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 + assert not legacy_calls, ( + "Streaming /v1/chat/completions must use _SameTaskStreamingResponse, " + "not Starlette's legacy task-group StreamingResponse. Lines: " + f"{legacy_calls}" + ) + assert same_task_calls >= 5 + + +def test_openai_passthrough_stream_avoids_starlette_task_group(): + top = _async_function("_openai_passthrough_stream") + legacy_calls = [] + same_task_calls = 0 + for sub in ast.walk(top): + if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)): continue - kwargs = {kw.arg for kw in sub.keywords if kw.arg} - assert "background" not in kwargs, ( - "StreamingResponse in openai_chat_completions must not pass " - "`background=` -- cleanup now lives in the generator's finally " - "block; a BackgroundTask would be skipped on abrupt disconnect" - ) + if sub.func.id == "StreamingResponse": + legacy_calls.append(sub.lineno) + if sub.func.id == "_SameTaskStreamingResponse": + same_task_calls += 1 + assert not legacy_calls, ( + "OpenAI passthrough streams must use _SameTaskStreamingResponse, " + "not Starlette's legacy task-group StreamingResponse. Lines: " + f"{legacy_calls}" + ) + assert same_task_calls >= 2 def test_direct_llama_server_streams_install_disconnect_watcher(): From b53015be858a671c3c43cd3be2a68c254aca3045 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 19:40:29 +0200 Subject: [PATCH 08/30] Address OpenAI stream review issues --- studio/backend/routes/inference.py | 92 ++++++++++++++++--- .../tests/test_openai_tool_passthrough.py | 74 ++++++++++++++- .../tests/test_responses_tool_passthrough.py | 5 +- .../test_stream_cancel_registration_timing.py | 8 ++ 4 files changed, 163 insertions(+), 16 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ee2aeaebf0..f0cd4d3b3e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1303,6 +1303,16 @@ async def _await_disconnect_then_close(request, resp, cancel_event) -> None: return +async def _await_disconnect_then_cancel(request, cancel_event) -> None: + """Set ``cancel_event`` when a same-task local stream disconnects.""" + try: + while not await request.is_disconnected(): + await asyncio.sleep(0.1) + cancel_event.set() + except asyncio.CancelledError: + return + + # Centralized local/server tool nudge. Keep render_html guidance gated to turns # where the canvas tool is actually present in the tool schema; otherwise # small local models can hallucinate a missing tool call instead of following @@ -5095,6 +5105,9 @@ async def openai_chat_completions( async def gguf_tool_stream(): gen = None + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -5232,6 +5245,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if gen is not None: try: gen.close() @@ -5283,6 +5301,9 @@ async def openai_chat_completions( _tracker.__enter__() async def gguf_stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: # First chunk: role first_chunk = ChatCompletionChunk( @@ -5391,6 +5412,11 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( @@ -5710,6 +5736,9 @@ async def openai_chat_completions( async def sf_tool_stream(): gen = None + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -5834,6 +5863,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if gen is not None: try: gen.close() @@ -5952,6 +5986,9 @@ async def openai_chat_completions( _tracker.__enter__() async def stream_chunks(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -6060,6 +6097,11 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( @@ -7115,8 +7157,6 @@ async def _responses_non_streaming( # the model produced content, so clients expecting a pure tool-call turn # (finish_reason="tool_calls") don't see a spurious empty message item. output_items: list[dict] = [] - if reasoning_text and not text and not tool_calls: - text = reasoning_text if reasoning_text: output_items.append(_responses_reasoning_output_item(reasoning_text)) if text: @@ -7421,6 +7461,7 @@ async def _responses_stream( client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout()) resp = None lines_iter = None + disconnect_watcher = None disconnect_event = threading.Event() try: req = client.build_request( @@ -7482,6 +7523,9 @@ async def _responses_stream( return lines_iter = resp.aiter_lines() + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_close(request, resp, disconnect_event) + ) async for raw_line in _aiter_llama_stream_items( lines_iter, cancel_event = disconnect_event, @@ -7643,6 +7687,12 @@ async def _responses_stream( ) return finally: + if disconnect_watcher is not None: + disconnect_watcher.cancel() + try: + await disconnect_watcher + except (asyncio.CancelledError, Exception): + pass if lines_iter is not None: try: await lines_iter.aclose() @@ -9600,11 +9650,13 @@ async def _openai_passthrough_stream( monitor_done = False saw_finish_reason = False saw_done = False + saw_tool_call_delta = False last_chunk_id = completion_id last_chunk_model = model_name last_chunk_created = int(time.time()) def _synthetic_finish_line() -> str: + finish_reason = "tool_calls" if saw_tool_call_delta else "stop" chunk = ChatCompletionChunk( id = last_chunk_id, created = last_chunk_created, @@ -9612,7 +9664,7 @@ async def _openai_passthrough_stream( choices = [ ChunkChoice( delta = ChoiceDelta(), - finish_reason = "stop", + finish_reason = finish_reason, ) ], ) @@ -9643,9 +9695,21 @@ async def _openai_passthrough_stream( ) yield finish_line + "\n\n" saw_finish_reason = True + _monitor_openai_sse_line( + monitor_id, + raw_line, + llama_backend.context_length, + ) yield raw_line + "\n\n" monitor_done = True break + # Honor parallel_tool_calls=false (best-effort): drop tool_call + # deltas with index>=1 so only the first call streams. Only + # lines carrying tool_calls are reparsed; everything else is + # relayed byte-for-byte. + if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: + raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + data_text = raw_line[6:].strip() try: chunk_data = json.loads(data_text) except json.JSONDecodeError: @@ -9660,14 +9724,12 @@ async def _openai_passthrough_stream( choices = chunk_data.get("choices") if isinstance(choices, list) and choices: choice = choices[0] - if isinstance(choice, dict) and choice.get("finish_reason"): - saw_finish_reason = True - # Honor parallel_tool_calls=false (best-effort): drop tool_call - # deltas with index>=1 so only the first call streams. Only - # lines carrying tool_calls are reparsed; everything else is - # relayed byte-for-byte. - if payload.parallel_tool_calls is False and '"tool_calls"' in raw_line: - raw_line = _cap_parallel_tool_calls_sse_line(raw_line) + if isinstance(choice, dict): + if choice.get("finish_reason"): + saw_finish_reason = True + delta = choice.get("delta") + if isinstance(delta, dict) and delta.get("tool_calls"): + saw_tool_call_delta = True monitor_event = _monitor_openai_sse_line( monitor_id, raw_line, @@ -9687,7 +9749,13 @@ async def _openai_passthrough_stream( llama_backend.context_length, ) yield finish_line + "\n\n" - yield "data: [DONE]\n\n" + done_line = "data: [DONE]" + _monitor_openai_sse_line( + monitor_id, + done_line, + llama_backend.context_length, + ) + yield done_line + "\n\n" monitor_done = True if not monitor_done: api_monitor.finish( diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 03f353a063..6d5d9badc1 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -2316,6 +2316,75 @@ class TestApiMonitorProviderAndCompletionStreams: assert '"finish_reason":"stop"' in body.replace(" ", "") assert "data: [DONE]" in body + assert monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_synthesizes_tool_call_finish_reason(self, monkeypatch): + async def _run(): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + yield ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' + '"id":"call_1","type":"function","function":{"name":"lookup",' + '"arguments":"{}"}}]}}]}' + ) + yield "data: [DONE]" + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + body = "".join([chunk async for chunk in response.body_iterator]) + compact = body.replace(" ", "") + + assert '"finish_reason":"tool_calls"' in compact + assert '"finish_reason":"stop"' not in compact + assert "data: [DONE]" in body + assert monitor.active_count() == 0 asyncio.run(_run()) @@ -2434,7 +2503,10 @@ class TestApiMonitorProviderAndCompletionStreams: async for chunk in response.body_iterator: chunks.append(chunk) - assert chunks == ['data: {"choices":[{"delta":{"content":"hello"}}]}\n\n'] + assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' + compact = "".join(chunks).replace(" ", "") + assert '"finish_reason":"stop"' in compact + assert chunks[-1] == "data: [DONE]\n\n" [entry] = monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == "hello" diff --git a/studio/backend/tests/test_responses_tool_passthrough.py b/studio/backend/tests/test_responses_tool_passthrough.py index 9a9cb024d6..4147746b54 100644 --- a/studio/backend/tests/test_responses_tool_passthrough.py +++ b/studio/backend/tests/test_responses_tool_passthrough.py @@ -992,7 +992,7 @@ class TestResponsesNonStreamingAdapter: assert [item["type"] for item in body["output"]] == ["message"] assert body["output"][0]["content"][0]["text"] == "33" - def test_reasoning_only_is_also_visible_message_text(self, monkeypatch): + def test_reasoning_only_stays_out_of_visible_message_text(self, monkeypatch): payload = ResponsesRequest(input = "hi", reasoning = {"effort": "high"}) body = self._run_with_message( monkeypatch, @@ -1000,9 +1000,8 @@ class TestResponsesNonStreamingAdapter: payload = payload, ) - assert [item["type"] for item in body["output"]] == ["reasoning", "message"] + assert [item["type"] for item in body["output"]] == ["reasoning"] assert body["output"][0]["content"][0]["text"] == "plan" - assert body["output"][1]["content"][0]["text"] == "plan" # ===================================================================== diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 35a7258a65..9df54ba23c 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -184,6 +184,14 @@ def test_openai_passthrough_stream_avoids_starlette_task_group(): assert same_task_calls >= 2 +def test_local_chat_streams_install_same_task_disconnect_watcher(): + top = _async_function("openai_chat_completions") + assert _calls_name(top, "_await_disconnect_then_cancel"), ( + "Local same-task streams must watch request disconnects themselves; " + "do not restore Starlette's task-group StreamingResponse for this." + ) + + def test_direct_llama_server_streams_install_disconnect_watcher(): required = { "openai_completions", From 5ba19c29778e67594ad5bb524c51cc138b38b71d Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 20:09:13 +0200 Subject: [PATCH 09/30] Clean up Studio OpenAI stream helpers --- .../core/inference/tool_call_parser.py | 204 +------- studio/backend/core/tool_healing.py | 129 +++-- studio/backend/routes/inference.py | 32 +- .../tests/test_openai_tool_passthrough.py | 493 ++++++------------ 4 files changed, 274 insertions(+), 584 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 78e8f14103..4f5310567f 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -7,21 +7,24 @@ Tolerates missing closing tags in either ``{json}`` or ``v...`` shape. """ -import json -import re - from core.tool_healing import ( _TC_END_TAG_RE, _TC_FUNC_CLOSE_RE, _TC_FUNC_START_RE, + _TC_GEMMA_END_TAG_RE, _TC_GEMMA_START_RE, _TC_JSON_START_RE, _TC_PARAM_CLOSE_RE, _TC_PARAM_START_RE, _TOOL_ALL_PATS, _TOOL_CLOSED_PATS, + _FUNC_CLOSE_TAG, + _PARAM_CLOSE_TAG, _balanced_brace_end, _gemma_arguments_to_json, + _inside_open_parameter, + parse_tool_calls_from_text, + strip_tool_call_markup as strip_tool_markup, ) @@ -75,201 +78,6 @@ RAG_SEARCH_CAP_NUDGE = ( ) -_TC_GEMMA_END_TAG_RE = re.compile(r"") -_PARAM_CLOSE_TAG = "" -_FUNC_CLOSE_TAG = "" - - -def _inside_open_parameter(content: str, pos: int) -> bool: - """Return True when ``pos`` falls inside an unclosed parameter value.""" - last_param_start = -1 - for match in _TC_PARAM_START_RE.finditer(content, 0, pos): - last_param_start = match.start() - if last_param_start < 0: - return False - last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) - last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) - return last_param_start > max(last_param_close, last_func_close) - - -def strip_tool_markup(text: str, *, final: bool = False) -> str: - """Strip tool-call XML from streamed text. - - ``final=False`` only removes closed pairs (used during streaming so - in-progress XML stays buffered). ``final=True`` also removes a - trailing unclosed run and trims the result. - """ - pats = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in pats: - text = pat.sub("", text) - return text.strip() if final else text - - -def parse_tool_calls_from_text( - content: str, - *, - id_offset: int = 0, - allow_incomplete: bool = True, -) -> list[dict]: - """Parse OpenAI-format ``tool_calls`` from model text. - - Returns a list of ``{"id", "type", "function": {"name", "arguments"}}`` - dicts. ``arguments`` is always a JSON string so callers can hand it - straight back into an OpenAI-style response. - - Handles three shapes: - - - JSON inside ```` tags: - ``{"name":"web_search","arguments":{"query":"..."}}`` - - Gemma 4 native call blocks: - ``<|tool_call>call:web_search{query:"..." }`` - - XML-style function blocks: - ``v`` - - ``allow_incomplete=True`` keeps the historical healing behavior for - missing closing tags. ``allow_incomplete=False`` accepts only - well-formed wrappers so disabled Auto-Heal can still parse valid - local tool protocol without repairing truncated output. - """ - tool_calls: list[dict] = [] - - # Pattern 1: {json}. Balanced-brace scan, skipping braces in - # JSON strings. - for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # opening { - i = _balanced_brace_end(content, brace_start) - if i < 0: - continue - if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_END_TAG_RE.match(tail_after_json) is None: - continue - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 1b: Gemma 4 native call block: - # <|tool_call>call:terminal{command:"ls"} - for m in _TC_GEMMA_START_RE.finditer(content): - brace_start = m.end() - 1 - i = _balanced_brace_end(content, brace_start) - if i < 0: - continue - if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_GEMMA_END_TAG_RE.match(tail_after_json) is None: - continue - try: - tool_calls.append( - { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": m.group(1), - "arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])), - }, - } - ) - except (json.JSONDecodeError, ValueError): - pass - - # Pattern 2: v... -- closing tags optional; - # isn't a body boundary since code values can contain it. - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - if not allow_incomplete: - # Bound the body at the closing tag rather than - # the end of the response, so a complete call followed by - # trailing prose is still accepted (matching the JSON-style - # path, which already tolerates trailing text). - # rfind picks the last , so a literal - # inside a code parameter value stays in the body. - close_idx = body.rfind(_FUNC_CLOSE_TAG) - if close_idx < 0: - continue - body = body[:close_idx] - else: - body = _TC_FUNC_CLOSE_RE.sub("", body) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - # Single param: take everything to body end so an embedded - # in code strings is preserved. - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = val.strip() - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = val.strip() - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - - return tool_calls - - def has_tool_signal(text: str) -> bool: """Return True if ``text`` contains any tool-call XML signal.""" return any(s in text for s in TOOL_XML_SIGNALS) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e13c56a6c3..855755e483 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -29,10 +29,13 @@ _TC_JSON_START_RE = re.compile(r"\s*\{") _TC_GEMMA_START_RE = re.compile(r"<\|tool_call>call:([\w-]+)\s*\{") _TC_FUNC_START_RE = re.compile(r"\s*") _TC_END_TAG_RE = re.compile(r"") +_TC_GEMMA_END_TAG_RE = re.compile(r"") _TC_FUNC_CLOSE_RE = re.compile(r"\s*\s*$") _TC_PARAM_START_RE = re.compile(r"\s*") _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' +_PARAM_CLOSE_TAG = "" +_FUNC_CLOSE_TAG = "" def _balanced_brace_end(content: str, brace_start: int) -> int: @@ -145,52 +148,72 @@ def _gemma_arguments_to_json(args_src: str) -> dict: return json.loads(src) -def parse_tool_calls_from_text(content: str) -> list[dict]: - """ - Parse tool calls from XML markup in content text. +def _inside_open_parameter(content: str, pos: int) -> bool: + """Return True when ``pos`` falls inside an unclosed parameter value.""" + last_param_start = -1 + for match in _TC_PARAM_START_RE.finditer(content, 0, pos): + last_param_start = match.start() + if last_param_start < 0: + return False + last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos) + last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos) + return last_param_start > max(last_param_close, last_func_close) + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, +) -> list[dict]: + """Parse OpenAI-format tool calls from model text. Handles formats like: {"name":"web_search","arguments":{"query":"..."}} <|tool_call>call:web_search{query:"..."} ... - Closing tags (, , ) are all - optional since models frequently omit them. """ - tool_calls = [] + tool_calls: list[dict] = [] - # Pattern 1: JSON inside tags. Balanced-brace extraction that - # skips braces inside JSON strings. for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 # position of the opening { + brace_start = m.end() - 1 i = _balanced_brace_end(content, brace_start) - if i >= 0: - json_str = content[brace_start : i + 1] - try: - obj = json.loads(json_str) - tc = { - "id": f"call_{len(tool_calls)}", - "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, - } - if isinstance(tc["function"]["arguments"], dict): - tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) - tool_calls.append(tc) - except (json.JSONDecodeError, ValueError): - pass + if i < 0: + continue + if not allow_incomplete: + tail_after_json = content[i + 1 :].lstrip() + if _TC_END_TAG_RE.match(tail_after_json) is None: + continue + json_str = content[brace_start : i + 1] + try: + obj = json.loads(json_str) + tc = { + "id": f"call_{id_offset + len(tool_calls)}", + "type": "function", + "function": { + "name": obj.get("name", ""), + "arguments": obj.get("arguments", {}), + }, + } + if isinstance(tc["function"]["arguments"], dict): + tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"]) + tool_calls.append(tc) + except (json.JSONDecodeError, ValueError): + pass - # Pattern 1b: Gemma 4 native <|tool_call>call:name{key:value}. for m in _TC_GEMMA_START_RE.finditer(content): brace_start = m.end() - 1 i = _balanced_brace_end(content, brace_start) if i < 0: continue + if not allow_incomplete: + tail_after_json = content[i + 1 :].lstrip() + if _TC_GEMMA_END_TAG_RE.match(tail_after_json) is None: + continue try: tool_calls.append( { - "id": f"call_{len(tool_calls)}", + "id": f"call_{id_offset + len(tool_calls)}", "type": "function", "function": { "name": m.group(1), @@ -201,17 +224,15 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: except (json.JSONDecodeError, ValueError): pass - # Pattern 2: XML-style value - # All closing tags optional; models frequently omit them. if not tool_calls: - # Step 1: Find positions and extract bodies. Use only - # or the next - # can appear in code values); trim a trailing afterwards. - func_starts = list(_TC_FUNC_START_RE.finditer(content)) + func_starts = [ + fm + for fm in _TC_FUNC_START_RE.finditer(content) + if not _inside_open_parameter(content, fm.start()) + ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) body_start = fm.end() - # Boundaries: next next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) end_tag = _TC_END_TAG_RE.search(content[body_start:]) if end_tag: @@ -220,36 +241,52 @@ def parse_tool_calls_from_text(content: str) -> list[dict]: body_end = len(content) body_end = min(body_end, next_func) body = content[body_start:body_end] - body = _TC_FUNC_CLOSE_RE.sub("", body) # trim closing + if not allow_incomplete: + close_idx = body.rfind(_FUNC_CLOSE_TAG) + if close_idx < 0: + continue + body = body[:close_idx] + else: + body = _TC_FUNC_CLOSE_RE.sub("", body) - # Step 2: Extract parameters from body. For single-parameter - # functions, use body end as the only boundary to avoid matching - # inside code strings. - arguments = {} + arguments: dict = {} param_starts = list(_TC_PARAM_START_RE.finditer(body)) if len(param_starts) == 1: - # Value is everything after the tag to end of body, less a - # trailing . pm = param_starts[0] val = body[pm.end() :] - val = _TC_PARAM_CLOSE_RE.sub("", val) + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + continue + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) arguments[pm.group(1)] = val.strip() else: + valid_params = True for pidx, pm in enumerate(param_starts): param_name = pm.group(1) val_start = pm.end() - # Value ends at next + if not allow_incomplete: + stripped_val = val.rstrip() + if not stripped_val.endswith(_PARAM_CLOSE_TAG): + valid_params = False + break + val = stripped_val[: -len(_PARAM_CLOSE_TAG)] + else: + val = _TC_PARAM_CLOSE_RE.sub("", val) arguments[param_name] = val.strip() + if not valid_params: + continue tc = { - "id": f"call_{len(tool_calls)}", + "id": f"call_{id_offset + len(tool_calls)}", "type": "function", "function": { "name": func_name, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f0cd4d3b3e..5fd91bab56 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1313,6 +1313,14 @@ async def _await_disconnect_then_cancel(request, cancel_event) -> None: return +async def _stop_local_disconnect_cancel_watcher(watcher) -> None: + watcher.cancel() + try: + await watcher + except (asyncio.CancelledError, Exception): + pass + + # Centralized local/server tool nudge. Keep render_html guidance gated to turns # where the canvas tool is actually present in the tool schema; otherwise # small local models can hallucinate a missing tool call instead of following @@ -5245,11 +5253,7 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: - disconnect_watcher.cancel() - try: - await disconnect_watcher - except (asyncio.CancelledError, Exception): - pass + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: try: gen.close() @@ -5412,11 +5416,7 @@ async def openai_chat_completions( error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: - disconnect_watcher.cancel() - try: - await disconnect_watcher - except (asyncio.CancelledError, Exception): - pass + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( @@ -5863,11 +5863,7 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: - disconnect_watcher.cancel() - try: - await disconnect_watcher - except (asyncio.CancelledError, Exception): - pass + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if gen is not None: try: gen.close() @@ -6097,11 +6093,7 @@ async def openai_chat_completions( } yield f"data: {json.dumps(error_chunk)}\n\n" finally: - disconnect_watcher.cancel() - try: - await disconnect_watcher - except (asyncio.CancelledError, Exception): - pass + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 6d5d9badc1..190bfd6a36 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1264,6 +1264,61 @@ class TestGgufVisionToolRouting: pass return payloads + def _run_gguf_case( + self, + monkeypatch, + *, + generate = None, + tool_generate = None, + payload_kwargs = None, + backend_kwargs = None, + ): + import routes.inference as inf_mod + + reset_tool_policy() + + def _plain(**_kwargs): + raise AssertionError("plain GGUF path should not be used") + + backend_data = { + "is_loaded": True, + "is_vision": False, + "supports_tools": tool_generate is not None, + "supports_reasoning": True, + "reasoning_always_on": True, + "_is_audio": False, + "model_identifier": "test-gguf", + "context_length": 4096, + "generate_chat_completion": generate or _plain, + } + if tool_generate is not None: + backend_data["generate_chat_completion_with_tools"] = tool_generate + if backend_kwargs: + backend_data.update(backend_kwargs) + backend = SimpleNamespace(**backend_data) + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) + + request_data = { + "model": "default", + "messages": [{"role": "user", "content": "hi"}], + } + if payload_kwargs: + request_data.update(payload_kwargs) + payload = ChatCompletionRequest(**request_data) + response = self._drive( + openai_chat_completions(payload, request = self._Request(), current_subject = "test") + ) + result = SimpleNamespace(response = response, monitor = monitor, backend = backend) + if request_data.get("stream"): + result.chunks = self._consume_response(response) + result.payloads = self._sse_payloads(result.chunks) + else: + result.body = json.loads(response.body) + return result + def test_image_request_with_enabled_tools_enters_gguf_tool_loop(self, monkeypatch): import routes.inference as inf_mod @@ -1410,10 +1465,6 @@ class TestGgufVisionToolRouting: assert monitor.active_count() == 0 def test_standard_gguf_stream_splits_reasoning_content(self, monkeypatch): - import routes.inference as inf_mod - - reset_tool_policy() - def _generate(**_kwargs): yield "plan" @@ -1425,44 +1476,20 @@ class TestGgufVisionToolRouting: "finish_reason": "stop", } - backend = SimpleNamespace( - is_loaded = True, - is_vision = False, - supports_tools = False, - supports_reasoning = True, - reasoning_always_on = True, - _is_audio = False, - model_identifier = "test-gguf", - context_length = 4096, - generate_chat_completion = _generate, + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True}, ) - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - - payload = ChatCompletionRequest( - model = "default", - stream = True, - messages = [{"role": "user", "content": "hi"}], - ) - - response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") - ) - payloads = self._sse_payloads(self._consume_response(response)) - deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" assert all("" not in d.get("content", "") for d in deltas) - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" def test_reasoning_capable_gguf_stream_splits_reasoning_by_default(self, monkeypatch): - import routes.inference as inf_mod - - reset_tool_policy() - def _generate(**_kwargs): yield "planvisible" yield { @@ -1471,43 +1498,20 @@ class TestGgufVisionToolRouting: "finish_reason": "stop", } - backend = SimpleNamespace( - is_loaded = True, - is_vision = False, - supports_tools = False, - supports_reasoning = True, - reasoning_always_on = False, - _is_audio = False, - model_identifier = "test-gguf", - context_length = 4096, - generate_chat_completion = _generate, + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True}, + backend_kwargs = {"reasoning_always_on": False}, ) - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - - payload = ChatCompletionRequest( - model = "default", - stream = True, - messages = [{"role": "user", "content": "hi"}], - ) - - response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") - ) - payloads = self._sse_payloads(self._consume_response(response)) - deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" assert "".join(d.get("content", "") for d in deltas) == "visible" - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" def test_reasoning_capable_gguf_stream_sanitizes_think_tags_when_disabled(self, monkeypatch): - import routes.inference as inf_mod - - reset_tool_policy() - def _generate(**_kwargs): yield "leakedvisible" yield { @@ -1516,48 +1520,21 @@ class TestGgufVisionToolRouting: "finish_reason": "stop", } - backend = SimpleNamespace( - is_loaded = True, - is_vision = False, - supports_tools = False, - supports_reasoning = True, - reasoning_always_on = False, - _is_audio = False, - model_identifier = "test-gguf", - context_length = 4096, - generate_chat_completion = _generate, + result = self._run_gguf_case( + monkeypatch, + generate = _generate, + payload_kwargs = {"stream": True, "enable_thinking": False}, + backend_kwargs = {"reasoning_always_on": False}, ) - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - - payload = ChatCompletionRequest( - model = "default", - stream = True, - enable_thinking = False, - messages = [{"role": "user", "content": "hi"}], - ) - - response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") - ) - payloads = self._sse_payloads(self._consume_response(response)) - deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] assert "".join(d.get("reasoning_content", "") for d in deltas) == "leaked" assert "".join(d.get("content", "") for d in deltas) == "visible" assert all("" not in d.get("content", "") for d in deltas) - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" def test_gguf_tool_stream_splits_reasoning_and_strips_gemma_tool_marker(self, monkeypatch): - import routes.inference as inf_mod - - reset_tool_policy() - - def _plain(**_kwargs): - raise AssertionError("plain GGUF path should not be used") - def _tools(**_kwargs): yield { "type": "content", @@ -1569,48 +1546,26 @@ class TestGgufVisionToolRouting: "finish_reason": "stop", } - backend = SimpleNamespace( - is_loaded = True, - is_vision = False, - supports_tools = True, - supports_reasoning = True, - reasoning_always_on = True, - _is_audio = False, - model_identifier = "test-gguf", - context_length = 4096, - generate_chat_completion = _plain, - generate_chat_completion_with_tools = _tools, + result = self._run_gguf_case( + monkeypatch, + tool_generate = _tools, + payload_kwargs = { + "stream": True, + "enable_tools": True, + "enabled_tools": ["terminal"], + "messages": [{"role": "user", "content": "list files"}], + }, ) - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - - payload = ChatCompletionRequest( - model = "default", - stream = True, - enable_tools = True, - enabled_tools = ["terminal"], - messages = [{"role": "user", "content": "list files"}], - ) - - response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") - ) - payloads = self._sse_payloads(self._consume_response(response)) - deltas = [p["choices"][0].get("delta", {}) for p in payloads if p.get("choices")] + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] assert "".join(d.get("reasoning_content", "") for d in deltas) == "plan" combined_content = "".join(d.get("content", "") for d in deltas) assert combined_content == "visible " assert "<|tool_call>" not in combined_content - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["reply"] == "visible " def test_non_streaming_gguf_splits_reasoning_content(self, monkeypatch): - import routes.inference as inf_mod - - reset_tool_policy() - def _generate(**_kwargs): yield "planvisible" yield { @@ -1619,35 +1574,13 @@ class TestGgufVisionToolRouting: "finish_reason": "stop", } - backend = SimpleNamespace( - is_loaded = True, - is_vision = False, - supports_tools = False, - supports_reasoning = True, - reasoning_always_on = True, - _is_audio = False, - model_identifier = "test-gguf", - context_length = 4096, - generate_chat_completion = _generate, - ) - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "get_llama_cpp_backend", lambda: backend) - - payload = ChatCompletionRequest( - model = "default", - messages = [{"role": "user", "content": "hi"}], - ) - - response = self._drive( - openai_chat_completions(payload, request = self._Request(), current_subject = "test") - ) - body = json.loads(response.body) + result = self._run_gguf_case(monkeypatch, generate = _generate) + body = result.body message = body["choices"][0]["message"] assert message["content"] == "visible" assert message["reasoning_content"] == "plan" - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["reply"] == "visible" def test_non_streaming_gguf_n_records_all_monitor_replies(self, monkeypatch): @@ -1812,6 +1745,61 @@ class TestApiMonitorProviderAndCompletionStreams: async def is_disconnected(self): return False + async def _run_passthrough_stream(self, monkeypatch, lines): + import routes.inference as inf_mod + + class Request: + async def is_disconnected(self): + return False + + async def fake_send(*_args, **_kwargs): + return httpx.Response(200, content = b"") + + async def fake_items(*_args, **_kwargs): + for line in lines: + yield line + + monitor = ApiMonitor(max_entries = 3) + monkeypatch.setattr(inf_mod, "api_monitor", monitor) + monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) + monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) + monitor_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "gguf", + prompt = "hi", + ) + payload = ChatCompletionRequest( + model = "default", + messages = [ChatMessage(role = "user", content = "hi")], + stream = True, + tools = [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + ) + + response = await _openai_passthrough_stream( + Request(), + threading.Event(), + SimpleNamespace( + base_url = "http://llama.test", + context_length = 4096, + _request_reasoning_kwargs = lambda *_args, **_kwargs: None, + ), + payload, + "gguf", + "chatcmpl-test", + monitor_id = monitor_id, + ) + chunks = [chunk async for chunk in response.body_iterator] + return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor) + def test_external_non_streaming_json_updates_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod @@ -2260,131 +2248,44 @@ class TestApiMonitorProviderAndCompletionStreams: def test_passthrough_stream_synthesizes_missing_finish_reason(self, monkeypatch): async def _run(): - import routes.inference as inf_mod - - class Request: - async def is_disconnected(self): - return False - - async def fake_send(*_args, **_kwargs): - return httpx.Response(200, content = b"") - - async def fake_items(*_args, **_kwargs): - yield 'data: {"id":"upstream","created":123,"model":"gguf","choices":[{"index":0,"delta":{"content":"hello"}}]}' - yield "data: [DONE]" - - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) - monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) - monitor_id = monitor.start( - endpoint = "/v1/chat/completions", - method = "POST", - model = "gguf", - prompt = "hi", - ) - payload = ChatCompletionRequest( - model = "default", - messages = [ChatMessage(role = "user", content = "hi")], - stream = True, - tools = [ - { - "type": "function", - "function": { - "name": "lookup", - "parameters": {"type": "object", "properties": {}}, - }, - } + result = await self._run_passthrough_stream( + monkeypatch, + [ + ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"content":"hello"}}]}' + ), + "data: [DONE]", ], ) - - response = await _openai_passthrough_stream( - Request(), - threading.Event(), - SimpleNamespace( - base_url = "http://llama.test", - context_length = 4096, - _request_reasoning_kwargs = lambda *_args, **_kwargs: None, - ), - payload, - "gguf", - "chatcmpl-test", - monitor_id = monitor_id, - ) - chunks = [chunk async for chunk in response.body_iterator] - body = "".join(chunks) + body = result.body assert '"finish_reason":"stop"' in body.replace(" ", "") assert "data: [DONE]" in body - assert monitor.active_count() == 0 + assert result.monitor.active_count() == 0 asyncio.run(_run()) def test_passthrough_stream_synthesizes_tool_call_finish_reason(self, monkeypatch): async def _run(): - import routes.inference as inf_mod - - class Request: - async def is_disconnected(self): - return False - - async def fake_send(*_args, **_kwargs): - return httpx.Response(200, content = b"") - - async def fake_items(*_args, **_kwargs): - yield ( - 'data: {"id":"upstream","created":123,"model":"gguf",' - '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' - '"id":"call_1","type":"function","function":{"name":"lookup",' - '"arguments":"{}"}}]}}]}' - ) - yield "data: [DONE]" - - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) - monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) - monitor_id = monitor.start( - endpoint = "/v1/chat/completions", - method = "POST", - model = "gguf", - prompt = "hi", - ) - payload = ChatCompletionRequest( - model = "default", - messages = [ChatMessage(role = "user", content = "hi")], - stream = True, - tools = [ - { - "type": "function", - "function": { - "name": "lookup", - "parameters": {"type": "object", "properties": {}}, - }, - } + result = await self._run_passthrough_stream( + monkeypatch, + [ + ( + 'data: {"id":"upstream","created":123,"model":"gguf",' + '"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,' + '"id":"call_1","type":"function","function":{"name":"lookup",' + '"arguments":"{}"}}]}}]}' + ), + "data: [DONE]", ], ) - - response = await _openai_passthrough_stream( - Request(), - threading.Event(), - SimpleNamespace( - base_url = "http://llama.test", - context_length = 4096, - _request_reasoning_kwargs = lambda *_args, **_kwargs: None, - ), - payload, - "gguf", - "chatcmpl-test", - monitor_id = monitor_id, - ) - body = "".join([chunk async for chunk in response.body_iterator]) - compact = body.replace(" ", "") + compact = result.body.replace(" ", "") assert '"finish_reason":"tool_calls"' in compact assert '"finish_reason":"stop"' not in compact - assert "data: [DONE]" in body - assert monitor.active_count() == 0 + assert "data: [DONE]" in result.body + assert result.monitor.active_count() == 0 asyncio.run(_run()) @@ -2449,68 +2350,20 @@ class TestApiMonitorProviderAndCompletionStreams: def test_passthrough_clean_eof_finalizes_monitor(self, monkeypatch): async def _run(): - import routes.inference as inf_mod - - class Request: - async def is_disconnected(self): - return False - - async def fake_send(*_args, **_kwargs): - return httpx.Response(200, content = b"") - - async def fake_items(*_args, **_kwargs): - yield 'data: {"choices":[{"delta":{"content":"hello"}}]}' - - monitor = ApiMonitor(max_entries = 3) - monkeypatch.setattr(inf_mod, "api_monitor", monitor) - monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send) - monkeypatch.setattr(inf_mod, "_aiter_llama_stream_items", fake_items) - monitor_id = monitor.start( - endpoint = "/v1/chat/completions", - method = "POST", - model = "gguf", - prompt = "hi", + result = await self._run_passthrough_stream( + monkeypatch, + ['data: {"choices":[{"delta":{"content":"hello"}}]}'], ) - payload = ChatCompletionRequest( - model = "default", - messages = [ChatMessage(role = "user", content = "hi")], - stream = True, - tools = [ - { - "type": "function", - "function": { - "name": "lookup", - "parameters": {"type": "object", "properties": {}}, - }, - } - ], - ) - - response = await _openai_passthrough_stream( - Request(), - threading.Event(), - SimpleNamespace( - base_url = "http://llama.test", - context_length = 4096, - _request_reasoning_kwargs = lambda *_args, **_kwargs: None, - ), - payload, - "gguf", - "chatcmpl-test", - monitor_id = monitor_id, - ) - chunks = [] - async for chunk in response.body_iterator: - chunks.append(chunk) + chunks = result.chunks assert chunks[0] == 'data: {"choices":[{"delta":{"content":"hello"}}]}\n\n' compact = "".join(chunks).replace(" ", "") assert '"finish_reason":"stop"' in compact assert chunks[-1] == "data: [DONE]\n\n" - [entry] = monitor.snapshot() + [entry] = result.monitor.snapshot() assert entry["status"] == "completed" assert entry["reply"] == "hello" - assert monitor.active_count() == 0 + assert result.monitor.active_count() == 0 asyncio.run(_run()) From 27e0228fdfff5ae26cde9664d306e3914b2da52b Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 20:40:29 +0200 Subject: [PATCH 10/30] Fix Studio passthrough cold stream timeout --- studio/backend/routes/inference.py | 10 +--- .../tests/test_llama_route_timeouts.py | 51 +++++++++++++------ 2 files changed, 38 insertions(+), 23 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5fd91bab56..da8f98b1f5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -845,11 +845,8 @@ async def _aiter_llama_stream_items( remaining_s = first_token_deadline - time.monotonic() if remaining_s <= 0: raise httpx.ReadTimeout("The model did not produce a first token in time.") - read_timeout_s = remaining_s - if request is not None: - read_timeout_s = min(read_timeout_s, _STREAM_DISCONNECT_POLL_TIMEOUT_S) if response is not None: - _set_stream_response_read_timeout(response, read_timeout_s) + _set_stream_response_read_timeout(response, remaining_s) # Keep httpx/httpcore's AnyIO cancel scope in this task. # asyncio.wait_for would drive __anext__ in a child task. async with _same_task_timeout(remaining_s): @@ -866,10 +863,7 @@ async def _aiter_llama_stream_items( ) if stall_remaining_s <= 0: raise httpx.ReadTimeout("The model stopped producing tokens mid-response.") - _set_stream_response_read_timeout( - response, - min(stall_remaining_s, _STREAM_DISCONNECT_POLL_TIMEOUT_S), - ) + _set_stream_response_read_timeout(response, stall_remaining_s) item = await async_iter.__anext__() except asyncio.TimeoutError as exc: if waiting_first_item: diff --git a/studio/backend/tests/test_llama_route_timeouts.py b/studio/backend/tests/test_llama_route_timeouts.py index b0954619dc..24866bd03e 100644 --- a/studio/backend/tests/test_llama_route_timeouts.py +++ b/studio/backend/tests/test_llama_route_timeouts.py @@ -101,38 +101,59 @@ def test_stream_first_item_deadline_uses_compat_timeout_without_task_hop(monkeyp asyncio.run(_run()) -def test_stream_wait_polls_disconnect_without_background_watcher(): +def test_stream_wait_stops_on_known_disconnect_before_read(): async def _run(): state = SimpleNamespace(disconnect_checks = 0) cancel_event = threading.Event() - response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) class _Request: async def is_disconnected(self): state.disconnect_checks += 1 - return state.disconnect_checks >= 2 + return True - class _SlowFirstItem: + class _Unread: async def __anext__(self): - await asyncio.sleep(0.02) - raise inf_mod.httpx.ReadTimeout("poll") + raise AssertionError("stream should stop before reading upstream") - started = time.monotonic() async for _ in inf_mod._aiter_llama_stream_items( - _SlowFirstItem(), + _Unread(), cancel_event = cancel_event, request = _Request(), - response = response, - first_token_deadline = started + 1, + first_token_deadline = time.monotonic() + 1, ): raise AssertionError("stream should stop after disconnect") assert cancel_event.is_set() - assert state.disconnect_checks >= 2 - assert time.monotonic() - started < 0.5 - assert response.request.extensions["timeout"]["read"] <= ( - inf_mod._STREAM_DISCONNECT_POLL_TIMEOUT_S - ) + assert state.disconnect_checks == 1 + + asyncio.run(_run()) + + +def test_stream_wait_does_not_shorten_upstream_read_for_disconnect_poll(): + async def _run(): + response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}})) + seen_read_timeouts = [] + + class _Request: + async def is_disconnected(self): + return False + + class _NoItem: + async def __anext__(self): + seen_read_timeouts.append(response.request.extensions["timeout"]["read"]) + raise StopAsyncIteration + + async for _ in inf_mod._aiter_llama_stream_items( + _NoItem(), + cancel_event = threading.Event(), + request = _Request(), + response = response, + first_token_deadline = time.monotonic() + 1, + ): + raise AssertionError("stream should end") + + assert seen_read_timeouts + assert seen_read_timeouts[0] > inf_mod._STREAM_DISCONNECT_POLL_TIMEOUT_S asyncio.run(_run()) From 5c4e7b53656dfdc6923e6e72bf4fbf416bac1b42 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 20:48:10 +0200 Subject: [PATCH 11/30] Fix tool parser compatibility exports lint --- .../core/inference/tool_call_parser.py | 40 ++++++++++--------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 4f5310567f..ca3d1e4cbc 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -7,25 +7,27 @@ Tolerates missing closing tags in either ``{json}`` or ``v...`` shape. """ -from core.tool_healing import ( - _TC_END_TAG_RE, - _TC_FUNC_CLOSE_RE, - _TC_FUNC_START_RE, - _TC_GEMMA_END_TAG_RE, - _TC_GEMMA_START_RE, - _TC_JSON_START_RE, - _TC_PARAM_CLOSE_RE, - _TC_PARAM_START_RE, - _TOOL_ALL_PATS, - _TOOL_CLOSED_PATS, - _FUNC_CLOSE_TAG, - _PARAM_CLOSE_TAG, - _balanced_brace_end, - _gemma_arguments_to_json, - _inside_open_parameter, - parse_tool_calls_from_text, - strip_tool_call_markup as strip_tool_markup, -) +from core import tool_healing as _tool_healing + + +_TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS + + +def parse_tool_calls_from_text( + content: str, + *, + id_offset: int = 0, + allow_incomplete: bool = True, +) -> list[dict]: + return _tool_healing.parse_tool_calls_from_text( + content, + id_offset = id_offset, + allow_incomplete = allow_incomplete, + ) + + +def strip_tool_markup(text: str, *, final: bool = False) -> str: + return _tool_healing.strip_tool_call_markup(text, final = final) # Prefixes the streaming buffer watches for to gate in-progress text. From 3ca278cb794c57c601de502a02e2b620e8b8c93e Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 21:16:29 +0200 Subject: [PATCH 12/30] Preserve audio stream disconnect cancellation --- studio/backend/routes/inference.py | 4 +++ .../test_stream_cancel_registration_timing.py | 27 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index da8f98b1f5..fb5426047b 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4697,6 +4697,9 @@ async def openai_chat_completions( _tracker.__enter__() async def audio_input_stream(): + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: first_chunk = ChatCompletionChunk( id = completion_id, @@ -4758,6 +4761,7 @@ async def openai_chat_completions( api_monitor.fail(monitor_id, _friendly_error(e)) yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n" finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) _tracker.__exit__(None, None, None) return _SameTaskStreamingResponse( diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 9df54ba23c..bc1688fef4 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -211,6 +211,33 @@ def test_direct_llama_server_streams_install_disconnect_watcher(): ) +def test_audio_input_stream_installs_disconnect_watcher(): + audio = _async_function("audio_input_stream") + has_watcher = False + has_cleanup = False + for sub in ast.walk(audio): + if isinstance(sub, ast.Call): + fn = sub.func + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "create_task" + and isinstance(fn.value, ast.Name) + and fn.value.id == "asyncio" + and sub.args + and isinstance(sub.args[0], ast.Call) + and isinstance(sub.args[0].func, ast.Name) + and sub.args[0].func.id == "_await_disconnect_then_cancel" + ): + has_watcher = True + if isinstance(fn, ast.Name) and fn.id == "_stop_local_disconnect_cancel_watcher": + has_cleanup = True + assert has_watcher, ( + "audio_input_stream must install a disconnect watcher so client " + "disconnects set cancel_event while asyncio.to_thread(next, ...) is blocked" + ) + assert has_cleanup, "audio_input_stream must stop its disconnect watcher in finally" + + # ── Behavioral helpers ─────────────────────────────────────── _WANTED = { From 6a57d3795ab435f29250576a7ea490c87745aca1 Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 21:54:50 +0200 Subject: [PATCH 13/30] Avoid synthetic finish after passthrough errors --- studio/backend/routes/inference.py | 16 +++++++- .../tests/test_openai_tool_passthrough.py | 39 +++++++++++++++++++ 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index fb5426047b..365b1c3f3a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -9640,6 +9640,7 @@ async def _openai_passthrough_stream( monitor_done = False saw_finish_reason = False saw_done = False + saw_stream_error = False saw_tool_call_delta = False last_chunk_id = completion_id last_chunk_model = model_name @@ -9676,7 +9677,11 @@ async def _openai_passthrough_stream( data_text = raw_line[6:].strip() if data_text == "[DONE]": saw_done = True - if not saw_finish_reason and not cancel_event.is_set(): + if ( + 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, @@ -9725,13 +9730,20 @@ async def _openai_passthrough_stream( raw_line, llama_backend.context_length, ) + if monitor_event == "error": + saw_stream_error = True # Relay verbatim to preserve llama-server's native id, # finish_reason, delta.tool_calls, and usage chunks. yield raw_line + "\n\n" if monitor_event == "done": monitor_done = True break - if not saw_done and not saw_finish_reason and not cancel_event.is_set(): + if ( + not saw_done + 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, diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 190bfd6a36..1a649f09fe 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -2289,6 +2289,45 @@ class TestApiMonitorProviderAndCompletionStreams: asyncio.run(_run()) + def test_passthrough_stream_error_done_skips_synthetic_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + [ + 'data: {"error":{"message":"boom","type":"server_error"}}', + "data: [DONE]", + ], + ) + compact = result.body.replace(" ", "") + + assert '"error":{"message":"boom","type":"server_error"}' in compact + assert '"finish_reason"' not in compact + assert "data: [DONE]" in result.body + [entry] = result.monitor.snapshot() + assert entry["status"] == "error" + assert entry["error"] == "boom" + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + + def test_passthrough_stream_error_eof_skips_synthetic_finish_reason(self, monkeypatch): + async def _run(): + result = await self._run_passthrough_stream( + monkeypatch, + ['data: {"error":{"message":"boom","type":"server_error"}}'], + ) + compact = result.body.replace(" ", "") + + assert '"error":{"message":"boom","type":"server_error"}' in compact + assert '"finish_reason"' not in compact + assert "data: [DONE]" not in result.body + [entry] = result.monitor.snapshot() + assert entry["status"] == "error" + assert entry["error"] == "boom" + assert result.monitor.active_count() == 0 + + asyncio.run(_run()) + def test_passthrough_non_streaming_cancel_finalizes_monitor(self, monkeypatch): async def _run(): import routes.inference as inf_mod From db8d03927b2953ef578e57577c49dc1511d316ff Mon Sep 17 00:00:00 2001 From: wasimysaid Date: Fri, 19 Jun 2026 22:49:06 +0200 Subject: [PATCH 14/30] Address stream cleanup and Gemma parser reviews --- studio/backend/core/tool_healing.py | 7 +-- studio/backend/routes/inference.py | 30 ++++++++--- studio/backend/tests/test_mcp_servers.py | 19 +++++++ .../tests/test_openai_tool_passthrough.py | 27 ++++++++++ .../test_stream_cancel_registration_timing.py | 50 +++++++++++++++++++ 5 files changed, 124 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 855755e483..b0e74f20fa 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -16,6 +16,7 @@ import re _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), re.compile(r"<\|tool_call>.*?", re.DOTALL), + re.compile(r""), re.compile(r".*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ @@ -38,13 +39,13 @@ _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -def _balanced_brace_end(content: str, brace_start: int) -> int: +def _balanced_brace_end(content: str, brace_start: int, *, gemma_quotes: bool = False) -> int: depth = 0 i = brace_start in_string = False in_gemma_string = False while i < len(content): - if content.startswith(_GEMMA_QUOTE, i): + if gemma_quotes and not in_string and content.startswith(_GEMMA_QUOTE, i): in_gemma_string = not in_gemma_string i += len(_GEMMA_QUOTE) continue @@ -203,7 +204,7 @@ def parse_tool_calls_from_text( for m in _TC_GEMMA_START_RE.finditer(content): brace_start = m.end() - 1 - i = _balanced_brace_end(content, brace_start) + i = _balanced_brace_end(content, brace_start, gemma_quotes = True) if i < 0: continue if not allow_incomplete: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 365b1c3f3a..f6feac72c5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -752,6 +752,9 @@ class _SameTaskStreamingResponse(StreamingResponse): try: await self.stream_response(send) except OSError: + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() raise ClientDisconnect() if self.background is not None: await self.background() @@ -5136,6 +5139,21 @@ async def openai_chat_completions( _stream_usage = None _stream_timings = None _stream_finish = None + + def _flush_reasoning_extractor(): + final_reasoning, final_visible = reasoning_extractor.finish() + chunks = [] + if final_reasoning: + chunks.append( + _gguf_chat_delta_line( + ChoiceDelta(reasoning_content = final_reasoning) + ) + ) + if final_visible: + api_monitor.append_reply(monitor_id, final_visible) + chunks.append(_gguf_chat_delta_line(ChoiceDelta(content = final_visible))) + return chunks + while True: if cancel_event.is_set(): break @@ -5154,6 +5172,8 @@ async def openai_chat_completions( # cumulative cursor so the next assistant turn # streams cleanly. if not event["text"]: + for chunk in _flush_reasoning_extractor(): + yield chunk prev_text = "" reasoning_extractor = _new_chat_reasoning_extractor() # Emit tool status as a custom SSE event (including @@ -5169,6 +5189,8 @@ async def openai_chat_completions( if event["type"] in ("tool_start", "tool_end"): if event["type"] == "tool_start": + for chunk in _flush_reasoning_extractor(): + yield chunk prev_text = "" reasoning_extractor = _new_chat_reasoning_extractor() yield f"data: {json.dumps(event)}\n\n" @@ -5201,12 +5223,8 @@ async def openai_chat_completions( api_monitor.append_reply(monitor_id, visible_delta) yield _gguf_chat_delta_line(ChoiceDelta(content = visible_delta)) - final_reasoning, final_visible = reasoning_extractor.finish() - if final_reasoning: - yield _gguf_chat_delta_line(ChoiceDelta(reasoning_content = final_reasoning)) - if final_visible: - api_monitor.append_reply(monitor_id, final_visible) - yield _gguf_chat_delta_line(ChoiceDelta(content = final_visible)) + for chunk in _flush_reasoning_extractor(): + yield chunk final_chunk = ChatCompletionChunk( id = completion_id, diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index d1a9662b1d..81153e26ba 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -431,6 +431,13 @@ def test_tool_healing_strip_handles_gemma_native_tool_call(): assert out == "before after" +def test_tool_healing_strip_handles_gemma_close_only_marker(): + from core.tool_healing import strip_tool_call_markup + + assert strip_tool_call_markup("before after") == "before after" + assert strip_tool_call_markup("before after", final = True) == "before after" + + def test_tool_healing_parser_handles_gemma_native_windows_path(): from core.tool_healing import parse_tool_calls_from_text import json as _json @@ -443,6 +450,18 @@ def test_tool_healing_parser_handles_gemma_native_windows_path(): assert _json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} +def test_tool_healing_json_parser_preserves_literal_gemma_quote_token(): + from core.tool_healing import parse_tool_calls_from_text + import json as _json + + text = "" + _json.dumps( + {"name": "python", "arguments": {"code": "print('<|\"|>')"}} + ) + "" + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert _json.loads(calls[0]["function"]["arguments"]) == {"code": "print('<|\"|>')"} + + def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch): """A tool call not in the per-request list must be refused by the GGUF agentic loop (mirroring the safetensors path).""" diff --git a/studio/backend/tests/test_openai_tool_passthrough.py b/studio/backend/tests/test_openai_tool_passthrough.py index 1a649f09fe..aaef9e4dcc 100644 --- a/studio/backend/tests/test_openai_tool_passthrough.py +++ b/studio/backend/tests/test_openai_tool_passthrough.py @@ -1565,6 +1565,33 @@ class TestGgufVisionToolRouting: [entry] = result.monitor.snapshot() assert entry["reply"] == "visible " + def test_gguf_tool_stream_flushes_held_text_before_status_reset(self, monkeypatch): + def _tools(**_kwargs): + yield {"type": "content", "text": "answer <"} + yield {"type": "status", "text": ""} + yield { + "type": "metadata", + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + "finish_reason": "stop", + } + + result = self._run_gguf_case( + monkeypatch, + tool_generate = _tools, + payload_kwargs = { + "stream": True, + "enable_tools": True, + "enabled_tools": ["terminal"], + "messages": [{"role": "user", "content": "say literal"}], + }, + ) + deltas = [p["choices"][0].get("delta", {}) for p in result.payloads if p.get("choices")] + + combined_content = "".join(d.get("content", "") for d in deltas) + assert combined_content == "answer <" + [entry] = result.monitor.snapshot() + assert entry["reply"] == "answer <" + def test_non_streaming_gguf_splits_reasoning_content(self, monkeypatch): def _generate(**_kwargs): yield "planvisible" diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index bc1688fef4..09e30bb730 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -275,6 +275,23 @@ def _load_registry_module(): return mod +def _load_same_task_response_module(): + for n in _TREE.body: + if isinstance(n, ast.ClassDef) and n.name == "_SameTaskStreamingResponse": + source = ast.get_source_segment(SRC, n) + break + else: + raise AssertionError("_SameTaskStreamingResponse missing") + mod = {} + exec( + "class StreamingResponse: pass\n" + "class ClientDisconnect(Exception): pass\n" + + source, + mod, + ) + return mod + + def _make_stream(tracker, raise_exc): async def gen(): try: @@ -379,6 +396,39 @@ def test_finally_cleanup_on_aclose(): assert "sid-abort" not in m["_CANCEL_REGISTRY"] +def test_same_task_response_closes_body_iterator_on_send_disconnect(): + m = _load_same_task_response_module() + closed = False + + async def body(): + nonlocal closed + try: + yield "data: first\n\n" + finally: + closed = True + + async def run(): + agen = body() + await agen.__anext__() + response = m["_SameTaskStreamingResponse"].__new__(m["_SameTaskStreamingResponse"]) + response.body_iterator = agen + response.background = None + + async def stream_response(_send): + raise OSError("client disconnected") + + response.stream_response = stream_response + try: + await response({}, None, lambda _message: None) + except m["ClientDisconnect"]: + pass + else: + raise AssertionError("expected ClientDisconnect") + + asyncio.run(run()) + assert closed + + def test_preset_cancel_event_exits_cleanly_with_done(): # Pending-replay: a stashed cancel pre-set cancel_event. The loop must break # cleanly with final_chunk + [DONE], not propagate GeneratorExit from the GGUF wrapper. From 67ccd784402dadd9a9b137027f144ef284f3f585 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 19 Jun 2026 20:49:37 +0000 Subject: [PATCH 15/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/tool_healing.py | 7 ++++++- studio/backend/tests/test_mcp_servers.py | 9 +++++---- tests/studio/test_stream_cancel_registration_timing.py | 4 +--- 3 files changed, 12 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index b0e74f20fa..fcc733796d 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -39,7 +39,12 @@ _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -def _balanced_brace_end(content: str, brace_start: int, *, gemma_quotes: bool = False) -> int: +def _balanced_brace_end( + content: str, + brace_start: int, + *, + gemma_quotes: bool = False, +) -> int: depth = 0 i = brace_start in_string = False diff --git a/studio/backend/tests/test_mcp_servers.py b/studio/backend/tests/test_mcp_servers.py index 81153e26ba..12239e7113 100644 --- a/studio/backend/tests/test_mcp_servers.py +++ b/studio/backend/tests/test_mcp_servers.py @@ -433,7 +433,6 @@ def test_tool_healing_strip_handles_gemma_native_tool_call(): def test_tool_healing_strip_handles_gemma_close_only_marker(): from core.tool_healing import strip_tool_call_markup - assert strip_tool_call_markup("before after") == "before after" assert strip_tool_call_markup("before after", final = True) == "before after" @@ -454,9 +453,11 @@ def test_tool_healing_json_parser_preserves_literal_gemma_quote_token(): from core.tool_healing import parse_tool_calls_from_text import json as _json - text = "" + _json.dumps( - {"name": "python", "arguments": {"code": "print('<|\"|>')"}} - ) + "" + text = ( + "" + + _json.dumps({"name": "python", "arguments": {"code": "print('<|\"|>')"}}) + + "" + ) calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 assert _json.loads(calls[0]["function"]["arguments"]) == {"code": "print('<|\"|>')"} diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 09e30bb730..33deb7af9d 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -284,9 +284,7 @@ def _load_same_task_response_module(): raise AssertionError("_SameTaskStreamingResponse missing") mod = {} exec( - "class StreamingResponse: pass\n" - "class ClientDisconnect(Exception): pass\n" - + source, + "class StreamingResponse: pass\nclass ClientDisconnect(Exception): pass\n" + source, mod, ) return mod From 2b7360176b313421b992d7d8f972c9e689cab463 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 20 Jun 2026 07:14:44 +0000 Subject: [PATCH 16/30] Gemma 4: parse bare-string tool args and keep safetensors tools for native <|tool_call> - Quote bare unquoted string values in Gemma native tool-call args (e.g. {location:Tokyo,unit:celsius}) so they parse; JSON scalars stay typed. - Stop _detect_safetensors_features from suppressing supports_tools for templates that emit Gemma native <|tool_call>, which the shared parser now reads. - Add tests for both. --- studio/backend/core/tool_healing.py | 16 ++++++++++++++++ studio/backend/routes/inference.py | 11 ++++++----- .../test_safetensors_capability_advertise.py | 14 ++++++++++++++ .../backend/tests/test_safetensors_tool_loop.py | 9 +++++++++ .../tests/test_tool_call_parser_strict.py | 12 ++++++++++++ 5 files changed, 57 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fcc733796d..4663512033 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -138,6 +138,22 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 + # Gemma may emit bare string values ({unit:celsius}); quote them so + # json.loads succeeds. JSON scalars/objects/arrays/quoted stay as-is. + ws = i + while i < len(src) and src[i].isspace(): + i += 1 + 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 ",}": + i += 1 + raw = src[v_start:i] + try: + json.loads(raw.strip()) + parts.append(raw) + except (json.JSONDecodeError, ValueError): + parts.append(json.dumps(raw.strip()) if raw.strip() else raw) else: parts.append(src[key_start:i]) return "".join(parts) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index f6feac72c5..17bcc084ee 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -1128,16 +1128,17 @@ def _detect_safetensors_features(backend, chat_template: Optional[str]) -> dict: "supports_tools": False, } ) - # Our safetensors loop only parses {json} and - # .... Llama uses <|python_tag|>, Mistral uses - # [TOOL_CALLS]; advertising tools for those enables a pill the parser - # can't honour. GGUF is unaffected -- llama-server normalises every - # format into structured deltas. + # Our safetensors loop only parses {json}, + # ..., and Gemma native <|tool_call>.... + # Llama uses <|python_tag|>, Mistral uses [TOOL_CALLS]; advertising tools for + # those enables a pill the parser can't honour. GGUF is unaffected -- + # llama-server normalises every format into structured deltas. if ( flags.get("supports_tools") and chat_template and "" not in chat_template and "" not in chat_template ): logger.info( "safetensors: template advertises tools but uses an " diff --git a/studio/backend/tests/test_safetensors_capability_advertise.py b/studio/backend/tests/test_safetensors_capability_advertise.py index 13cb6bbd46..671af93708 100644 --- a/studio/backend/tests/test_safetensors_capability_advertise.py +++ b/studio/backend/tests/test_safetensors_capability_advertise.py @@ -203,6 +203,20 @@ def test_detect_safetensors_features_function_xml_format_keeps_tools_on(): assert flags["supports_tools"] is True +def test_detect_safetensors_features_gemma_native_tool_call_keeps_tools_on(): + """Gemma 4 emits <|tool_call>call:name{...}, which the shared + parser now reads, so the gate must not suppress tools for it.""" + from routes.inference import _detect_safetensors_features + + tpl_with_gemma_native = ( + "{%- if tools -%}Tool call format: " + "<|tool_call>call:name{key:value}{%- endif -%}" + ) + backend = SimpleNamespace(active_model_name = "unsloth/gemma-4-12b-it") + flags = _detect_safetensors_features(backend, tpl_with_gemma_native) + assert flags["supports_tools"] is True + + # Qwen3.5 family pin: the live GGUF + safetensors templates both wrap tool # calls as ``\n...``. Faithful slice so the # classifier never silently regresses for this family. diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 6b470d87a7..6e450e6cbc 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -99,6 +99,15 @@ class TestParser: assert result[0]["function"]["name"] == "terminal" assert json.loads(result[0]["function"]["arguments"]) == {"command": "echo {foo:bar}"} + def test_gemma_native_tool_call_bare_string_values(self): + text = '<|tool_call>call:get_weather{location:Tokyo,unit:celsius}' + result = parse_tool_calls_from_text(text) + assert len(result) == 1 + assert json.loads(result[0]["function"]["arguments"]) == { + "location": "Tokyo", + "unit": "celsius", + } + def test_xml_function_call(self): text = "print('hi')" result = parse_tool_calls_from_text(text) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 799f38710f..61c625f6fe 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -139,6 +139,18 @@ class TestGemmaNativeStyle: assert len(calls) == 1 assert json.loads(calls[0]["function"]["arguments"]) == {"path": r"C:\Users\wasim\repo"} + def test_bare_unquoted_string_values_are_accepted(self): + # Gemma can emit enum/string args unquoted; bare JSON scalars stay typed. + text = '<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}' + calls = parse_tool_calls_from_text(text, allow_incomplete = False) + assert len(calls) == 1 + assert json.loads(calls[0]["function"]["arguments"]) == { + "location": "Tokyo", + "unit": "celsius", + "days": 3, + "live": True, + } + class TestHealingPathUnaffected: def test_auto_heal_still_repairs_unclosed_function(self): From 4b68cac412ab4a017e5e1e86d5dfb7621cc7f1a2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 20 Jun 2026 07:15:26 +0000 Subject: [PATCH 17/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_safetensors_tool_loop.py | 2 +- studio/backend/tests/test_tool_call_parser_strict.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_safetensors_tool_loop.py b/studio/backend/tests/test_safetensors_tool_loop.py index 6e450e6cbc..0b6869275a 100644 --- a/studio/backend/tests/test_safetensors_tool_loop.py +++ b/studio/backend/tests/test_safetensors_tool_loop.py @@ -100,7 +100,7 @@ class TestParser: assert json.loads(result[0]["function"]["arguments"]) == {"command": "echo {foo:bar}"} def test_gemma_native_tool_call_bare_string_values(self): - text = '<|tool_call>call:get_weather{location:Tokyo,unit:celsius}' + text = "<|tool_call>call:get_weather{location:Tokyo,unit:celsius}" result = parse_tool_calls_from_text(text) assert len(result) == 1 assert json.loads(result[0]["function"]["arguments"]) == { diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 61c625f6fe..931d8a705d 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -141,7 +141,9 @@ class TestGemmaNativeStyle: def test_bare_unquoted_string_values_are_accepted(self): # Gemma can emit enum/string args unquoted; bare JSON scalars stay typed. - text = '<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}' + text = ( + "<|tool_call>call:get_weather{location:Tokyo,unit:celsius,days:3,live:true}" + ) calls = parse_tool_calls_from_text(text, allow_incomplete = False) assert len(calls) == 1 assert json.loads(calls[0]["function"]["arguments"]) == { From 491586c55678c51cc9f15c3bf32c7d75925c7445 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 04:59:29 +0000 Subject: [PATCH 18/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 046968cb71..3ee0319a7d 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -818,6 +818,8 @@ class _SameTaskStreamingResponse(StreamingResponse): raise ClientDisconnect() if self.background is not None: await self.background() + + async def _aclose_stream_resources( *, watchers = (), From b3e244d658a6edc90e802626a2ea350e5d92238e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 11:03:21 +0000 Subject: [PATCH 19/30] 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} From 0083826dde39989df4973bfe46b49c927979be0d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 11:49:46 +0000 Subject: [PATCH 20/30] Emit the terminal finish_reason chunk in GGUF streams The OpenAI chat-completions GGUF tool stream and plain stream both built a final ChatCompletionChunk carrying finish_reason but never yielded it, so clients received the optional usage chunk and [DONE] with no chunk carrying finish_reason. OpenAI-compatible consumers rely on that terminal choice to distinguish stop/length/tool_calls. Yield it before the usage chunk and [DONE], matching the other streaming paths. --- studio/backend/routes/inference.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5c88387e9b..a7b04c7c53 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -5306,6 +5306,10 @@ async def openai_chat_completions( ) ], ) + # Emit the terminal chunk carrying finish_reason before the + # optional usage chunk and [DONE], so OpenAI-compatible + # clients can detect stop/length/tool_calls. + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" usage_line = _openai_stream_usage_chunk( payload, completion_id, @@ -5458,6 +5462,10 @@ async def openai_chat_completions( ) ], ) + # Emit the terminal chunk carrying finish_reason before the + # optional usage chunk and [DONE], so OpenAI-compatible + # clients can detect stop/length/tool_calls. + yield f"data: {final_chunk.model_dump_json(exclude_none = True)}\n\n" usage_line = _openai_stream_usage_chunk( payload, completion_id, From eca61f7f5e67d8440f639b2b30f1c9e34d62f357 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 12:24:41 +0000 Subject: [PATCH 21/30] Parse tool calls in document order and skip nested markers both ways Unify the JSON- and Gemma-format tool-call passes into a single position-ordered scan: - Calls are now emitted in byte order across both formats, so a mixed output like `<|tool_call>call:create{...} ... {"name":"read",...}` executes create before read, matching the order they appear in (tools run in returned order). - A candidate that starts inside an already-accepted call's span is skipped, in both directions: a JSON marker inside a Gemma argument and a Gemma marker inside a JSON argument are treated as data, not promoted to a second executable tool call. Extends tests/test_gemma_tool_parse_edge_cases.py with the ordering and JSON-in-Gemma nesting cases. --- studio/backend/core/tool_healing.py | 89 +++++++++---------- .../tests/test_gemma_tool_parse_edge_cases.py | 23 +++++ 2 files changed, 63 insertions(+), 49 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index e91b541c0d..5c99773b95 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -207,63 +207,54 @@ 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]] = [] - + # Collect JSON- and Gemma-format candidates with their byte spans, then + # accept them in document order. Both order and spans matter: + # * tools execute in returned order, so a call appearing earlier in the + # text must be emitted first even across the two formats; + # * a tool-call marker INSIDE another call's argument string is data, not a + # call, so a candidate starting within an already accepted span is + # skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker + # nested in a JSON arg alike, regardless of which format is outer). + candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): - brace_start = m.end() - 1 - i = _balanced_brace_end(content, brace_start) - if i < 0: + end = _balanced_brace_end(content, m.end() - 1) + if end >= 0: + candidates.append((m.start(), end, "json", m)) + for m in _TC_GEMMA_START_RE.finditer(content): + end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True) + if end >= 0: + candidates.append((m.start(), end, "gemma", m)) + candidates.sort(key = lambda c: c[0]) + + consumed: list[tuple[int, int]] = [] + for start, end, kind, m in candidates: + if any(s <= start < e for s, e in consumed): continue if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_END_TAG_RE.match(tail_after_json) is None: + tail = content[end + 1 :].lstrip() + close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE + if close_re.match(tail) is None: continue - json_str = content[brace_start : i + 1] try: - obj = json.loads(json_str) - tc = { + if kind == "json": + obj = json.loads(content[m.end() - 1 : end + 1]) + name = obj.get("name", "") + arguments = obj.get("arguments", {}) + if isinstance(arguments, dict): + arguments = json.dumps(arguments) + else: + name = m.group(1) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end])) + except (json.JSONDecodeError, ValueError): + continue + tool_calls.append( + { "id": f"call_{id_offset + len(tool_calls)}", "type": "function", - "function": { - "name": obj.get("name", ""), - "arguments": obj.get("arguments", {}), - }, + "function": {"name": name, "arguments": arguments}, } - 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: - continue - if not allow_incomplete: - tail_after_json = content[i + 1 :].lstrip() - if _TC_GEMMA_END_TAG_RE.match(tail_after_json) is None: - continue - try: - tool_calls.append( - { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": m.group(1), - "arguments": json.dumps(_gemma_arguments_to_json(content[m.end() : i])), - }, - } - ) - consumed.append((m.start(), i + 1)) - except (json.JSONDecodeError, ValueError): - pass + ) + consumed.append((start, end + 1)) if not tool_calls: func_starts = [ 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 45e3277a4b..eb3c31c8a7 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -62,3 +62,26 @@ def test_two_separate_gemma_calls_both_parse(): assert [c["function"]["name"] for c in calls] == ["a", "b"], calls assert _args(calls[0]) == {"x": 1} assert _args(calls[1]) == {"y": 2} + + +def test_mixed_format_calls_preserve_document_order(): + # A Gemma-native call precedes a JSON-format call in the text; tools execute + # in returned order, so `create` must come before `read`. + content = ( + '<|tool_call>call:create{path:a} then ' + '{"name":"read","arguments":{"path":"a"}}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["create", "read"], calls + + +def test_json_marker_inside_gemma_argument_is_not_a_second_call(): + # The reverse of the JSON-outer case: a JSON-style marker inside a Gemma + # call's quoted argument is code text, not a second `terminal` call. + content = ( + '<|tool_call>call:python{code:<|"|>' + 'print({"name":"terminal","arguments":{"command":"ls"}})' + '<|"|>}' + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls From 11564256b9107f0088df80825b97c6a9366827c3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 12:26:01 +0000 Subject: [PATCH 22/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_gemma_tool_parse_edge_cases.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 eb3c31c8a7..e828b6dcea 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -68,7 +68,7 @@ def test_mixed_format_calls_preserve_document_order(): # A Gemma-native call precedes a JSON-format call in the text; tools execute # in returned order, so `create` must come before `read`. content = ( - '<|tool_call>call:create{path:a} then ' + "<|tool_call>call:create{path:a} then " '{"name":"read","arguments":{"path":"a"}}' ) calls = parse_tool_calls_from_text(content) From 520df9fe9dffa244a7bc1f4c83175e83e9326e05 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 12:59:00 +0000 Subject: [PATCH 23/30] Quote bare Gemma array elements; order finish before trailing usage - _quote_gemma_object_keys skipped array values, so a Gemma call with a bare-string array argument like labels:[bug,ui] produced invalid JSON and the whole tool call was dropped. Array values are now scanned and bare string elements quoted, while numbers, quoted strings, and JSON literals are preserved. - In the OpenAI passthrough stream, a trailing usage-only chunk (stream_options.include_usage) that arrived before any finish chunk was relayed before the synthetic finish, producing usage -> finish -> [DONE]. Emit the synthetic finish before that usage chunk so the order matches the other streams (finish -> usage -> [DONE]). Extends tests/test_gemma_tool_parse_edge_cases.py with the bare-array cases. --- studio/backend/core/tool_healing.py | 86 ++++++++++++++++++- studio/backend/routes/inference.py | 19 ++++ .../tests/test_gemma_tool_parse_edge_cases.py | 15 ++++ 3 files changed, 119 insertions(+), 1 deletion(-) 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"]} From 1e58c3707dab1b1c68564ec6fe58edd96d2f6073 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 15:14:42 +0000 Subject: [PATCH 24/30] Harden Gemma array parsing, XML-parameter guard, and stream teardown Address five review findings on the Gemma tool-call and OpenAI passthrough streaming paths: - parse_tool_calls_from_text collected JSON and Gemma markers without the _inside_open_parameter guard, so a marker embedded in an existing value was promoted to a separate tool call. Candidates that start inside an open XML parameter are now skipped, matching the guard the XML-style parser already applies. - _quote_gemma_array_elements preserved array elements starting with { or [ verbatim, so an array of objects (items:[{path:a}]) or a nested array failed json.loads and the whole call was dropped. Object and nested-array elements are now normalised recursively. - _openai_passthrough_stream synthesized a finish chunk before a trailing usage-only chunk and set saw_finish_reason, which made the EOF guard skip the [DONE] sentinel. The EOF path now emits [DONE] whenever the upstream omitted it, even after a finish chunk was already synthesized. - /generate/stream drove generation through asyncio.to_thread with no disconnect watcher, so a client disconnect during a long generation went unnoticed until the next send. It now runs _await_disconnect_then_cancel against the request, matching the other local streaming endpoints. - _SameTaskStreamingResponse closed the body iterator with aclose() on a send-side disconnect, raising GeneratorExit so the generators' cancellation handlers (which finish the api_monitor entry) never ran. It now throws CancelledError, falling back to aclose() when athrow is unavailable. Extends tests/test_gemma_tool_parse_edge_cases.py with array-of-objects, nested-array, and marker-inside-XML-parameter cases. --- studio/backend/core/tool_healing.py | 33 ++++++++-- studio/backend/routes/inference.py | 64 +++++++++++++------ .../tests/test_gemma_tool_parse_edge_cases.py | 43 +++++++++++++ 3 files changed, 116 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 14d87f69e9..7ef1870519 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -136,16 +136,32 @@ def _split_top_level_commas(src: str) -> list: 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.""" + """Normalise the elements of a Gemma array value so json.loads succeeds. + + Gemma may emit ``labels:[bug,ui]`` without per-element quotes, or arrays of + objects (``items:[{path:a}]``) whose keys/values also lack quotes; left + as-is json.loads fails and the whole call is dropped. Bare string elements + are quoted, object and nested-array elements are normalised recursively, and + 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 '"{[': + if not stripped or stripped[0] == '"': out.append(element) continue + if stripped[0] == "{": + # Object element: quote its keys/bare values like a top-level object. + out.append(_quote_gemma_object_keys(stripped)) + continue + if stripped[0] == "[": + # Nested array: normalise its elements too. + inner_end = _balanced_bracket_end(stripped, 0) + if inner_end == len(stripped) - 1: + out.append("[" + _quote_gemma_array_elements(stripped[1:inner_end]) + "]") + else: + out.append(element) + continue try: json.loads(stripped) out.append(element) @@ -301,10 +317,17 @@ def parse_tool_calls_from_text( # nested in a JSON arg alike, regardless of which format is outer). candidates = [] # (start, brace_end, kind, match) for m in _TC_JSON_START_RE.finditer(content): + # A marker that begins inside an open value + # is that parameter's data, not its own call; skip it (same guard the + # XML-style parser below applies to nested = 0: candidates.append((m.start(), end, "json", m)) for m in _TC_GEMMA_START_RE.finditer(content): + if _inside_open_parameter(content, m.start()): + continue end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = True) if end >= 0: candidates.append((m.start(), end, "gemma", m)) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 01ed86207a..1f50f2815a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -812,9 +812,22 @@ class _SameTaskStreamingResponse(StreamingResponse): try: await self.stream_response(send) except OSError: - aclose = getattr(self.body_iterator, "aclose", None) - if aclose is not None: - await aclose() + # Client disconnected mid-send. Throw CancelledError into the body + # generator instead of aclose() (which raises GeneratorExit): the + # generators run their `except asyncio.CancelledError` handler, which + # finishes the api_monitor entry as "cancelled", whereas GeneratorExit + # skips it and only runs `finally`, leaving the monitor entry active. + # Fall back to aclose() for iterators without athrow. + athrow = getattr(self.body_iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + pass + else: + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() raise ClientDisconnect() if self.background is not None: await self.background() @@ -3332,7 +3345,9 @@ async def get_api_monitor_entry(entry_id: str, current_subject: str = Depends(ge @router.post("/generate/stream") async def generate_stream( - request: GenerateRequest, current_subject: str = Depends(get_current_subject) + request: GenerateRequest, + fastapi_request: Request, + current_subject: str = Depends(get_current_subject), ): """ Generate a chat response with Server-Sent Events (SSE) streaming. @@ -3382,6 +3397,13 @@ async def generate_stream( async def stream(): gen = None completed = False + # Cancel the generation when the client disconnects. The generator only + # awaits asyncio.to_thread(next, gen, ...), so without a concurrent + # watcher a disconnect during a long prefill/generation would go + # unnoticed until the next send and the backend would keep generating. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(fastapi_request, cancel_event) + ) try: gen = backend.generate_chat_response( messages = request.messages, @@ -3396,12 +3418,15 @@ async def generate_stream( ) _DONE = object() while True: + if cancel_event.is_set(): + break chunk = await asyncio.to_thread(next, gen, _DONE) if chunk is _DONE: + completed = True break yield f"data: {json.dumps({'content': chunk})}\n\n" - completed = True - yield "data: [DONE]\n\n" + if completed: + yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() @@ -3413,6 +3438,7 @@ async def generate_stream( logger.error(f"Error during generation: {e}", exc_info = True) yield f"data: {json.dumps({'error': _friendly_error(e)})}\n\n" finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) if not completed and not cancel_event.is_set(): cancel_event.set() backend.reset_generation_state() @@ -9624,19 +9650,19 @@ async def _openai_passthrough_stream( if monitor_event == "done": monitor_done = True break - if ( - not saw_done - 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" + if not saw_done and not saw_stream_error and not cancel_event.is_set(): + # Synthesize a finish chunk only if one was not already + # emitted (e.g. before a trailing usage-only chunk), but + # always close with [DONE] whenever the upstream omitted it, + # so the stream ends on the [DONE] sentinel either way. + if not saw_finish_reason: + finish_line = _synthetic_finish_line() + _monitor_openai_sse_line( + monitor_id, + finish_line, + llama_backend.context_length, + ) + yield finish_line + "\n\n" done_line = "data: [DONE]" _monitor_openai_sse_line( monitor_id, 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 9dd8607bd5..f6c2fdd2c6 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -100,3 +100,46 @@ def test_array_keeps_numbers_and_quoted_elements(): '<|tool_call>call:f{nums:[1,2],tags:[<|"|>a,b<|"|>,c]}' ) assert _args(calls[0]) == {"nums": [1, 2], "tags": ["a,b", "c"]} + + +def test_array_of_objects_is_normalised(): + # Arrays of objects are a common tool-schema shape; their (unquoted) keys and + # bare values must be normalised too, not left verbatim, or the call drops. + calls = parse_tool_calls_from_text( + "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" + ) + assert len(calls) == 1, calls + assert _args(calls[0]) == { + "items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}] + } + + +def test_nested_array_elements_are_normalised(): + calls = parse_tool_calls_from_text( + "<|tool_call>call:grid{cells:[[a,b],[c,d]]}" + ) + assert _args(calls[0]) == {"cells": [["a", "b"], ["c", "d"]]} + + +def test_gemma_marker_inside_xml_parameter_is_not_a_second_call(): + # An XML-style call whose value contains a + # Gemma marker: the marker is the parameter's data, not a separate terminal + # call, so only the python call must be returned. + content = ( + "" + "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 + assert "terminal" in _args(calls[0])["code"] + + +def test_json_marker_inside_xml_parameter_is_not_a_second_call(): + content = ( + "" + 'run({"name":"terminal","arguments":{"command":"ls"}})' + "" + ) + calls = parse_tool_calls_from_text(content) + assert [c["function"]["name"] for c in calls] == ["python"], calls From 6618213da549e6b22a58551a2c4b05fca43b09ce Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:16:42 +0000 Subject: [PATCH 25/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_gemma_tool_parse_edge_cases.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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 f6c2fdd2c6..b198f426df 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -109,15 +109,11 @@ def test_array_of_objects_is_normalised(): "<|tool_call>call:batch{items:[{path:a,mode:r},{path:b,mode:w}]}" ) assert len(calls) == 1, calls - assert _args(calls[0]) == { - "items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}] - } + assert _args(calls[0]) == {"items": [{"path": "a", "mode": "r"}, {"path": "b", "mode": "w"}]} def test_nested_array_elements_are_normalised(): - calls = parse_tool_calls_from_text( - "<|tool_call>call:grid{cells:[[a,b],[c,d]]}" - ) + calls = parse_tool_calls_from_text("<|tool_call>call:grid{cells:[[a,b],[c,d]]}") assert _args(calls[0]) == {"cells": [["a", "b"], ["c", "d"]]} From e1fdb140b4bad80767c2e74c8db5c12130bb65a1 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 15:37:00 +0000 Subject: [PATCH 26/30] Watch disconnects on Anthropic streams; keep timestamps in Gemma values Two follow-ups on the streaming and tool-parse paths: - _anthropic_tool_stream and _anthropic_plain_stream drove generation through asyncio.to_thread(next, gen, ...) and only polled is_disconnected() between events, so a client disconnect during prefill or a long generation/tool step held the decode slot until the next event or a failed send. Both now run the _await_disconnect_then_cancel watcher used by the other local streams, stop it in finally, and break promptly when cancel_event is set. - _GEMMA_NEXT_KEY_RE treated any comma followed by word-chars-then-colon as the next key, so a bare value such as "meet at 10:00, 11:00 tomorrow" was split into bogus keys. The next-key token must now be identifier-shaped (start with a letter or underscore), so a comma before a timestamp, ratio, or other numeric-then-colon text stays part of the value. Adds a timestamp-in-bare-value regression test. --- studio/backend/core/tool_healing.py | 7 ++++-- studio/backend/routes/inference.py | 24 +++++++++++++++++-- .../tests/test_gemma_tool_parse_edge_cases.py | 13 ++++++++++ 3 files changed, 40 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 7ef1870519..c647738a3a 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -39,8 +39,11 @@ _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*:") +# `location:New York, NY`), so it must not terminate the value. The key token +# must be identifier-shaped (start with a letter or underscore); a comma +# followed by digits-then-colon is value text such as a timestamp or ratio +# (`meet at 10:00, 11:00 tomorrow`), not a new key. +_GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") def _balanced_brace_end( diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 1f50f2815a..6b31662885 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -8471,9 +8471,17 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() + # Concurrent disconnect watcher: the loop only polls is_disconnected() + # between events, so a client disconnect during a long prefill or + # generation step would otherwise hold the decode slot until the next + # event or a failed send. The watcher sets cancel_event so the backend + # stops promptly. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return event = await asyncio.to_thread(next, gen, _sentinel) @@ -8521,6 +8529,8 @@ async def _anthropic_tool_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) stop_reason = openai_finish_to_anthropic_stop( captured_finish_reason, had_tool_calls = ends_on_tool_use @@ -8557,9 +8567,17 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() + # Concurrent disconnect watcher: the loop only polls is_disconnected() + # between chunks, so a client disconnect during a long prefill or + # generation step would otherwise hold the decode slot until the next + # chunk or a failed send. The watcher sets cancel_event so the backend + # stops promptly. + disconnect_watcher = asyncio.create_task( + _await_disconnect_then_cancel(request, cancel_event) + ) try: while True: - if await request.is_disconnected(): + if cancel_event.is_set() or await request.is_disconnected(): cancel_event.set() return cumulative = await asyncio.to_thread(next, gen, _sentinel) @@ -8582,6 +8600,8 @@ async def _anthropic_plain_stream( if _error_event is not None: yield _error_event return + finally: + await _stop_local_disconnect_cancel_watcher(disconnect_watcher) stop_reason = openai_finish_to_anthropic_stop(captured_finish_reason, had_tool_calls = False) for line in emitter.finish(stop_reason = stop_reason, stop_sequence = None): 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 b198f426df..15ca81cb5c 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -45,6 +45,19 @@ def test_normal_multi_key_arguments_still_split(): assert _args(calls[0]) == {"a": 1, "b": "hello", "c": "x,y"} +def test_bare_value_with_timestamps_after_comma_is_kept(): + # A comma followed by digits-then-colon (a timestamp/ratio) is value text, + # not a new key, so the whole query must be preserved as one argument. + calls = parse_tool_calls_from_text( + "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" + ) + assert len(calls) == 1, calls + assert _args(calls[0]) == { + "query": "meet at 10:00, 11:00 tomorrow", + "priority": "high", + } + + 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. From d376756e9c4d7aba2c25768d64c143321d99ccf4 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:37:34 +0000 Subject: [PATCH 27/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_gemma_tool_parse_edge_cases.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 15ca81cb5c..411e6de1f3 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -52,10 +52,7 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) assert len(calls) == 1, calls - assert _args(calls[0]) == { - "query": "meet at 10:00, 11:00 tomorrow", - "priority": "high", - } + assert _args(calls[0]) == {"query": "meet at 10:00, 11:00 tomorrow", "priority": "high"} def test_marker_inside_json_argument_is_not_a_second_call(): From b0dbe438671a0d01be1f88e54fb475ebd433fff2 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 16:24:09 +0000 Subject: [PATCH 28/30] Guard nested markers, reset on disconnect, clean unstarted streams Three follow-ups on the tool-parse and streaming paths: - parse_tool_calls_from_text only skipped markers that fell inside a span it had already parsed successfully, so when an unquoted Gemma argument contained a literal marker (code:<|tool_call>call:terminal{...}) the outer object failed to normalize, its span was never recorded, and the inner marker was promoted to a standalone terminal call. Candidates nested inside any other candidate's brace span are now skipped regardless of whether the enclosing candidate parsed, so a marker in malformed outer data is never executed. - /generate/stream skipped backend.reset_generation_state() when the disconnect watcher set cancel_event between chunks: the loop broke and the finally's reset is guarded on cancel_event being unset. A subprocess backend kept decoding after the client left. The cancel-break path now resets the backend. - _SameTaskStreamingResponse threw CancelledError / called aclose() on the body iterator on a send-side disconnect, but neither runs the try/finally of a generator that never started (early disconnect on http.response.start), so the passthrough's eagerly-opened upstream httpx stream and cancel-registry entry leaked. It now tracks whether the body started and, when it did not, runs an optional unstarted_cleanup hook; the OpenAI passthrough wires it to close the upstream resp/client and exit the cancel tracker. Adds a nested-unquoted-marker regression test. --- studio/backend/core/tool_healing.py | 12 ++- studio/backend/routes/inference.py | 77 +++++++++++++++---- .../tests/test_gemma_tool_parse_edge_cases.py | 10 +++ 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index c647738a3a..fe26d48c7f 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -336,9 +336,14 @@ def parse_tool_calls_from_text( candidates.append((m.start(), end, "gemma", m)) candidates.sort(key = lambda c: c[0]) - consumed: list[tuple[int, int]] = [] - for start, end, kind, m in candidates: - if any(s <= start < e for s, e in consumed): + spans = [(s, e) for s, e, _kind, _m in candidates] + for idx, (start, end, kind, m) in enumerate(candidates): + # Skip a candidate nested inside another candidate's brace span: it is + # the enclosing call's argument data, not its own call. Checked against + # every candidate span (not only the ones that parsed successfully), so a + # marker inside an outer call that later fails to normalize is still + # never promoted to its own executable tool call. + if any(s <= start and end <= e for j, (s, e) in enumerate(spans) if j != idx): continue if not allow_incomplete: tail = content[end + 1 :].lstrip() @@ -364,7 +369,6 @@ def parse_tool_calls_from_text( "function": {"name": name, "arguments": arguments}, } ) - consumed.append((start, end + 1)) if not tool_calls: func_starts = [ diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 6b31662885..51700e4017 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -808,26 +808,61 @@ def _same_task_timeout(timeout_s: float): class _SameTaskStreamingResponse(StreamingResponse): """StreamingResponse without Starlette's legacy AnyIO task-group wrapper.""" + def __init__(self, *args, unstarted_cleanup = None, **kwargs) -> None: + super().__init__(*args, **kwargs) + # Async callable invoked when the client disconnects before the body + # iterator is ever advanced. A generator that never started cannot run + # its own try/finally, so a stream that acquires resources before its + # first yield (the passthrough opens an upstream httpx stream eagerly) + # passes this to release them. + self._unstarted_cleanup = unstarted_cleanup + async def __call__(self, scope, receive, send) -> None: + # Track whether the body iterator was ever advanced: send() only emits a + # body message after the generator yields its first chunk, so a failure + # before then means it never entered its try/finally. + body_started = False + + async def _tracking_send(message) -> None: + nonlocal body_started + if message.get("type") == "http.response.body": + body_started = True + await send(message) + try: - await self.stream_response(send) + await self.stream_response(_tracking_send) except OSError: - # Client disconnected mid-send. Throw CancelledError into the body - # generator instead of aclose() (which raises GeneratorExit): the - # generators run their `except asyncio.CancelledError` handler, which - # finishes the api_monitor entry as "cancelled", whereas GeneratorExit - # skips it and only runs `finally`, leaving the monitor entry active. - # Fall back to aclose() for iterators without athrow. - athrow = getattr(self.body_iterator, "athrow", None) - if athrow is not None: - try: - await athrow(asyncio.CancelledError()) - except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): - pass + # Client disconnected mid-send. + if body_started: + # The generator produced at least one chunk and is suspended in + # its try/finally. Throw CancelledError into it (not aclose's + # GeneratorExit) so its `except asyncio.CancelledError` handler + # runs and finishes any api_monitor entry; GeneratorExit would + # skip it and only run `finally`. Fall back to aclose() without + # athrow. + athrow = getattr(self.body_iterator, "athrow", None) + if athrow is not None: + try: + await athrow(asyncio.CancelledError()) + except (asyncio.CancelledError, StopAsyncIteration, RuntimeError): + pass + else: + aclose = getattr(self.body_iterator, "aclose", None) + if aclose is not None: + await aclose() else: + # http.response.start failed before the body iterator advanced, + # so its try/finally never armed and aclose()/athrow() are no-ops + # on an unstarted generator. Release any resources acquired + # before the first yield via the explicit cleanup hook. aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() + if self._unstarted_cleanup is not None: + try: + await self._unstarted_cleanup() + except Exception: + pass raise ClientDisconnect() if self.background is not None: await self.background() @@ -3419,6 +3454,13 @@ async def generate_stream( _DONE = object() while True: if cancel_event.is_set(): + # The disconnect watcher set cancel_event between chunks. + # Reset the backend here: closing the Python generator does + # not signal a subprocess backend, so without this it keeps + # decoding after the client is gone. The finally's reset is + # guarded on cancel_event being unset, so it will not run + # again for this path. + backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) if chunk is _DONE: @@ -9726,6 +9768,14 @@ async def _openai_passthrough_stream( ) _tracker.__exit__(None, None, None) + async def _unstarted_cleanup() -> None: + # Client disconnected before the body stream started, so _stream()'s + # finally never ran. Release the eagerly-opened upstream resp/client + # and the cancel-registry entry here; the watchers and line iterator + # are created inside _stream(), so there is nothing else to close. + await _aclose_stream_resources(resp = resp, client = client) + _tracker.__exit__(None, None, None) + return _SameTaskStreamingResponse( _stream(), media_type = "text/event-stream", @@ -9734,6 +9784,7 @@ async def _openai_passthrough_stream( "Connection": "close", "X-Accel-Buffering": "no", }, + unstarted_cleanup = _unstarted_cleanup, ) except BaseException: _tracker.__exit__(None, None, None) 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 411e6de1f3..8df8d37a52 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -97,6 +97,16 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_nested_gemma_marker_in_unquoted_arg_does_not_run_inner_call(): + # An UNQUOTED Gemma value containing a literal marker: the outer object fails + # to normalize (the inner braces/marker break the JSON), but the inner marker + # is nested in the outer candidate span, so it must not be promoted to a + # standalone `terminal` call. The safe outcome is no executed tool call. + content = "<|tool_call>call:python{code:<|tool_call>call:terminal{command:ls}}" + calls = parse_tool_calls_from_text(content) + assert "terminal" not in [c["function"]["name"] for c in calls], 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. From 75b65b48345ba8938af0d0b923bef6bb13c430be Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 22 Jun 2026 16:24:45 +0000 Subject: [PATCH 29/30] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/routes/inference.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 51700e4017..5f467cc83a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -808,7 +808,12 @@ def _same_task_timeout(timeout_s: float): class _SameTaskStreamingResponse(StreamingResponse): """StreamingResponse without Starlette's legacy AnyIO task-group wrapper.""" - def __init__(self, *args, unstarted_cleanup = None, **kwargs) -> None: + def __init__( + self, + *args, + unstarted_cleanup = None, + **kwargs, + ) -> None: super().__init__(*args, **kwargs) # Async callable invoked when the client disconnects before the body # iterator is ever advanced. A generator that never started cannot run From 35e18e7680f97eb3f067f278e5539b7ee37bb986 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 23 Jun 2026 13:28:49 +0000 Subject: [PATCH 30/30] Quote-aware Gemma strip, symmetric unstarted cleanup, ReDoS anchor Address review findings on the tool-strip and streaming paths: - strip_tool_call_markup stripped Gemma-native spans with a plain regex that stops at the first , so a literal close marker inside a <|"|>-quoted argument truncated the span and leaked its suffix into visible text. A brace/quote-aware _strip_gemma_native_spans now removes complete spans (keeping an incomplete one unless final), matching the parser's own balance logic. - The Gemma close pattern this PR added (<\|tool_call>.*?) had no \Z fallback, so a run of unclosed markers backtracked from every open position (quadratic, and the streaming stripper re-scans per token). It is now anchored to (?:|\Z) like routes/inference.py's _TOOL_XML_RE, linear with identical output on well-formed input. - _SameTaskStreamingResponse added unstarted_cleanup for the OpenAI passthrough, but the local GGUF/safetensors streams that enter _TrackedCancel before returning only unregister in the generator finally, which never runs if the client disconnects before the body iterator starts, leaking cancel-registry entries. Each such stream now passes unstarted_cleanup to exit its tracker. - __call__ reads _unstarted_cleanup via getattr so a response built through __new__ (the cancel-timing test) without __init__ does not raise AttributeError; the test also sets the attribute explicitly. - Document that the verbatim /v1/chat/completions passthrough delegates /<|tool_call> splitting to llama-server (--jinja, --reasoning-format auto) and is intentionally not re-parsed locally, noting the llama.cpp dependency. Adds a regression test for the close-marker-inside-quoted-argument strip. --- studio/backend/core/tool_healing.py | 56 ++++++++++++++++++- studio/backend/routes/inference.py | 40 ++++++++++++- .../tests/test_gemma_tool_parse_edge_cases.py | 13 +++++ .../test_stream_cancel_registration_timing.py | 1 + 4 files changed, 106 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..fe8b94a659 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -13,15 +13,25 @@ import re # Pre-compiled patterns for tool XML stripping. The hyphen in the name # char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, # issue-number) parse alongside the built-ins. +# +# The Gemma close marker is anchored to ``(?:|\Z)`` (the safe form +# routes/inference.py's _TOOL_XML_RE uses): the plain ``<\|tool_call>.*?`` +# this PR introduced backtracks from every open position on a run of unclosed +# markers (quadratic, and strip_tool_markup_streaming re-scans the cumulative +# buffer per token), whereas the ``\Z`` alternative lets the first open consume +# to EOF in one linear pass. strip_tool_call_markup additionally strips Gemma +# spans via the brace/quote-aware _strip_gemma_native_spans, so a literal close +# marker inside a <|"|>-quoted argument cannot truncate the span and leak its +# suffix; the regex below is the streaming-stripper fallback. +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), + _TC_GEMMA_CLOSED_PAT, re.compile(r""), re.compile(r".*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), - re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] @@ -443,6 +453,41 @@ def parse_tool_calls_from_text( return tool_calls +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native ``<|tool_call>call:NAME{...}`` + spans, brace- and quote-balanced so a literal ```` inside a + ``<|"|>``-quoted argument does not truncate the span and leak its suffix + (which the plain ``.*?`` regex does). A span without a balanced closing + ``}`` or a trailing close marker is incomplete: dropped to EOF when + ``final`` (the response is over), otherwise kept verbatim so a call that is + still streaming is not stripped mid-token. + """ + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + if final: + out.append(text[cursor:start]) + cursor = len(text) + continue + tail = text[brace_end + 1 :] + leading_ws = len(tail) - len(tail.lstrip()) + close = _TC_GEMMA_END_TAG_RE.match(tail, leading_ws) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + continue + out.append(text[cursor:start]) + cursor = brace_end + 1 + close.end() + out.append(text[cursor:]) + return "".join(out) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -450,7 +495,14 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. """ + # Gemma-native spans are stripped brace/quote-aware first; the regex form is + # not quote-aware and would truncate a span at a close marker inside a quoted + # argument. Skip that regex below and let the remaining patterns handle the + # JSON/XML formats and any orphan close marker. + text = _strip_gemma_native_spans(text, final = final) patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in patterns: + if pat is _TC_GEMMA_CLOSED_PAT: + continue text = pat.sub("", text) return text.strip() if final else text diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a90201803a..01e60be873 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -863,9 +863,13 @@ class _SameTaskStreamingResponse(StreamingResponse): aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - if self._unstarted_cleanup is not None: + # getattr (not self._unstarted_cleanup) so a response built via + # __new__ (some tests, pickling) without __init__ does not raise + # AttributeError here. + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -873,6 +877,20 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """Build an ``unstarted_cleanup`` for a local stream that entered ``tracker`` + (a ``_TrackedCancel``) before returning the response. The generator exits the + tracker in its ``finally``, but that never runs if the client disconnects + before the body iterator starts, leaking the cancel-registry entry. This + exits the tracker on that pre-start path only (mutually exclusive with the + generator's finally, so it never double-exits).""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -4953,6 +4971,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5422,6 +5441,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5571,6 +5591,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5958,6 +5979,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6150,6 +6172,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -9493,6 +9516,19 @@ async def _openai_passthrough_stream( response ``id``, ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so the client sees a standard OpenAI response. + + Reasoning/tool-call extraction here is delegated to llama-server: this path + forwards to its ``/v1/chat/completions`` (Studio launches with ``--jinja`` + and ``--reasoning-format auto``), which parses Gemma-native ```` into + ``reasoning_content`` and ``<|tool_call>`` into structured ``tool_calls`` + server-side, so the relayed ``delta.content`` carries no raw markup. This is + deliberately NOT re-parsed with the local reasoning extractor / Gemma parser + (verified end to end on the current llama.cpp build), unlike Studio's own + ``/completion``-level generation paths, which must parse the raw text + themselves. The dependency is on llama.cpp's chat parser: if a future build + or chat template stops splitting ````/``<|tool_call>``, raw markup + would relay into ``content`` and this path would need the local extractor as + a safety net. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( 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 8df8d37a52..d573522bcc 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -22,6 +22,7 @@ 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 +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -159,3 +160,15 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # A literal inside a <|"|>-quoted argument must not truncate the + # span: the parser keeps it as data, and stripping must remove the whole span + # (brace/quote-aware), not stop at the inner marker and leak the suffix. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 33deb7af9d..73e60a5b0f 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -411,6 +411,7 @@ def test_same_task_response_closes_body_iterator_on_send_disconnect(): response = m["_SameTaskStreamingResponse"].__new__(m["_SameTaskStreamingResponse"]) response.body_iterator = agen response.background = None + response._unstarted_cleanup = None async def stream_response(_send): raise OSError("client disconnected")