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