From b0dbe438671a0d01be1f88e54fb475ebd433fff2 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 22 Jun 2026 16:24:09 +0000 Subject: [PATCH] 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.