From 9d53b00d1d0315d3ab2ca0e6e53ef5ef6c59d764 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 23 Jun 2026 13:28:49 +0000 Subject: [PATCH 01/21] 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 fe26d48c7fa..fe8b94a659c 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -13,15 +13,25 @@ # 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 2e2c38933e5..5006e88d9fd 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -863,9 +863,13 @@ async def _tracking_send(message) -> None: 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 @@ async def _tracking_send(message) -> None: 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 audio_input_stream(): return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5427,6 +5446,7 @@ def _flush_reasoning_extractor(): return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5576,6 +5596,7 @@ async def gguf_stream_chunks(): return _SameTaskStreamingResponse( gguf_stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5963,6 +5984,7 @@ async def sf_tool_stream(): 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", @@ -6155,6 +6177,7 @@ async def stream_chunks(): return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -9498,6 +9521,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 8df8d37a522..d573522bcc2 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 @@ 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 33deb7af9d6..73e60a5b0f7 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -411,6 +411,7 @@ async def run(): 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") From f2c578ffdece342d4b4ce9e10b8e99cbf36dfe6d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 23 Jun 2026 14:05:35 +0000 Subject: [PATCH 02/21] Tighten comments on the tool-strip and streaming paths Compress the verbose comment blocks added with the Gemma tool-call / streaming work to crisp one or two liners, drop restatements of obvious code, and shorten docstrings, keeping the load-bearing rationale (ReDoS anchor, quote-aware strip, unstarted-cleanup, llama.cpp passthrough dependency). Code is unchanged (verified comment-only via AST/ast signature, docstrings stripped). --- studio/backend/core/tool_healing.py | 87 +++++++--------------- studio/backend/routes/inference.py | 109 ++++++++++------------------ 2 files changed, 64 insertions(+), 132 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe8b94a659c..081d933ea55 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -10,19 +10,11 @@ import json 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. +# Tool XML stripping patterns. Hyphen in the name class matches dashed MCP names +# (mcp__srv__list-issues). The Gemma marker is anchored to (?:|\Z) +# (like _TOOL_XML_RE) so an unclosed run strips linearly, not quadratically; +# strip_tool_call_markup uses the quote-aware _strip_gemma_native_spans instead, +# leaving this regex as the streaming fallback. _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), @@ -47,12 +39,9 @@ _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. 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. +# A bare Gemma value ends at `}` or a comma starting the next `key:` pair. The +# key must be identifier-shaped, so a comma before digits-then-colon stays in +# the value (`location:New York, NY`; `meet at 10:00, 11:00`). _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") @@ -149,14 +138,10 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """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.""" + """Normalise a Gemma array value so json.loads succeeds: quote bare string + elements, recurse into object/nested-array elements, leave quoted strings, + numbers, and JSON literals as-is. Gemma emits these unquoted + (``labels:[bug,ui]``, ``items:[{path:a}]``) and the call would otherwise drop.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -164,11 +149,9 @@ def _quote_gemma_array_elements(body: str) -> str: 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]) + "]") @@ -245,15 +228,13 @@ 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. + # Quote bare string values ({unit:celsius}); JSON scalars/objects/ + # arrays/quoted strings 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] == "[": - # 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:]) @@ -263,9 +244,7 @@ def _quote_gemma_object_keys(src: str) -> str: 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. - # `New York, NY`) does not terminate it. + # Bare value: up to `}` or a comma that starts the next key:pair. while i < len(src): if src[i] == "}": break @@ -320,19 +299,11 @@ def parse_tool_calls_from_text( ... """ tool_calls: list[dict] = [] - # 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). + # Collect JSON- and Gemma-format candidates with byte spans, then accept in + # document order (tools run in returned order). A marker inside an open + # value is that parameter's data, so skip it. 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 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. + spans, brace/quote-balanced so a literal ```` inside a + ``<|"|>``-quoted argument cannot truncate the span and leak its suffix. An + incomplete span is dropped to EOF when ``final``, else kept (still streaming). """ out: list[str] = [] cursor = 0 @@ -495,10 +461,7 @@ 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. + # Gemma spans first (quote-aware); skip the non-quote-aware Gemma regex below. text = _strip_gemma_native_spans(text, final = final) patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in patterns: diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5006e88d9fd..5c222fc8050 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -815,17 +815,14 @@ def __init__( **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. + # Released when the client disconnects before the body iterator starts: + # its try/finally never runs, so a stream that opens resources before the + # first yield (the passthrough's upstream httpx stream) passes this. 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. + # send() emits a body message only after the first chunk, so no body + # message means the generator never entered its try/finally. body_started = False async def _tracking_send(message) -> None: @@ -836,15 +833,11 @@ async def _tracking_send(message) -> None: try: await self.stream_response(_tracking_send) - except OSError: - # Client disconnected mid-send. + except OSError: # 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. + # Generator is suspended in its try/finally: throw CancelledError + # (not aclose's GeneratorExit) so its handler finishes the + # api_monitor entry. Fall back to aclose() without athrow. athrow = getattr(self.body_iterator, "athrow", None) if athrow is not None: try: @@ -856,16 +849,12 @@ async def _tracking_send(message) -> 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. + # Generator never started; aclose()/athrow() are no-ops on it, so + # release eager resources via the hook. getattr guards a response + # built through __new__ without __init__ (tests, pickling). aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - # 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: @@ -878,12 +867,9 @@ async def _tracking_send(message) -> None: 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).""" + """unstarted_cleanup for a local stream that entered ``tracker`` before + returning: exits it on pre-start disconnect, when the generator's finally + (which normally does so) never runs. Mutually exclusive with that finally.""" async def _cleanup() -> None: tracker.__exit__(None, None, None) @@ -3478,12 +3464,10 @@ async def 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. + # Watcher set cancel_event between chunks. Reset here: + # closing the generator does not signal a subprocess backend, + # so it would keep decoding. The finally's reset is guarded on + # cancel_event being unset, so it will not double-run. backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) @@ -8547,11 +8531,9 @@ async def _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. + # Watcher to cancel on disconnect: the in-loop poll only fires between + # events, so a disconnect during a long prefill/step would otherwise hold + # the decode slot until the next event or a failed send. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -8643,11 +8625,9 @@ async def _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. + # Watcher to cancel on disconnect: the in-loop poll only fires between + # chunks, so a disconnect during a long prefill/step would otherwise hold + # the decode slot until the next chunk or a failed send. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -9522,18 +9502,12 @@ async def _openai_passthrough_stream( ``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. + Reasoning/tool-call splitting is delegated to llama-server's + ``/v1/chat/completions`` (Studio runs ``--jinja --reasoning-format auto``), + so ``delta.content`` carries no raw ````/``<|tool_call>`` markup and + is deliberately not re-parsed locally (unlike the ``/completion`` paths). If + a future llama.cpp build stops splitting it, 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( @@ -9720,11 +9694,9 @@ def _synthetic_finish_line() -> str: 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. + # Detect an error chunk independently of API monitoring + # (skip_api_monitor returns early), else the synthetic + # finish would fire after a failed stream. if _monitor_openai_error_message(chunk_data): saw_stream_error = True monitor_event = _monitor_openai_sse_line( @@ -9760,10 +9732,8 @@ def _synthetic_finish_line() -> str: monitor_done = True break 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. + # Synthesize finish only if not already emitted (e.g. before + # a trailing usage chunk), but always close with [DONE]. if not saw_finish_reason: finish_line = _synthetic_finish_line() _monitor_openai_sse_line( @@ -9816,10 +9786,9 @@ def _synthetic_finish_line() -> str: _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. + # Disconnect before the stream started: _stream()'s finally never + # ran, so release the eager resp/client and cancel entry here (its + # watchers/line iterator are created inside _stream()). await _aclose_stream_resources(resp = resp, client = client) _tracker.__exit__(None, None, None) From 3215482668a86a63eb1476473c86252e2d85efcb Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 24 Jun 2026 01:06:06 +0000 Subject: [PATCH 03/21] Harden Gemma parse/strip: span-aware XML fallback and quote-aware streaming - Security: the XML fallback in parse_tool_calls_from_text scanned the whole content for markers and only skipped those inside an open XML parameter, not those inside a collected JSON/Gemma candidate span. A balanced but unparsable Gemma call whose argument data contained XML tool markup (<|tool_call>call:outer{code:...}) therefore fell through to the fallback and returned an executable terminal call. The fallback now also excludes markers inside any candidate span, including ones that failed to parse. - strip_tool_call_markup no longer skips the generic Gemma regex after running the quote-aware _strip_gemma_native_spans, so a closed Gemma span the helper cannot match (malformed, e.g. <|tool_call>{"name":"x"}) is still stripped instead of leaking its opener and payload into visible text. - _strip_gemma_native_spans stops at the first unbalanced start instead of re-scanning every later start to EOF, keeping it linear on a run of unclosed markers rather than quadratic. - The GGUF and safetensors streaming strippers run _strip_gemma_native_spans before the regex patterns, so a well-formed streamed call whose quoted argument contains a literal close marker no longer leaks its suffix into incremental display. Adds regression tests for the nested-XML escape and the malformed-span strip. --- studio/backend/core/inference/llama_cpp.py | 4 ++++ .../core/inference/safetensors_agentic.py | 4 ++++ .../core/inference/tool_call_parser.py | 1 + studio/backend/core/tool_healing.py | 15 +++++++++---- .../tests/test_gemma_tool_parse_edge_cases.py | 22 +++++++++++++++++++ 5 files changed, 42 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8a7fad9af4f..7b245daa559 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -40,6 +40,7 @@ ) from core.tool_healing import ( _TOOL_ALL_PATS, + _strip_gemma_native_spans, strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret @@ -7780,6 +7781,9 @@ def _strip_tool_markup( def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text + # Quote-aware Gemma spans first, else a literal inside a + # quoted argument truncates the regex match and leaks the suffix. + text = _strip_gemma_native_spans(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 0c96378d6c9..eaed55507ae 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -22,6 +22,7 @@ from core.inference.tool_call_parser import ( _TOOL_ALL_PATS, + _strip_gemma_native_spans, BUDGET_EXHAUSTED_NUDGE, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, @@ -60,6 +61,9 @@ def strip_tool_markup_streaming( """Strip open-ended tool XML from display text without trimming whitespace.""" if not (auto_heal_tool_calls or tool_protocol_active): return text + # Quote-aware Gemma spans first, else a literal inside a quoted + # argument truncates the regex match and leaks the suffix into display. + text = _strip_gemma_native_spans(text, final = True) for pat in _TOOL_ALL_PATS: text = pat.sub("", text) return text diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index ca3d1e4cbc0..2302fa2bcf8 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -11,6 +11,7 @@ _TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS +_strip_gemma_native_spans = _tool_healing._strip_gemma_native_spans def parse_tool_calls_from_text( diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 081d933ea55..9a8a41d1451 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -350,10 +350,15 @@ def parse_tool_calls_from_text( ) if not tool_calls: + # Exclude markers inside any collected JSON/Gemma candidate + # span, even one that failed to parse: such a marker is that call's + # argument data, and promoting it here would let nested XML escape from a + # malformed Gemma call into an executable tool call. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) + and not any(s <= fm.start() <= e for s, e in spans) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -436,10 +441,13 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: continue brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) if brace_end < 0: + # Unbalanced: no complete span from here on. Drop the rest if final, + # else keep it, and stop -- later starts are inside this unclosed run, + # so rescanning them would re-walk to EOF each time (quadratic). if final: out.append(text[cursor:start]) cursor = len(text) - continue + break tail = text[brace_end + 1 :] leading_ws = len(tail) - len(tail.lstrip()) close = _TC_GEMMA_END_TAG_RE.match(tail, leading_ws) @@ -461,11 +469,10 @@ 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 spans first (quote-aware); skip the non-quote-aware Gemma regex below. + # Well-formed Gemma spans first (quote-aware); the regex patterns below then + # mop up any malformed Gemma span the helper could not match. 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/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index d573522bcc2..4df046ff31a 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -172,3 +172,25 @@ def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): 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" + + +def test_nested_xml_in_malformed_gemma_call_does_not_execute(): + # A balanced but unparsable Gemma call whose argument data contains XML tool + # markup must not let that escape into an executable call via the + # XML fallback (the Gemma candidate span covers it even though it failed). + text = ( + "<|tool_call>call:outer{code:id" + ", broken:{x}}" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_malformed_closed_gemma_span_is_stripped(): + # A closed Gemma span the quote-aware helper cannot match (no call:NAME{) + # must still be stripped, not leak its opener/payload into visible text. + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) From a06bbd564ce03746cda3b6759ceba57dc0378e28 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 25 Jun 2026 04:08:18 +0000 Subject: [PATCH 04/21] Avoid remainder copy in _strip_gemma_native_spans Match the Gemma close marker with re pos directly on the buffer instead of slicing tail = text[brace_end + 1:] on every span. The streaming strippers re-scan a growing cumulative buffer per token, so the per-span remainder copy was quadratic. Behavior is unchanged. --- studio/backend/core/tool_healing.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 9a8a41d1451..4046880b9d8 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -448,16 +448,19 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: out.append(text[cursor:start]) cursor = len(text) break - tail = text[brace_end + 1 :] - leading_ws = len(tail) - len(tail.lstrip()) - close = _TC_GEMMA_END_TAG_RE.match(tail, leading_ws) + # Match the close marker via re pos (no remainder copy): streaming + # re-scans a growing buffer per token, so slicing here is quadratic. + close_idx = brace_end + 1 + while close_idx < len(text) and text[close_idx].isspace(): + close_idx += 1 + close = _TC_GEMMA_END_TAG_RE.match(text, close_idx) 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() + cursor = close.end() out.append(text[cursor:]) return "".join(out) From 75ce5bb569ab2d5f0d81f9a7a28ec62b16331ed5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 25 Jun 2026 12:31:07 +0000 Subject: [PATCH 05/21] Exclude unclosed Gemma/JSON starts from the XML tool-call fallback The nested-XML guard only skipped markers inside recorded candidate spans, but a span is recorded only when the braces balance. An unbalanced call such as <|tool_call>call:outer{code:... recorded no span, so the fallback still promoted the inner to an executable terminal call. Treat unclosed JSON/Gemma starts as exclusion spans through EOF before scanning. Standalone calls with no preceding unclosed start still parse. Regression tests added. --- studio/backend/core/tool_healing.py | 16 ++++++++++---- .../tests/test_gemma_tool_parse_edge_cases.py | 21 +++++++++++++++++++ 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 4046880b9d8..784251502d4 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -303,18 +303,25 @@ def parse_tool_calls_from_text( # document order (tools run in returned order). A marker inside an open # value is that parameter's data, so skip it. candidates = [] # (start, brace_end, kind, match) + # Unclosed starts (braces never balance): everything after is the call's + # argument data through EOF, so the XML fallback must treat it as excluded. + unclosed_starts = [] for m in _TC_JSON_START_RE.finditer(content): if _inside_open_parameter(content, m.start()): continue end = _balanced_brace_end(content, m.end() - 1) if end >= 0: candidates.append((m.start(), end, "json", m)) + else: + unclosed_starts.append(m.start()) 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)) + else: + unclosed_starts.append(m.start()) candidates.sort(key = lambda c: c[0]) spans = [(s, e) for s, e, _kind, _m in candidates] @@ -351,14 +358,15 @@ def parse_tool_calls_from_text( if not tool_calls: # Exclude markers inside any collected JSON/Gemma candidate - # span, even one that failed to parse: such a marker is that call's - # argument data, and promoting it here would let nested XML escape from a - # malformed Gemma call into an executable tool call. + # span (its argument data, even if it failed to parse) and inside any + # unclosed start through EOF; otherwise nested XML in a malformed call + # escapes into an executable tool call. + exclusion_spans = spans + [(s, len(content)) for s in unclosed_starts] func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in spans) + and not any(s <= fm.start() <= e for s, e in exclusion_spans) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) 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 4df046ff31a..a76edd769bd 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -187,6 +187,27 @@ def test_nested_xml_in_malformed_gemma_call_does_not_execute(): assert "terminal" not in [c["function"]["name"] for c in calls], calls +def test_unbalanced_gemma_call_with_xml_does_not_execute(): + # An unbalanced Gemma call (braces never close) records no candidate span, + # so the XML fallback must exclude its trailing through EOF + # rather than promote it to an executable call. + text = ( + "<|tool_call>call:outer{code:" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_standalone_function_xml_still_parses(): + # The exclusion must not over-block: a real call with no + # preceding unclosed Gemma/JSON start is still a valid tool call. + text = "id" + calls = parse_tool_calls_from_text(text) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + def test_malformed_closed_gemma_span_is_stripped(): # A closed Gemma span the quote-aware helper cannot match (no call:NAME{) # must still be stripped, not leak its opener/payload into visible text. From 75a5258c93be79c76d34f3df1f3a1cdd80ba4205 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 27 Jun 2026 03:06:21 +0000 Subject: [PATCH 06/21] Skip doomed tool-strip passes to avoid quadratic rescans The lazy closed-pair strip patterns (.*?, .*?) rescan to EOF from every opener when their close token is absent, which is O(n^2) and re-runs per streamed token. Add strip_tool_patterns, which skips a pass whose close token is not present in the text; output is identical to the per-pattern loop (verified by fuzz), and a degenerate run drops from ~minutes to milliseconds. Used by strip_tool_call_markup and the GGUF/safetensors streaming strippers. --- studio/backend/core/inference/llama_cpp.py | 5 +- .../core/inference/safetensors_agentic.py | 5 +- .../core/inference/tool_call_parser.py | 1 + studio/backend/core/tool_healing.py | 27 ++++++-- studio/backend/tests/test_tool_strip_guard.py | 64 +++++++++++++++++++ 5 files changed, 92 insertions(+), 10 deletions(-) create mode 100644 studio/backend/tests/test_tool_strip_guard.py diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index dd38363081d..3613322d863 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -42,6 +42,7 @@ _TOOL_ALL_PATS, _strip_gemma_native_spans, strip_tool_call_markup, + strip_tool_patterns, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -7864,9 +7865,7 @@ def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: # Quote-aware Gemma spans first, else a literal inside a # quoted argument truncates the regex match and leaks the suffix. text = _strip_gemma_native_spans(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + return strip_tool_patterns(text, _TOOL_ALL_PATS) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index eaed55507ae..11c98dbc455 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -29,6 +29,7 @@ TOOL_XML_SIGNALS, parse_tool_calls_from_text, strip_tool_markup, + strip_tool_patterns, ) from core.inference.tool_loop_controller import ( ToolLoopController, @@ -64,9 +65,7 @@ def strip_tool_markup_streaming( # Quote-aware Gemma spans first, else a literal inside a quoted # argument truncates the regex match and leaks the suffix into display. text = _strip_gemma_native_spans(text, final = True) - for pat in _TOOL_ALL_PATS: - text = pat.sub("", text) - return text + return strip_tool_patterns(text, _TOOL_ALL_PATS) def _strip_tool_markup_final( diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index 2302fa2bcf8..cfc30ea575e 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -12,6 +12,7 @@ _TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS _strip_gemma_native_spans = _tool_healing._strip_gemma_native_spans +strip_tool_patterns = _tool_healing.strip_tool_patterns def parse_tool_calls_from_text( diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 784251502d4..962cf85af6e 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -15,17 +15,37 @@ # (like _TOOL_XML_RE) so an unclosed run strips linearly, not quadratically; # strip_tool_call_markup uses the quote-aware _strip_gemma_native_spans instead, # leaving this regex as the streaming fallback. +_TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) +_TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, _TC_GEMMA_CLOSED_PAT, re.compile(r""), - re.compile(r".*?", re.DOTALL), + _TC_FUNC_CLOSED_PAT, ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +# A lazy closed-pair pattern rescans to EOF from every opener when its close +# token is absent (O(n^2); O(n^3) re-run per streamed token). Skip the doomed +# sweep when the token is missing -- identical output, no backtracking. +_PAT_REQUIRED_TOKEN = { + _TC_JSON_CLOSED_PAT: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """Apply strip ``patterns`` in order, skipping a closed-pair pass whose close + token is absent (avoids its quadratic no-match rescan).""" + for pat in patterns: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + text = pat.sub("", text) + return text # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") @@ -484,6 +504,5 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: # mop up any malformed Gemma span the helper could not match. text = _strip_gemma_native_spans(text, final = final) patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) + text = strip_tool_patterns(text, patterns) return text.strip() if final else text diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py new file mode 100644 index 00000000000..ee69f348321 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""strip_tool_patterns must produce identical output to the plain per-pattern +loop, while skipping a lazy closed-pair sweep whose close token is absent (the +O(n^2) no-match rescan, O(n^3) when re-run per streamed token).""" + +import random +import sys +import time +from pathlib import Path + +_BACKEND_ROOT = Path(__file__).resolve().parents[1] +if str(_BACKEND_ROOT) not in sys.path: + sys.path.insert(0, str(_BACKEND_ROOT)) + +from core.tool_healing import ( + _TOOL_ALL_PATS, + _TOOL_CLOSED_PATS, + strip_tool_call_markup, + strip_tool_patterns, +) + + +def _naive(text, patterns): + for pat in patterns: + text = pat.sub("", text) + return text + + +_TOKENS = [ + "", "", "<|tool_call>", "", + "", "", "", + "", "", "call:fn{", "}", "{", '<|"|>', + "A", " ", "\n", "id", "x:1", "", +] + + +def test_guard_matches_plain_loop_on_fuzz(): + rng = random.Random(1234) + for patterns in (_TOOL_ALL_PATS, _TOOL_CLOSED_PATS): + for _ in range(20000): + s = "".join(rng.choice(_TOKENS) for _ in range(rng.randint(0, 10))) + assert strip_tool_patterns(s, patterns) == _naive(s, patterns), (s, patterns) + + +def test_strip_markup_representative_cases_unchanged(): + assert strip_tool_call_markup("a {} b") == "a b" + assert ( + strip_tool_call_markup("a 1 b") == "a b" + ) + # Non-final keeps an unclosed block; final strips it to EOF. + assert strip_tool_call_markup("a {partial") == "a {partial" + assert strip_tool_call_markup("a {partial", final = True) == "a" + + +def test_no_quadratic_blowup_on_unclosed_markers(): + # Many openers with no close token: the guard skips the lazy closed-pair + # sweep, keeping this linear. Unguarded this took minutes. + big = "" * 20000 + "" * 20000 + t0 = time.perf_counter() + out = strip_tool_call_markup(big, final = True) + assert time.perf_counter() - t0 < 2.0 + assert out == "" From 939614efbd2dc33778233272a4ebec65c9b87c14 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 27 Jun 2026 03:07:38 +0000 Subject: [PATCH 07/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/tool_healing.py | 1 + studio/backend/tests/test_tool_strip_guard.py | 28 ++++++++++++++----- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 962cf85af6e..faa7fbc325d 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -47,6 +47,7 @@ def strip_tool_patterns(text: str, patterns) -> str: text = pat.sub("", text) return text + # 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*\{") diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py index ee69f348321..ef5bdaf769e 100644 --- a/studio/backend/tests/test_tool_strip_guard.py +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -29,10 +29,26 @@ def _naive(text, patterns): _TOKENS = [ - "", "", "<|tool_call>", "", - "", "", "", - "", "", "call:fn{", "}", "{", '<|"|>', - "A", " ", "\n", "id", "x:1", "", + "", + "", + "<|tool_call>", + "", + "", + "", + "", + "", + "", + "call:fn{", + "}", + "{", + '<|"|>', + "A", + " ", + "\n", + "id", + "x:1", + "", ] @@ -46,9 +62,7 @@ def test_guard_matches_plain_loop_on_fuzz(): def test_strip_markup_representative_cases_unchanged(): assert strip_tool_call_markup("a {} b") == "a b" - assert ( - strip_tool_call_markup("a 1 b") == "a b" - ) + assert strip_tool_call_markup("a 1 b") == "a b" # Non-final keeps an unclosed block; final strips it to EOF. assert strip_tool_call_markup("a {partial") == "a {partial" assert strip_tool_call_markup("a {partial", final = True) == "a" From 2bd13b9a95570e99329d1e307084ec8b99b3e21f Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 27 Jun 2026 07:12:44 +0000 Subject: [PATCH 08/21] Use full tool-call envelopes to close nested-XML escape variants Key the parser and stripper off the full <|tool_call>... / ... envelope (start to close marker, searched after the braces; EOF if unclosed) instead of just the braces: - XML between the closing brace and the close marker (call:outer{broken:{x}}...) is now inside the envelope, so the fallback no longer promotes it to a tool call. - A balanced inner call inside an unclosed outer (call:outer{code:<|tool_call>call:terminal{...}) is skipped via the envelope nested check, not just the XML fallback. - strip_tool_call_markup searches for the close marker after the braces, so junk before is stripped through the close and text after it is preserved instead of truncated to EOF; a no-close run stops early (linear). Regression tests added; standalone XML and well-formed calls unaffected. --- studio/backend/core/tool_healing.py | 92 +++++++++---------- .../tests/test_gemma_tool_parse_edge_cases.py | 30 ++++++ 2 files changed, 75 insertions(+), 47 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index faa7fbc325d..0086f0f0dec 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -320,53 +320,54 @@ def parse_tool_calls_from_text( ... """ tool_calls: list[dict] = [] - # Collect JSON- and Gemma-format candidates with byte spans, then accept in - # document order (tools run in returned order). A marker inside an open - # value is that parameter's data, so skip it. - candidates = [] # (start, brace_end, kind, match) - # Unclosed starts (braces never balance): everything after is the call's - # argument data through EOF, so the XML fallback must treat it as excluded. - unclosed_starts = [] - for m in _TC_JSON_START_RE.finditer(content): - if _inside_open_parameter(content, m.start()): - continue - end = _balanced_brace_end(content, m.end() - 1) - if end >= 0: - candidates.append((m.start(), end, "json", m)) - else: - unclosed_starts.append(m.start()) - 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)) - else: - unclosed_starts.append(m.start()) - candidates.sort(key = lambda c: c[0]) - - spans = [(s, e) for s, e, _kind, _m in candidates] - for idx, (start, end, kind, m) in enumerate(candidates): - # Skip a candidate nested in another candidate's span: it is the outer - # call's argument data. Checked against every span (not just parsed - # ones), so a marker in an outer call that fails to parse is never run. - if any(s <= start and end <= e for j, (s, e) in enumerate(spans) if j != idx): + # Collect JSON/Gemma markers with their full envelope [start, env_end): a + # balanced call runs to its close marker (searched after the braces, since + # junk can sit between } and the marker), an unbalanced/unclosed one runs to + # EOF. A marker starting inside another's envelope is that call's argument + # data, never its own call. A marker inside an open + # value is that parameter's data, so skip it. + markers = [] # (start, brace_end, env_end, kind, match); brace_end < 0 = unclosed + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE + for m in start_re.finditer(content): + if _inside_open_parameter(content, m.start()): + continue + brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) + if brace_end < 0: + env_end = len(content) + else: + cm = close_re.search(content, brace_end + 1) + env_end = cm.end() if cm else len(content) + markers.append((m.start(), brace_end, env_end, kind, m)) + markers.sort(key = lambda c: c[0]) + + envelopes = [(s, e) for s, _be, e, _k, _m in markers] + for idx, (start, brace_end, env_end, kind, m) in enumerate(markers): + # Skip a marker nested in another marker's envelope -- it is the outer + # call's argument data (covers an unclosed outer, whose envelope is EOF), + # so a nested or post-brace marker is never run as its own call. + if any(s <= start < e for j, (s, e) in enumerate(envelopes) if j != idx): continue + if brace_end < 0: + continue # unclosed: not parseable; its envelope still excludes XML below if not allow_incomplete: - tail = content[end + 1 :].lstrip() + tail = content[brace_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 try: if kind == "json": - obj = json.loads(content[m.end() - 1 : end + 1]) + obj = json.loads(content[m.end() - 1 : brace_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])) + arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue tool_calls.append( @@ -378,16 +379,14 @@ def parse_tool_calls_from_text( ) if not tool_calls: - # Exclude markers inside any collected JSON/Gemma candidate - # span (its argument data, even if it failed to parse) and inside any - # unclosed start through EOF; otherwise nested XML in a malformed call - # escapes into an executable tool call. - exclusion_spans = spans + [(s, len(content)) for s in unclosed_starts] + # Exclude markers inside any marker's envelope (its argument + # data, even if the call failed to parse), so nested XML in a malformed + # call cannot escape into an executable tool call. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in exclusion_spans) + and not any(s <= fm.start() < e for s, e in envelopes) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -477,17 +476,16 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: out.append(text[cursor:start]) cursor = len(text) break - # Match the close marker via re pos (no remainder copy): streaming - # re-scans a growing buffer per token, so slicing here is quadratic. - close_idx = brace_end + 1 - while close_idx < len(text) and text[close_idx].isspace(): - close_idx += 1 - close = _TC_GEMMA_END_TAG_RE.match(text, close_idx) + # Search for the close marker after the braces (not just immediately + # after): junk between } and is malformed-call markup, so + # strip through the close and keep any text after it. None anywhere after + # means nothing closes from here on, so stop (keeps this linear). + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) if close is None: if final: out.append(text[cursor:start]) cursor = len(text) - continue + break out.append(text[cursor:start]) cursor = close.end() out.append(text[cursor:]) 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 a76edd769bd..56a859b1b86 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -208,6 +208,36 @@ def test_standalone_function_xml_still_parses(): assert [c["function"]["name"] for c in calls] == ["terminal"], calls +def test_xml_between_braces_and_close_marker_does_not_execute(): + # Balanced-but-unparsable outer call with XML after the braces but before the + # close marker: the envelope runs to the close marker, so here is + # the outer call's data, not an executable tool call. + text = ( + "<|tool_call>call:outer{broken:{x}}" + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): + # A balanced inner call inside an unclosed outer call's argument data must be + # skipped, not accepted, even though its own braces balance. + text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_preserves_text_after_malformed_gemma_close(): + # A valid call:name{...} prefix with junk before its close is a + # malformed closed span: strip through the close, keep the text after it. + text = "pre <|tool_call>call:t{a:1} note post" + assert strip_tool_call_markup(text) == "pre post" + assert strip_tool_call_markup(text, final = True) == "pre post" + + def test_malformed_closed_gemma_span_is_stripped(): # A closed Gemma span the quote-aware helper cannot match (no call:NAME{) # must still be stripped, not leak its opener/payload into visible text. From 9cca79bd0f7511cdb520578a52c63da857c838ab Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sat, 27 Jun 2026 12:40:39 +0000 Subject: [PATCH 09/21] Fix non-final Gemma strip and missing-close recovery for PR #6611 Split the nested-skip from the XML fallback exclusion: nesting is decided by each marker's brace region, so a balanced call after one with a missing close marker is recovered instead of being swallowed to EOF. Only the XML fallback keeps the search-to-close envelope, so trailing nested markup still cannot escape as an executable call. Use a closed-only Gemma pattern in the non-final strip list so an incomplete block is preserved (matching the JSON and function paths); the final list keeps the close-or-EOF Gemma pattern in its original position, so streaming display output is byte-for-byte unchanged. Add regression tests for both cases. --- studio/backend/core/tool_healing.py | 77 +++++++++++-------- .../tests/test_gemma_tool_parse_edge_cases.py | 19 +++++ 2 files changed, 64 insertions(+), 32 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 0086f0f0dec..bd75c237f84 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -11,28 +11,38 @@ import re # Tool XML stripping patterns. Hyphen in the name class matches dashed MCP names -# (mcp__srv__list-issues). The Gemma marker is anchored to (?:|\Z) -# (like _TOOL_XML_RE) so an unclosed run strips linearly, not quadratically; -# strip_tool_call_markup uses the quote-aware _strip_gemma_native_spans instead, -# leaving this regex as the streaming fallback. +# (mcp__srv__list-issues). The non-final list uses a closed-only Gemma pattern so +# an incomplete block is preserved (matching JSON/function); the final list keeps +# the same pattern order but lets the Gemma span run to its close marker OR EOF, +# then sweeps any unclosed JSON/function remainder to EOF. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) -_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) +# Final-only variant: an unclosed Gemma run also strips to EOF. The \Z alternative +# makes the lazy match short-circuit on the first opener, so it stays linear. +_TC_GEMMA_CLOSED_OR_EOF_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) +_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ _TC_JSON_CLOSED_PAT, _TC_GEMMA_CLOSED_PAT, - re.compile(r""), + _TC_GEMMA_END_PAT, _TC_FUNC_CLOSED_PAT, ] -_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ +_TOOL_ALL_PATS = [ + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_OR_EOF_PAT, + _TC_GEMMA_END_PAT, + _TC_FUNC_CLOSED_PAT, re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] # A lazy closed-pair pattern rescans to EOF from every opener when its close # token is absent (O(n^2); O(n^3) re-run per streamed token). Skip the doomed -# sweep when the token is missing -- identical output, no backtracking. +# sweep when the token is missing -- identical output, no backtracking. The +# close-or-EOF variant needs no guard: its \Z short-circuits the first opener. _PAT_REQUIRED_TOKEN = { _TC_JSON_CLOSED_PAT: "", + _TC_GEMMA_CLOSED_PAT: "", _TC_FUNC_CLOSED_PAT: "", } @@ -320,39 +330,34 @@ def parse_tool_calls_from_text( ... """ tool_calls: list[dict] = [] - # Collect JSON/Gemma markers with their full envelope [start, env_end): a - # balanced call runs to its close marker (searched after the braces, since - # junk can sit between } and the marker), an unbalanced/unclosed one runs to - # EOF. A marker starting inside another's envelope is that call's argument - # data, never its own call. A marker inside an open - # value is that parameter's data, so skip it. - markers = [] # (start, brace_end, env_end, kind, match); brace_end < 0 = unclosed + # Collect JSON/Gemma markers. Nesting is decided by the brace region only -- + # [start, brace_end], or [start, EOF) when the braces never close -- so a + # marker starting inside another's braces is that call's argument data, while + # a later balanced call after one with a missing close is still recovered (not + # swallowed to EOF). The XML fallback instead excludes each marker's envelope + # (start to its close marker, searched after the braces, or EOF) so trailing + # markup before the close cannot escape. A marker inside an open + # value is that parameter's data, so skip it. + markers = [] # (start, brace_end, kind, match); brace_end < 0 = unclosed for start_re, gemma, kind in ( (_TC_JSON_START_RE, False, "json"), (_TC_GEMMA_START_RE, True, "gemma"), ): - close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE for m in start_re.finditer(content): if _inside_open_parameter(content, m.start()): continue brace_end = _balanced_brace_end(content, m.end() - 1, gemma_quotes = gemma) - if brace_end < 0: - env_end = len(content) - else: - cm = close_re.search(content, brace_end + 1) - env_end = cm.end() if cm else len(content) - markers.append((m.start(), brace_end, env_end, kind, m)) + markers.append((m.start(), brace_end, kind, m)) markers.sort(key = lambda c: c[0]) - envelopes = [(s, e) for s, _be, e, _k, _m in markers] - for idx, (start, brace_end, env_end, kind, m) in enumerate(markers): - # Skip a marker nested in another marker's envelope -- it is the outer - # call's argument data (covers an unclosed outer, whose envelope is EOF), - # so a nested or post-brace marker is never run as its own call. - if any(s <= start < e for j, (s, e) in enumerate(envelopes) if j != idx): + brace_regions = [(s, be if be >= 0 else len(content)) for s, be, _k, _m in markers] + for idx, (start, brace_end, kind, m) in enumerate(markers): + # Skip a marker whose start falls in another marker's brace region: it is + # that outer call's argument data, not its own call. + if any(s <= start <= e for j, (s, e) in enumerate(brace_regions) if j != idx): continue if brace_end < 0: - continue # unclosed: not parseable; its envelope still excludes XML below + continue # unclosed: not parseable; the fallback still excludes its XML if not allow_incomplete: tail = content[brace_end + 1 :].lstrip() close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE @@ -379,9 +384,17 @@ def parse_tool_calls_from_text( ) if not tool_calls: - # Exclude markers inside any marker's envelope (its argument - # data, even if the call failed to parse), so nested XML in a malformed - # call cannot escape into an executable tool call. + # Exclude markers inside any marker's envelope -- its data + # through the close marker (searched after the braces) or EOF, even if the + # call failed to parse -- so nested XML cannot escape as a real call. + envelopes = [] + for s, be, kind, _m in markers: + if be < 0: + envelopes.append((s, len(content))) + else: + close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE + cm = close_re.search(content, be + 1) + envelopes.append((s, cm.end() if cm else len(content))) func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) 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 56a859b1b86..a11a00af258 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -245,3 +245,22 @@ def test_malformed_closed_gemma_span_is_stripped(): strip_tool_call_markup('before <|tool_call>{"name":"x"} after') == "before after" ) + + +def test_valid_call_after_missing_close_is_recovered(): + # A balanced call missing its own close marker must not swallow a later valid + # closed call: nesting is decided by the brace region, not an envelope that + # would run to EOF, so the second call is still recovered. + text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" + names_inc = [c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)] + assert "b" in names_inc, names_inc + names_strict = [c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)] + assert names_strict == ["b"], names_strict + + +def test_strip_non_final_keeps_incomplete_gemma_block(): + # Non-final must preserve an incomplete Gemma block (matching JSON/function), + # while final strips the unclosed remainder to EOF. + text = "before <|tool_call>call:t{" + assert strip_tool_call_markup(text) == text + assert strip_tool_call_markup(text, final = True) == "before" From beff0f3dee22dab65208259d6cbeea13a719d7df Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 27 Jun 2026 12:41:15 +0000 Subject: [PATCH 10/21] [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, 6 insertions(+), 2 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 a11a00af258..a2c57766584 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -252,9 +252,13 @@ def test_valid_call_after_missing_close_is_recovered(): # closed call: nesting is decided by the brace region, not an envelope that # would run to EOF, so the second call is still recovered. text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" - names_inc = [c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True)] + names_inc = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) + ] assert "b" in names_inc, names_inc - names_strict = [c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False)] + names_strict = [ + c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = False) + ] assert names_strict == ["b"], names_strict From 3a55738325aa6e300a172a0bbb9ff63a6bf2eedb Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 1 Jul 2026 11:22:03 +0000 Subject: [PATCH 11/21] Block gap-nested tool markers and fix XML strip order for PR #6611 Decide candidate nesting by a per-marker coverage region paired with a per-format stack (a close after the braces pops the nearest still-open marker of that format). A closed outer call now covers up to its own close marker, so a JSON or Gemma tool marker smuggled between the outer braces and that close is treated as data instead of being executed. An outer that balances but has no close of its own covers only its brace region, so a later sibling after an omitted close marker is still recovered (adjacent calls use an exclusive end bound so the next call is not misread as nested). Strip every closed pair (JSON, Gemma, function) before any to-EOF sweep, so a closed function call whose parameter text contains a bare Gemma opener is removed as a unit and the to-EOF sweep can no longer drop the visible text after the close. Add regression tests for both. --- studio/backend/core/tool_healing.py | 85 +++++++++++++------ .../tests/test_gemma_tool_parse_edge_cases.py | 32 +++++++ 2 files changed, 92 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index bd75c237f84..177874aab64 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -11,35 +11,30 @@ import re # Tool XML stripping patterns. Hyphen in the name class matches dashed MCP names -# (mcp__srv__list-issues). The non-final list uses a closed-only Gemma pattern so -# an incomplete block is preserved (matching JSON/function); the final list keeps -# the same pattern order but lets the Gemma span run to its close marker OR EOF, -# then sweeps any unclosed JSON/function remainder to EOF. +# (mcp__srv__list-issues). Both lists strip every closed pair first (JSON, Gemma, +# function), so a closed call is removed as a unit before any to-EOF sweep can +# reach markup nested inside it; the final list then adds greedy .*$ sweeps that +# drop an unclosed opener's remainder to EOF. The non-final list keeps incomplete +# blocks (no EOF sweeps), matching the parser's allow_incomplete = False path. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) -# Final-only variant: an unclosed Gemma run also strips to EOF. The \Z alternative -# makes the lazy match short-circuit on the first opener, so it stays linear. -_TC_GEMMA_CLOSED_OR_EOF_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ _TC_JSON_CLOSED_PAT, _TC_GEMMA_CLOSED_PAT, - _TC_GEMMA_END_PAT, _TC_FUNC_CLOSED_PAT, -] -_TOOL_ALL_PATS = [ - _TC_JSON_CLOSED_PAT, - _TC_GEMMA_CLOSED_OR_EOF_PAT, _TC_GEMMA_END_PAT, - _TC_FUNC_CLOSED_PAT, +] +_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ + re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] # A lazy closed-pair pattern rescans to EOF from every opener when its close # token is absent (O(n^2); O(n^3) re-run per streamed token). Skip the doomed # sweep when the token is missing -- identical output, no backtracking. The -# close-or-EOF variant needs no guard: its \Z short-circuits the first opener. +# to-EOF sweeps are greedy, so they short-circuit on the first opener. _PAT_REQUIRED_TOKEN = { _TC_JSON_CLOSED_PAT: "", _TC_GEMMA_CLOSED_PAT: "", @@ -316,6 +311,45 @@ def _inside_open_parameter(content: str, pos: int) -> bool: return last_param_start > max(last_param_close, last_func_close) +def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """Coverage region ``[start, end]`` for each marker, used to skip markers that + are another call's data. Each marker's own close is paired to it with a + per-format stack (a close after the braces pops the nearest still-open marker + of that format), so an inner call's close is not mistaken for the outer's: + + - braces never balance -> covers ``[start, EOF)`` (the whole ambiguous tail); + - braces balance and a matching close is found -> covers ``[start, close_end]`` + so a marker smuggled between the braces and that close is treated as data; + - braces balance but no close is paired -> covers only the brace region, so a + later sibling after an omitted close marker is still recovered as its own call. + """ + n = len(content) + events = [] # (position, order) with order 0 = braces-done, 1 = close marker + for idx, (_start, brace_end, _kind, _m) in enumerate(markers): + if brace_end >= 0: + events.append((brace_end, 0, _kind, idx)) + for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): + for cm in close_re.finditer(content): + events.append((cm.start(), 1, kind, cm.end())) + events.sort(key = lambda e: (e[0], e[1])) + waiting = {"json": [], "gemma": []} + close_end_for: dict[int, int] = {} + for _pos, order, kind, payload in events: + if order == 0: + waiting[kind].append(payload) # marker index, now awaiting its close + elif waiting[kind]: + close_end_for[waiting[kind].pop()] = payload # innermost open marker closes here + coverage = [] + for idx, (start, brace_end, _kind, _m) in enumerate(markers): + if brace_end < 0: + coverage.append((start, n)) + elif idx in close_end_for: + coverage.append((start, close_end_for[idx])) + else: + coverage.append((start, brace_end)) + return coverage + + def parse_tool_calls_from_text( content: str, *, @@ -330,13 +364,12 @@ def parse_tool_calls_from_text( ... """ tool_calls: list[dict] = [] - # Collect JSON/Gemma markers. Nesting is decided by the brace region only -- - # [start, brace_end], or [start, EOF) when the braces never close -- so a - # marker starting inside another's braces is that call's argument data, while - # a later balanced call after one with a missing close is still recovered (not - # swallowed to EOF). The XML fallback instead excludes each marker's envelope - # (start to its close marker, searched after the braces, or EOF) so trailing - # markup before the close cannot escape. A marker inside an open + # Collect JSON/Gemma markers, then decide nesting by each marker's coverage + # region (see _marker_coverage): a closed outer covers up to its own close + # marker, so a JSON/Gemma marker smuggled between the outer braces and that + # close is treated as data, not executed; an outer that balances but has no + # close of its own covers only its brace region, so a later sibling after an + # omitted close is still recovered. A marker inside an open # value is that parameter's data, so skip it. markers = [] # (start, brace_end, kind, match); brace_end < 0 = unclosed for start_re, gemma, kind in ( @@ -350,11 +383,13 @@ def parse_tool_calls_from_text( markers.append((m.start(), brace_end, kind, m)) markers.sort(key = lambda c: c[0]) - brace_regions = [(s, be if be >= 0 else len(content)) for s, be, _k, _m in markers] + coverage = _marker_coverage(content, markers) for idx, (start, brace_end, kind, m) in enumerate(markers): - # Skip a marker whose start falls in another marker's brace region: it is - # that outer call's argument data, not its own call. - if any(s <= start <= e for j, (s, e) in enumerate(brace_regions) if j != idx): + # Skip a marker whose start falls in another marker's coverage: it is that + # outer call's data (its braces, or the gap up to its close), not a call. + # The end is exclusive: a marker starting exactly where another's close + # ends is the next sibling (adjacent calls), not nested. + if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue if brace_end < 0: continue # unclosed: not parseable; the fallback still excludes its XML 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 a2c57766584..d397f13c802 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -268,3 +268,35 @@ def test_strip_non_final_keeps_incomplete_gemma_block(): text = "before <|tool_call>call:t{" assert strip_tool_call_markup(text) == text assert strip_tool_call_markup(text, final = True) == "before" + + +def test_json_call_between_gemma_braces_and_close_does_not_execute(): + # A malformed outer Gemma call can carry a fully-formed JSON tool call after + # its balanced brace but before its close. That inner marker sits + # inside the outer call's coverage (up to its close), so it is data, not an + # executable call, in both allow_incomplete modes. + text = ( + "<|tool_call>call:outer{broken:{x}}" + '{"name":"terminal","arguments":{"command":"id"}}' + "" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): + # Same escape but the smuggled inner marker is Gemma-native, not JSON. + text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert "terminal" not in [c["function"]["name"] for c in calls], calls + + +def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): + # A closed whose parameter text contains a bare + # <|tool_call> must be stripped as a unit; the to-EOF Gemma sweep must not eat + # the trailing visible text after . + text = 'before print("<|tool_call>") after' + assert strip_tool_call_markup(text, final = True) == "before after" + assert strip_tool_call_markup(text) == "before after" From 4502c4e1b4022f6ca01358fc442896bbce387386 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 11:23:12 +0000 Subject: [PATCH 12/21] [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 | 4 +++- 1 file changed, 3 insertions(+), 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 d397f13c802..a0ce44d25ed 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -297,6 +297,8 @@ def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): # A closed whose parameter text contains a bare # <|tool_call> must be stripped as a unit; the to-EOF Gemma sweep must not eat # the trailing visible text after . - text = 'before print("<|tool_call>") after' + text = ( + 'before print("<|tool_call>") after' + ) assert strip_tool_call_markup(text, final = True) == "before after" assert strip_tool_call_markup(text) == "before after" From b3b1723d42f0b2d8749daa782182c8f691230dc7 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 1 Jul 2026 11:58:56 +0000 Subject: [PATCH 13/21] Strip closed tool blocks before the Gemma final sweep for PR #6611 The final display strip ran the quote-aware Gemma helper before the closed JSON/function patterns. A closed ... or ... block whose argument data held a call-form Gemma opener (e.g. a "<|tool_call>call:t{" string) was read as an incomplete Gemma span and truncated to EOF, dropping the block's close and any visible text after it. Strip closed JSON/function blocks first, so such a block is removed as a unit before the helper runs. Centralize the final strip order in a shared strip_tool_markup_final so strip_tool_call_markup and both streaming display wrappers (safetensors, llama_cpp) stay in sync, and apply the same closed-block pre-pass to the non-final path. Add regression tests for the JSON and function variants. --- studio/backend/core/inference/llama_cpp.py | 9 ++---- .../core/inference/safetensors_agentic.py | 9 ++---- .../core/inference/tool_call_parser.py | 1 + studio/backend/core/tool_healing.py | 32 +++++++++++++++---- .../tests/test_gemma_tool_parse_edge_cases.py | 13 ++++++++ 5 files changed, 44 insertions(+), 20 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 2961564c2c8..37bfc91f5c8 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -39,10 +39,8 @@ strip_split_mode_only, ) from core.tool_healing import ( - _TOOL_ALL_PATS, - _strip_gemma_native_spans, strip_tool_call_markup, - strip_tool_patterns, + strip_tool_markup_final, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -8071,10 +8069,7 @@ def _strip_tool_markup( def _strip_tool_markup_streaming(text: str, *, force: bool = False) -> str: if not (auto_heal_tool_calls or force): return text - # Quote-aware Gemma spans first, else a literal inside a - # quoted argument truncates the regex match and leaks the suffix. - text = _strip_gemma_native_spans(text, final = True) - return strip_tool_patterns(text, _TOOL_ALL_PATS) + return strip_tool_markup_final(text) def _build_metadata_event(usage, timings, finish_reason): """Final usage+timings metadata event for the given pass, merging its diff --git a/studio/backend/core/inference/safetensors_agentic.py b/studio/backend/core/inference/safetensors_agentic.py index 11c98dbc455..b00bc1cb3d3 100644 --- a/studio/backend/core/inference/safetensors_agentic.py +++ b/studio/backend/core/inference/safetensors_agentic.py @@ -21,15 +21,13 @@ from loggers import get_logger from core.inference.tool_call_parser import ( - _TOOL_ALL_PATS, - _strip_gemma_native_spans, BUDGET_EXHAUSTED_NUDGE, RAG_MAX_SEARCHES_PER_TURN, RAG_SEARCH_CAP_NUDGE, TOOL_XML_SIGNALS, parse_tool_calls_from_text, strip_tool_markup, - strip_tool_patterns, + strip_tool_markup_final, ) from core.inference.tool_loop_controller import ( ToolLoopController, @@ -62,10 +60,7 @@ def strip_tool_markup_streaming( """Strip open-ended tool XML from display text without trimming whitespace.""" if not (auto_heal_tool_calls or tool_protocol_active): return text - # Quote-aware Gemma spans first, else a literal inside a quoted - # argument truncates the regex match and leaks the suffix into display. - text = _strip_gemma_native_spans(text, final = True) - return strip_tool_patterns(text, _TOOL_ALL_PATS) + return strip_tool_markup_final(text) def _strip_tool_markup_final( diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index cfc30ea575e..b994e18a1e0 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -13,6 +13,7 @@ _TOOL_ALL_PATS = _tool_healing._TOOL_ALL_PATS _strip_gemma_native_spans = _tool_healing._strip_gemma_native_spans strip_tool_patterns = _tool_healing.strip_tool_patterns +strip_tool_markup_final = _tool_healing.strip_tool_markup_final def parse_tool_calls_from_text( diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 177874aab64..6686f44f3f2 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -31,6 +31,11 @@ re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +# Closed JSON/function blocks, stripped before the quote-aware Gemma helper so a +# Gemma opener sitting in their argument data (e.g. a "<|tool_call>call:t{" +# string) cannot make the helper truncate the whole block, and text after it, to +# EOF. Both are guarded below, so the pre-pass is skipped when absent. +_TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] # A lazy closed-pair pattern rescans to EOF from every opener when its close # token is absent (O(n^2); O(n^3) re-run per streamed token). Skip the doomed # sweep when the token is missing -- identical output, no backtracking. The @@ -540,6 +545,19 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: return "".join(out) +def strip_tool_markup_final(text: str) -> str: + """Final display strip, shared by ``strip_tool_call_markup`` and the streaming + wrappers so all three order the passes the same way. Closed JSON/function + blocks go first, so a Gemma opener in their argument data cannot make the + quote-aware helper truncate the block (and text after it) to EOF; then + well-formed Gemma spans (quote-aware); then the regex sweeps mop up malformed + spans and drop any unclosed remainder to EOF. Surrounding whitespace is kept. + """ + text = strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + text = _strip_gemma_native_spans(text, final = True) + return strip_tool_patterns(text, _TOOL_ALL_PATS) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -547,9 +565,11 @@ 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. """ - # Well-formed Gemma spans first (quote-aware); the regex patterns below then - # mop up any malformed Gemma span the helper could not match. - text = _strip_gemma_native_spans(text, final = final) - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - text = strip_tool_patterns(text, patterns) - return text.strip() if final else text + if final: + return strip_tool_markup_final(text).strip() + # Non-final: closed JSON/function blocks first (same reason as the final path), + # then keep any still-incomplete Gemma span (quote-aware), then the closed + # patterns mop up the rest without touching incomplete blocks. + text = strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + text = _strip_gemma_native_spans(text, final = False) + return strip_tool_patterns(text, _TOOL_CLOSED_PATS) 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 a0ce44d25ed..377ad29be13 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -302,3 +302,16 @@ def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): ) assert strip_tool_call_markup(text, final = True) == "before after" assert strip_tool_call_markup(text) == "before after" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # A closed JSON/function block whose argument data holds a call-form Gemma + # opener (e.g. "<|tool_call>call:t{") must be removed as a unit: the + # quote-aware helper must not treat the inner opener as an incomplete span and + # truncate the block (and the visible text after it) to EOF. + xml = "<|tool_call>call:t{" + json_block = '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + for block in (xml, json_block): + text = "before " + block + " after" + assert strip_tool_call_markup(text, final = True) == "before after", block + assert strip_tool_call_markup(text) == "before after", block From bddc1af122c78664c3e3ef216013962068619df0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 11:59:29 +0000 Subject: [PATCH 14/21] [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 | 4 +++- 1 file changed, 3 insertions(+), 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 377ad29be13..047359fdc8b 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -310,7 +310,9 @@ def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener() # quote-aware helper must not treat the inner opener as an incomplete span and # truncate the block (and the visible text after it) to EOF. xml = "<|tool_call>call:t{" - json_block = '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + json_block = ( + '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' + ) for block in (xml, json_block): text = "before " + block + " after" assert strip_tool_call_markup(text, final = True) == "before after", block From 51669478bdf392ecd53b95ce124e4e07b5a06587 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 1 Jul 2026 12:37:19 +0000 Subject: [PATCH 15/21] Recover XML/JSON siblings after a close-less tool marker for PR #6611 Two fixes so the XML fallback and marker coverage recover a later valid call after an earlier marker omits its close, matching the candidate loop: Reuse the candidate marker-coverage in the XML fallback instead of a separate search-to-close-or-EOF envelope. A balanced but close-less marker now covers only its brace region there too, so a following sibling is recovered rather than filtered as nested data; an unbalanced marker still covers to EOF and a closed one still covers through its close, so nested XML stays blocked. Ignore a close token that falls inside another call's balanced braces when pairing closes in _marker_coverage. Such a token is that call's quoted argument data, so it no longer pops an earlier close-less marker and extends its coverage over a later valid sibling. Add regression tests for both. --- studio/backend/core/tool_healing.py | 25 ++++++++------- .../tests/test_gemma_tool_parse_edge_cases.py | 32 +++++++++++++++++++ 2 files changed, 45 insertions(+), 12 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 6686f44f3f2..3ff8108e059 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -329,12 +329,19 @@ def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: later sibling after an omitted close marker is still recovered as its own call. """ n = len(content) + brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] events = [] # (position, order) with order 0 = braces-done, 1 = close marker for idx, (_start, brace_end, _kind, _m) in enumerate(markers): if brace_end >= 0: events.append((brace_end, 0, _kind, idx)) for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): for cm in close_re.finditer(content): + # A close token inside another call's balanced braces is that call's + # quoted argument data, not a structural close. Ignore it, else it + # could pop an earlier close-less marker and extend its coverage over a + # later valid sibling (dropping that sibling). + if any(s < cm.start() < be for s, be in brace_regions): + continue events.append((cm.start(), 1, kind, cm.end())) events.sort(key = lambda e: (e[0], e[1])) waiting = {"json": [], "gemma": []} @@ -424,22 +431,16 @@ def parse_tool_calls_from_text( ) if not tool_calls: - # Exclude markers inside any marker's envelope -- its data - # through the close marker (searched after the braces) or EOF, even if the - # call failed to parse -- so nested XML cannot escape as a real call. - envelopes = [] - for s, be, kind, _m in markers: - if be < 0: - envelopes.append((s, len(content))) - else: - close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE - cm = close_re.search(content, be + 1) - envelopes.append((s, cm.end() if cm else len(content))) + # Exclude markers inside any marker's coverage -- its braces, + # or the gap up to its close -- even if that call failed to parse, so + # nested XML cannot escape as a real call. Reusing the same coverage as the + # candidate loop keeps this consistent: a after a balanced + # close-less marker is a sibling (recovered), not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() < e for s, e in envelopes) + and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) 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 047359fdc8b..da5a3d5d6cb 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -317,3 +317,35 @@ def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener() text = "before " + block + " after" assert strip_tool_call_markup(text, final = True) == "before after", block assert strip_tool_call_markup(text) == "before after", block + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # A balanced but unparsable Gemma marker with no close tag, followed by a valid + # XML call: the malformed marker only covers its brace region, so + # the sibling function call is recovered, not filtered as nested data. + text = ( + "<|tool_call>call:bad{broken:{x}} " + "id" + ) + for allow_incomplete in (True, False): + calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) + assert [c["function"]["name"] for c in calls] == ["terminal"], calls + + +def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): + # The next valid call carries the literal close token inside a quoted argument. + # That token is the later call's data, not a structural close for the earlier + # close-less marker, so it must not extend the earlier marker's coverage over + # the later call and drop it. + gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' + names = [c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)] + assert names == ["b"], names + json_text = ( + '{"name":"a","arguments":{}} ' + '{"name":"b","arguments":{"x":""}}' + ) + names_j = [ + c["function"]["name"] + for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + ] + assert "b" in names_j, names_j From b87e24776593110816931d46fb0c50457ab32046 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:38:22 +0000 Subject: [PATCH 16/21] [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 | 7 ++++--- 1 file changed, 4 insertions(+), 3 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 da5a3d5d6cb..cb7d11dbef4 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -338,14 +338,15 @@ def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered # close-less marker, so it must not extend the earlier marker's coverage over # the later call and drop it. gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' - names = [c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False)] + names = [ + c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) + ] assert names == ["b"], names json_text = ( '{"name":"a","arguments":{}} ' '{"name":"b","arguments":{"x":""}}' ) names_j = [ - c["function"]["name"] - for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) + c["function"]["name"] for c in parse_tool_calls_from_text(json_text, allow_incomplete = False) ] assert "b" in names_j, names_j From 7f448a4356d20afdd5fe88bd8be06f71743fa1cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 4 Jul 2026 10:18:30 +0000 Subject: [PATCH 17/21] Make the closed-block strip pre-pass Gemma-span-aware The final display strip ran the closed JSON/function regex pre-pass before removing Gemma-native spans, so a literal quoted inside a Gemma argument plus any later (a real call's close or even prose) was deleted across the Gemma boundary. That mangled the Gemma close marker, the quote-aware helper then saw an unclosed opener, and the whole visible tail after the call was truncated. The pre-pass now skips matches that start inside a complete Gemma span (that text is the span's argument data) and resumes scanning at the end of the covering span, so a real function-XML call after the Gemma call is still stripped. The original ordering rationale is preserved: a Gemma opener inside a JSON or function argument still cannot truncate that block, covered by regression tests for both directions. --- studio/backend/core/tool_healing.py | 71 +++++++++++++++++-- .../tests/test_tool_call_parser_strict.py | 44 ++++++++++++ 2 files changed, 109 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index afc2d102156..37076cf07b5 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -567,15 +567,74 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: return "".join(out) +def _gemma_span_ranges(text: str) -> list: + """``(start, end)`` of each complete Gemma-native span, same walk as + ``_strip_gemma_native_spans`` without stripping (brace/quote-balanced, + close marker searched after the braces).""" + ranges: list[tuple] = [] + 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: + break + close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) + if close is None: + break + ranges.append((start, close.end())) + cursor = close.end() + return ranges + + +def _strip_closed_blocks_outside_gemma(text: str) -> str: + """Closed JSON/function pre-pass that leaves matches STARTING inside a + complete Gemma span alone. A literal ```` quoted in a Gemma + argument plus a later real ```` would otherwise be deleted + across the Gemma boundary, mangling the Gemma close marker so the + quote-aware strip truncates the whole tail. A skipped match resumes the + scan at the END of the covering Gemma span (not the match end), so a real + function-XML call after the span is still stripped.""" + ranges = _gemma_span_ranges(text) + if not ranges: + return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + for pat in _TOOL_CLOSED_BLOCK_PATS: + token = _PAT_REQUIRED_TOKEN.get(pat) + if token is not None and token not in text: + continue + out: list[str] = [] + pos = 0 + while True: + m = pat.search(text, pos) + if m is None: + out.append(text[pos:]) + break + covering = next((r for r in ranges if r[0] <= m.start() < r[1]), None) + if covering is not None: + out.append(text[pos : covering[1]]) + pos = covering[1] + continue + out.append(text[pos : m.start()]) + pos = m.end() + new_text = "".join(out) + if new_text != text: + text = new_text + ranges = _gemma_span_ranges(text) + return text + + def strip_tool_markup_final(text: str) -> str: """Final display strip, shared by ``strip_tool_call_markup`` and the streaming wrappers so all three order the passes the same way. Closed JSON/function - blocks go first, so a Gemma opener in their argument data cannot make the - quote-aware helper truncate the block (and text after it) to EOF; then - well-formed Gemma spans (quote-aware); then the regex sweeps mop up malformed - spans and drop any unclosed remainder to EOF. Surrounding whitespace is kept. + blocks go first (so a Gemma opener in their argument data cannot make the + quote-aware helper truncate the block and its tail to EOF), but Gemma-aware: + a match starting inside a complete Gemma span is that span's argument data, + not a block to delete across the boundary. Then well-formed Gemma spans + (quote-aware); then the regex sweeps mop up malformed spans and drop any + unclosed remainder to EOF. Surrounding whitespace is kept. """ - text = strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + text = _strip_closed_blocks_outside_gemma(text) text = _strip_gemma_native_spans(text, final = True) return strip_tool_patterns(text, _TOOL_ALL_PATS) @@ -592,6 +651,6 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: # Non-final: closed JSON/function blocks first (same reason as the final path), # then keep any still-incomplete Gemma span (quote-aware), then the closed # patterns mop up the rest without touching incomplete blocks. - text = strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) + text = _strip_closed_blocks_outside_gemma(text) text = _strip_gemma_native_spans(text, final = False) return strip_tool_patterns(text, _TOOL_CLOSED_PATS) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 39fdd151be3..791776bfc09 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -197,3 +197,47 @@ def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): assert text[span[0] : span[1]] == ( "cats" ) + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + Gemma span: a literal quoted in a Gemma argument plus a later + real would otherwise mangle the Gemma close marker and truncate + the whole visible tail.""" + + def test_literal_function_in_gemma_arg_with_later_real_call(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " ls" + " after" + ) + assert strip_tool_call_markup(text, final = True) == "before after" + + def test_literal_function_in_gemma_arg_with_prose_closer(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' + " then use to close. after" + ) + out = strip_tool_call_markup(text, final = True) + assert out.startswith("before") + assert out.endswith("after") + assert "call:python" not in out + + def test_gemma_opener_inside_json_arg_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + + text = '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + assert strip_tool_call_markup(text, final = True) == "after" + + def test_gemma_opener_inside_function_param_still_strips_block(self): + from core.tool_healing import strip_tool_call_markup + + text = ( + 'x = "<|tool_call>call:t{"' + " after" + ) + assert strip_tool_call_markup(text, final = True) == "after" From 52a935e09fc01c584b080823789d3f9e84b38558 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 4 Jul 2026 10:18:59 +0000 Subject: [PATCH 18/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_tool_call_parser_strict.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 791776bfc09..f59e7a930c3 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -207,7 +207,6 @@ class TestGemmaAwareClosedBlockPrePass: def test_literal_function_in_gemma_arg_with_later_real_call(self): from core.tool_healing import strip_tool_call_markup - text = ( 'before <|tool_call>call:python{code:<|"|>print("")<|"|>}' " ls" @@ -229,13 +228,13 @@ def test_literal_function_in_gemma_arg_with_prose_closer(self): def test_gemma_opener_inside_json_arg_still_strips_block(self): from core.tool_healing import strip_tool_call_markup - - text = '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + text = ( + '{"name":"t","arguments":{"code":"<|tool_call>call:x{"}} after' + ) assert strip_tool_call_markup(text, final = True) == "after" def test_gemma_opener_inside_function_param_still_strips_block(self): from core.tool_healing import strip_tool_call_markup - text = ( 'x = "<|tool_call>call:t{"' " after" From 71624800a52d8670ed8eb8a2d898a1cf04678e08 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:21:15 +0000 Subject: [PATCH 19/21] Trim comments in the Gemma streaming and strip pipeline to essentials --- studio/backend/core/tool_healing.py | 142 ++++++------------ studio/backend/routes/inference.py | 14 +- .../tests/test_gemma_tool_parse_edge_cases.py | 92 +++--------- .../tests/test_tool_call_parser_strict.py | 4 +- studio/backend/tests/test_tool_strip_guard.py | 8 +- 5 files changed, 76 insertions(+), 184 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 37076cf07b5..ce58ee6cdad 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -10,12 +10,10 @@ import json import re -# Tool XML stripping patterns. Hyphen in the name class matches dashed MCP names -# (mcp__srv__list-issues). Both lists strip every closed pair first (JSON, Gemma, -# function), so a closed call is removed as a unit before any to-EOF sweep can -# reach markup nested inside it; the final list then adds greedy .*$ sweeps that -# drop an unclosed opener's remainder to EOF. The non-final list keeps incomplete -# blocks (no EOF sweeps), matching the parser's allow_incomplete = False path. +# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs +# strip first, so a closed call is removed as a unit before any to-EOF sweep can +# reach markup nested inside it; only the final list adds the .*$ EOF sweeps, so +# the non-final list keeps incomplete blocks. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) @@ -31,15 +29,11 @@ re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] -# Closed JSON/function blocks, stripped before the quote-aware Gemma helper so a -# Gemma opener sitting in their argument data (e.g. a "<|tool_call>call:t{" -# string) cannot make the helper truncate the whole block, and text after it, to -# EOF. Both are guarded below, so the pre-pass is skipped when absent. +# Stripped before the quote-aware Gemma helper so a Gemma opener quoted in +# their argument data cannot make the helper truncate the block and its tail. _TOOL_CLOSED_BLOCK_PATS = [_TC_JSON_CLOSED_PAT, _TC_FUNC_CLOSED_PAT] -# A lazy closed-pair pattern rescans to EOF from every opener when its close -# token is absent (O(n^2); O(n^3) re-run per streamed token). Skip the doomed -# sweep when the token is missing -- identical output, no backtracking. The -# to-EOF sweeps are greedy, so they short-circuit on the first opener. +# A lazy closed-pair pattern whose close token is absent rescans to EOF from +# every opener (quadratic, re-run per streamed token); skip that doomed pass. _PAT_REQUIRED_TOKEN = { _TC_JSON_CLOSED_PAT: "", _TC_GEMMA_CLOSED_PAT: "", @@ -48,8 +42,7 @@ def strip_tool_patterns(text: str, patterns) -> str: - """Apply strip ``patterns`` in order, skipping a closed-pair pass whose close - token is absent (avoids its quadratic no-match rescan).""" + """Apply ``patterns`` in order, skipping closed-pair passes with no close token.""" for pat in patterns: token = _PAT_REQUIRED_TOKEN.get(pat) if token is not None and token not in text: @@ -70,9 +63,8 @@ def strip_tool_patterns(text: str, patterns) -> str: _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -# A bare Gemma value ends at `}` or a comma starting the next `key:` pair. The -# key must be identifier-shaped, so a comma before digits-then-colon stays in -# the value (`location:New York, NY`; `meet at 10:00, 11:00`). +# A bare Gemma value ends at `}` or a comma starting the next identifier-shaped +# `key:`; commas before non-keys stay in the value (`New York, NY`, `10:00, 11:00`). _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w-]*\s*:") @@ -169,10 +161,8 @@ def _split_top_level_commas(src: str) -> list: def _quote_gemma_array_elements(body: str) -> str: - """Normalise a Gemma array value so json.loads succeeds: quote bare string - elements, recurse into object/nested-array elements, leave quoted strings, - numbers, and JSON literals as-is. Gemma emits these unquoted - (``labels:[bug,ui]``, ``items:[{path:a}]``) and the call would otherwise drop.""" + """Normalise a Gemma array value (``labels:[bug,ui]``) so json.loads succeeds: + quote bare strings, recurse into objects/arrays, keep quoted/JSON literals.""" out: list[str] = [] for element in _split_top_level_commas(body): stripped = element.strip() @@ -259,8 +249,7 @@ def _quote_gemma_object_keys(src: str) -> str: parts.append(src[i:colon_pos]) parts.append(":") i = colon_pos + 1 - # Quote bare string values ({unit:celsius}); JSON scalars/objects/ - # arrays/quoted strings stay as-is. + # Quote bare string values ({unit:celsius}); JSON stays as-is. ws = i while i < len(src) and src[i].isspace(): i += 1 @@ -317,17 +306,11 @@ def _inside_open_parameter(content: str, pos: int) -> bool: def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: - """Coverage region ``[start, end]`` for each marker, used to skip markers that - are another call's data. Each marker's own close is paired to it with a - per-format stack (a close after the braces pops the nearest still-open marker - of that format), so an inner call's close is not mistaken for the outer's: - - - braces never balance -> covers ``[start, EOF)`` (the whole ambiguous tail); - - braces balance and a matching close is found -> covers ``[start, close_end]`` - so a marker smuggled between the braces and that close is treated as data; - - braces balance but no close is paired -> covers only the brace region, so a - later sibling after an omitted close marker is still recovered as its own call. - """ + """Coverage ``[start, end]`` per marker, used to skip markers that are another + call's data. Closes pair to markers via a per-format stack so an inner close + is not mistaken for the outer's. Unbalanced braces cover to EOF; balanced with + a paired close cover through it (markers before the close are data); balanced + without one cover only the braces, so a later sibling is still recovered.""" n = len(content) brace_regions = [(s, be) for (s, be, _k, _m) in markers if be >= 0] events = [] # (position, order) with order 0 = braces-done, 1 = close marker @@ -336,10 +319,8 @@ def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: events.append((brace_end, 0, _kind, idx)) for kind, close_re in (("json", _TC_END_TAG_RE), ("gemma", _TC_GEMMA_END_TAG_RE)): for cm in close_re.finditer(content): - # A close token inside another call's balanced braces is that call's - # quoted argument data, not a structural close. Ignore it, else it - # could pop an earlier close-less marker and extend its coverage over a - # later valid sibling (dropping that sibling). + # A close inside another call's balanced braces is quoted data; it + # must not pop an earlier close-less marker and swallow a sibling. if any(s < cm.start() < be for s, be in brace_regions): continue events.append((cm.start(), 1, kind, cm.end())) @@ -383,13 +364,9 @@ def parse_tool_calls_from_text( """ tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect JSON/Gemma markers, then decide nesting by each marker's coverage - # region (see _marker_coverage): a closed outer covers up to its own close - # marker, so a JSON/Gemma marker smuggled between the outer braces and that - # close is treated as data, not executed; an outer that balances but has no - # close of its own covers only its brace region, so a later sibling after an - # omitted close is still recovered. A marker inside an open - # value is that parameter's data, so skip it. + # Collect JSON/Gemma markers; nesting is decided by _marker_coverage (a + # marker inside another call's coverage is data, not executed). Markers + # inside an open value are that parameter's data. markers = [] # (start, brace_end, kind, match); brace_end < 0 = unclosed for start_re, gemma, kind in ( (_TC_JSON_START_RE, False, "json"), @@ -405,10 +382,8 @@ def parse_tool_calls_from_text( coverage = _marker_coverage(content, markers) parsed_items = [] # (start, span_end, name, arguments) in document order for idx, (start, brace_end, kind, m) in enumerate(markers): - # Skip a marker whose start falls in another marker's coverage: it is that - # outer call's data (its braces, or the gap up to its close), not a call. - # The end is exclusive: a marker starting exactly where another's close - # ends is the next sibling (adjacent calls), not nested. + # A marker starting inside another's coverage is that call's data. The + # end is exclusive so a marker at a close's end is an adjacent sibling. if any(s <= start < e for j, (s, e) in enumerate(coverage) if j != idx): continue if brace_end < 0: @@ -430,8 +405,7 @@ def parse_tool_calls_from_text( arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end])) except (json.JSONDecodeError, ValueError): continue - # Span for with_spans callers: through the close tag when present so the - # healed strip removes the whole wrapped call, else just the braces. + # Span reaches through the close tag when present, else just the braces. span_end = brace_end + 1 close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE ws = len(content[span_end:]) - len(content[span_end:].lstrip()) @@ -440,12 +414,10 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) - # Function-XML calls parse alongside marker calls (mixed formats promote in - # document order -- the #6801 contract). Exclude markers inside - # any marker's coverage -- its braces, or the gap up to its close -- even if - # that call failed to parse, so nested XML cannot escape as a real call. A - # after a balanced close-less marker is a sibling (recovered), - # not swallowed to EOF. + # Function-XML calls promote in document order alongside marker calls (the + # #6801 contract). A inside any marker's coverage is excluded -- + # even if that marker failed to parse -- so nested XML cannot escape; one + # after a balanced close-less marker is a sibling, not swallowed to EOF. func_starts = [ fm for fm in _TC_FUNC_START_RE.finditer(content) @@ -531,11 +503,9 @@ def parse_tool_calls_from_text( def _strip_gemma_native_spans(text: str, *, final: bool) -> str: - """Remove complete Gemma-native ``<|tool_call>call:NAME{...}`` - spans, brace/quote-balanced so a literal ```` inside a - ``<|"|>``-quoted argument cannot truncate the span and leak its suffix. An - incomplete span is dropped to EOF when ``final``, else kept (still streaming). - """ + """Remove complete Gemma-native spans, brace/quote-balanced so a literal + ```` in a quoted argument cannot truncate the span. An incomplete + span is dropped to EOF when ``final``, else kept (still streaming).""" out: list[str] = [] cursor = 0 for match in _TC_GEMMA_START_RE.finditer(text): @@ -544,17 +514,14 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: continue brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) if brace_end < 0: - # Unbalanced: no complete span from here on. Drop the rest if final, - # else keep it, and stop -- later starts are inside this unclosed run, - # so rescanning them would re-walk to EOF each time (quadratic). + # Unbalanced: nothing completes from here on. Drop the rest if final, + # else keep it; stop either way (rescanning would be quadratic). if final: out.append(text[cursor:start]) cursor = len(text) break - # Search for the close marker after the braces (not just immediately - # after): junk between } and is malformed-call markup, so - # strip through the close and keep any text after it. None anywhere after - # means nothing closes from here on, so stop (keeps this linear). + # Junk between } and is malformed-call markup: strip through + # the close, keep text after it. No close anywhere means stop (linear). close = _TC_GEMMA_END_TAG_RE.search(text, brace_end + 1) if close is None: if final: @@ -568,9 +535,8 @@ def _strip_gemma_native_spans(text: str, *, final: bool) -> str: def _gemma_span_ranges(text: str) -> list: - """``(start, end)`` of each complete Gemma-native span, same walk as - ``_strip_gemma_native_spans`` without stripping (brace/quote-balanced, - close marker searched after the braces).""" + """``(start, end)`` of each complete Gemma-native span; same walk as + ``_strip_gemma_native_spans`` without stripping.""" ranges: list[tuple] = [] cursor = 0 for match in _TC_GEMMA_START_RE.finditer(text): @@ -589,13 +555,10 @@ def _gemma_span_ranges(text: str) -> list: def _strip_closed_blocks_outside_gemma(text: str) -> str: - """Closed JSON/function pre-pass that leaves matches STARTING inside a - complete Gemma span alone. A literal ```` quoted in a Gemma - argument plus a later real ```` would otherwise be deleted - across the Gemma boundary, mangling the Gemma close marker so the - quote-aware strip truncates the whole tail. A skipped match resumes the - scan at the END of the covering Gemma span (not the match end), so a real - function-XML call after the span is still stripped.""" + """Closed JSON/function pre-pass that skips matches starting inside a complete + Gemma span: deleting across the span boundary would mangle the Gemma close and + truncate the tail. A skipped match resumes at the covering span's end, so a + real function-XML call after the span is still stripped.""" ranges = _gemma_span_ranges(text) if not ranges: return strip_tool_patterns(text, _TOOL_CLOSED_BLOCK_PATS) @@ -625,15 +588,10 @@ def _strip_closed_blocks_outside_gemma(text: str) -> str: def strip_tool_markup_final(text: str) -> str: - """Final display strip, shared by ``strip_tool_call_markup`` and the streaming - wrappers so all three order the passes the same way. Closed JSON/function - blocks go first (so a Gemma opener in their argument data cannot make the - quote-aware helper truncate the block and its tail to EOF), but Gemma-aware: - a match starting inside a complete Gemma span is that span's argument data, - not a block to delete across the boundary. Then well-formed Gemma spans - (quote-aware); then the regex sweeps mop up malformed spans and drop any - unclosed remainder to EOF. Surrounding whitespace is kept. - """ + """Final display strip, shared with the streaming wrappers so all paths order + the passes identically: Gemma-aware closed JSON/function blocks first, then + well-formed Gemma spans (quote-aware), then the regex sweeps mop up malformed + spans and drop any unclosed remainder to EOF. Whitespace is kept.""" text = _strip_closed_blocks_outside_gemma(text) text = _strip_gemma_native_spans(text, final = True) return strip_tool_patterns(text, _TOOL_ALL_PATS) @@ -648,9 +606,7 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """ if final: return strip_tool_markup_final(text).strip() - # Non-final: closed JSON/function blocks first (same reason as the final path), - # then keep any still-incomplete Gemma span (quote-aware), then the closed - # patterns mop up the rest without touching incomplete blocks. + # Non-final: same ordering as the final path, but incomplete blocks are kept. text = _strip_closed_blocks_outside_gemma(text) text = _strip_gemma_native_spans(text, final = False) return strip_tool_patterns(text, _TOOL_CLOSED_PATS) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 82e97c206ae..520176e69fa 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -893,9 +893,8 @@ async def _tracking_send(message) -> None: def _tracked_cancel_unstarted_cleanup(tracker): - """unstarted_cleanup for a local stream that entered ``tracker`` before - returning: exits it on pre-start disconnect, when the generator's finally - (which normally does so) never runs. Mutually exclusive with that finally.""" + """unstarted_cleanup that exits ``tracker`` on a pre-start disconnect, when + the generator's finally (which normally exits it) never runs.""" async def _cleanup() -> None: tracker.__exit__(None, None, None) @@ -10879,12 +10878,9 @@ async def _openai_passthrough_stream( ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so the client sees a standard OpenAI response. - Reasoning/tool-call splitting is delegated to llama-server's - ``/v1/chat/completions`` (Studio runs ``--jinja --reasoning-format auto``), - so ``delta.content`` carries no raw ````/``<|tool_call>`` markup and - is deliberately not re-parsed locally (unlike the ``/completion`` paths). If - a future llama.cpp build stops splitting it, this path would need the local - extractor as a safety net. + Reasoning/tool-call splitting is delegated to llama-server (``--jinja + --reasoning-format auto``), so ``delta.content`` carries no raw markup and is + deliberately not re-parsed locally, unlike the ``/completion`` paths. """ 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 cb7d11dbef4..9aa55874421 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -1,15 +1,8 @@ # 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. -""" +"""Gemma-native tool-call parsing edge cases: commas inside bare string values, +and markers inside another call's argument data staying data.""" from __future__ import annotations @@ -41,14 +34,11 @@ def test_bare_string_argument_with_comma_is_kept(): 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_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. + # A comma before digits-then-colon (timestamp/ratio) is value text, not a key. calls = parse_tool_calls_from_text( "<|tool_call>call:remind{query:meet at 10:00, 11:00 tomorrow,priority:high}" ) @@ -57,8 +47,6 @@ def test_bare_value_with_timestamps_after_comma_is_kept(): 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}"}}' @@ -76,8 +64,6 @@ def test_two_separate_gemma_calls_both_parse(): 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"}}' @@ -87,8 +73,6 @@ def test_mixed_format_calls_preserve_document_order(): 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"}})' @@ -99,18 +83,14 @@ def test_json_marker_inside_gemma_argument_is_not_a_second_call(): 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. + # The outer object fails to normalize, but the nested marker is covered by + # its span; safe outcome is no executed call at all. 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. 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"]} @@ -124,8 +104,6 @@ def test_array_keeps_numbers_and_quoted_elements(): 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}]}" ) @@ -139,9 +117,6 @@ def test_nested_array_elements_are_normalised(): 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}" @@ -163,9 +138,7 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): 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. + # Parse keeps the quoted close marker as data; strip removes the whole span. text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' calls = parse_tool_calls_from_text(text) assert len(calls) == 1, calls @@ -175,9 +148,7 @@ def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): def test_nested_xml_in_malformed_gemma_call_does_not_execute(): - # A balanced but unparsable Gemma call whose argument data contains XML tool - # markup must not let that escape into an executable call via the - # XML fallback (the Gemma candidate span covers it even though it failed). + # The failed Gemma candidate's span still covers its nested . text = ( "<|tool_call>call:outer{code:id" ", broken:{x}}" @@ -188,9 +159,7 @@ def test_nested_xml_in_malformed_gemma_call_does_not_execute(): def test_unbalanced_gemma_call_with_xml_does_not_execute(): - # An unbalanced Gemma call (braces never close) records no candidate span, - # so the XML fallback must exclude its trailing through EOF - # rather than promote it to an executable call. + # Unclosed braces cover to EOF, so the trailing is excluded. text = ( "<|tool_call>call:outer{code:" "id" @@ -201,17 +170,13 @@ def test_unbalanced_gemma_call_with_xml_does_not_execute(): def test_standalone_function_xml_still_parses(): - # The exclusion must not over-block: a real call with no - # preceding unclosed Gemma/JSON start is still a valid tool call. text = "id" calls = parse_tool_calls_from_text(text) assert [c["function"]["name"] for c in calls] == ["terminal"], calls def test_xml_between_braces_and_close_marker_does_not_execute(): - # Balanced-but-unparsable outer call with XML after the braces but before the - # close marker: the envelope runs to the close marker, so here is - # the outer call's data, not an executable tool call. + # Coverage runs to the close marker, so in the gap is data. text = ( "<|tool_call>call:outer{broken:{x}}" "id" @@ -222,8 +187,6 @@ def test_xml_between_braces_and_close_marker_does_not_execute(): def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): - # A balanced inner call inside an unclosed outer call's argument data must be - # skipped, not accepted, even though its own braces balance. text = "<|tool_call>call:outer{code:<|tool_call>call:terminal{command:id}" for allow_incomplete in (True, False): calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) @@ -231,16 +194,13 @@ def test_balanced_inner_call_inside_unclosed_outer_does_not_execute(): def test_strip_preserves_text_after_malformed_gemma_close(): - # A valid call:name{...} prefix with junk before its close is a - # malformed closed span: strip through the close, keep the text after it. + # Junk before the close is a malformed span: strip through it, keep the tail. text = "pre <|tool_call>call:t{a:1} note post" assert strip_tool_call_markup(text) == "pre post" assert strip_tool_call_markup(text, final = True) == "pre post" def test_malformed_closed_gemma_span_is_stripped(): - # A closed Gemma span the quote-aware helper cannot match (no call:NAME{) - # must still be stripped, not leak its opener/payload into visible text. assert ( strip_tool_call_markup('before <|tool_call>{"name":"x"} after') == "before after" @@ -248,9 +208,7 @@ def test_malformed_closed_gemma_span_is_stripped(): def test_valid_call_after_missing_close_is_recovered(): - # A balanced call missing its own close marker must not swallow a later valid - # closed call: nesting is decided by the brace region, not an envelope that - # would run to EOF, so the second call is still recovered. + # A close-less call covers only its braces, so the later call is recovered. text = "<|tool_call>call:a{x:1} <|tool_call>call:b{y:2}" names_inc = [ c["function"]["name"] for c in parse_tool_calls_from_text(text, allow_incomplete = True) @@ -263,18 +221,13 @@ def test_valid_call_after_missing_close_is_recovered(): def test_strip_non_final_keeps_incomplete_gemma_block(): - # Non-final must preserve an incomplete Gemma block (matching JSON/function), - # while final strips the unclosed remainder to EOF. text = "before <|tool_call>call:t{" assert strip_tool_call_markup(text) == text assert strip_tool_call_markup(text, final = True) == "before" def test_json_call_between_gemma_braces_and_close_does_not_execute(): - # A malformed outer Gemma call can carry a fully-formed JSON tool call after - # its balanced brace but before its close. That inner marker sits - # inside the outer call's coverage (up to its close), so it is data, not an - # executable call, in both allow_incomplete modes. + # A JSON call between the outer's braces and its close is covered data. text = ( "<|tool_call>call:outer{broken:{x}}" '{"name":"terminal","arguments":{"command":"id"}}' @@ -286,7 +239,7 @@ def test_json_call_between_gemma_braces_and_close_does_not_execute(): def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): - # Same escape but the smuggled inner marker is Gemma-native, not JSON. + # Same escape with a Gemma-native inner marker. text = "<|tool_call>call:outer{broken:{x}}<|tool_call>call:terminal{command:id}" for allow_incomplete in (True, False): calls = parse_tool_calls_from_text(text, allow_incomplete = allow_incomplete) @@ -294,9 +247,7 @@ def test_gemma_call_between_gemma_braces_and_close_does_not_execute(): def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): - # A closed whose parameter text contains a bare - # <|tool_call> must be stripped as a unit; the to-EOF Gemma sweep must not eat - # the trailing visible text after . + # The to-EOF Gemma sweep must not eat visible text after . text = ( 'before print("<|tool_call>") after' ) @@ -305,10 +256,7 @@ def test_strip_final_keeps_text_after_closed_xml_with_inner_gemma_opener(): def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): - # A closed JSON/function block whose argument data holds a call-form Gemma - # opener (e.g. "<|tool_call>call:t{") must be removed as a unit: the - # quote-aware helper must not treat the inner opener as an incomplete span and - # truncate the block (and the visible text after it) to EOF. + # A call-form Gemma opener quoted in a closed block must not truncate it. xml = "<|tool_call>call:t{" json_block = ( '{"name":"python","arguments":{"code":"<|tool_call>call:t{"}}' @@ -320,9 +268,7 @@ def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener() def test_function_sibling_after_close_less_gemma_marker_is_recovered(): - # A balanced but unparsable Gemma marker with no close tag, followed by a valid - # XML call: the malformed marker only covers its brace region, so - # the sibling function call is recovered, not filtered as nested data. + # The close-less marker covers only its braces; the XML sibling is recovered. text = ( "<|tool_call>call:bad{broken:{x}} " "id" @@ -333,10 +279,8 @@ def test_function_sibling_after_close_less_gemma_marker_is_recovered(): def test_valid_call_after_close_less_marker_with_quoted_close_token_is_recovered(): - # The next valid call carries the literal close token inside a quoted argument. - # That token is the later call's data, not a structural close for the earlier - # close-less marker, so it must not extend the earlier marker's coverage over - # the later call and drop it. + # A close token quoted in the later call must not extend the earlier + # close-less marker's coverage over that call. gemma = '<|tool_call>call:a{x:1} <|tool_call>call:b{note:<|"|><|"|>}' names = [ c["function"]["name"] for c in parse_tool_calls_from_text(gemma, allow_incomplete = False) diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index f59e7a930c3..e4e60c89952 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -201,9 +201,7 @@ def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self): class TestGemmaAwareClosedBlockPrePass: """The closed JSON/function strip pre-pass must not delete across a complete - Gemma span: a literal quoted in a Gemma argument plus a later - real would otherwise mangle the Gemma close marker and truncate - the whole visible tail.""" + Gemma span (a quoted plus a later real ).""" def test_literal_function_in_gemma_arg_with_later_real_call(self): from core.tool_healing import strip_tool_call_markup diff --git a/studio/backend/tests/test_tool_strip_guard.py b/studio/backend/tests/test_tool_strip_guard.py index ef5bdaf769e..dfa31018826 100644 --- a/studio/backend/tests/test_tool_strip_guard.py +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -"""strip_tool_patterns must produce identical output to the plain per-pattern -loop, while skipping a lazy closed-pair sweep whose close token is absent (the -O(n^2) no-match rescan, O(n^3) when re-run per streamed token).""" +"""strip_tool_patterns must match the plain per-pattern loop while skipping the +quadratic no-match rescan of a closed-pair sweep whose close token is absent.""" import random import sys @@ -69,8 +68,7 @@ def test_strip_markup_representative_cases_unchanged(): def test_no_quadratic_blowup_on_unclosed_markers(): - # Many openers with no close token: the guard skips the lazy closed-pair - # sweep, keeping this linear. Unguarded this took minutes. + # Unguarded, this took minutes. big = "" * 20000 + "" * 20000 t0 = time.perf_counter() out = strip_tool_call_markup(big, final = True) From 45078844b13952d60c2ef330bfc884d1a32d1c29 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 08:35:13 +0000 Subject: [PATCH 20/21] Tighten comments in the Gemma strip and streaming disconnect paths --- studio/backend/core/tool_healing.py | 5 ++--- studio/backend/routes/inference.py | 17 +++++++---------- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index ce58ee6cdad..706a84e578f 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -11,9 +11,8 @@ import re # Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs -# strip first, so a closed call is removed as a unit before any to-EOF sweep can -# reach markup nested inside it; only the final list adds the .*$ EOF sweeps, so -# the non-final list keeps incomplete blocks. +# strip first so a closed call goes as a unit before any to-EOF sweep reaches +# nested markup; only the final list adds the .*$ EOF sweeps. _TC_JSON_CLOSED_PAT = re.compile(r".*?", re.DOTALL) _TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?", re.DOTALL) _TC_FUNC_CLOSED_PAT = re.compile(r".*?", re.DOTALL) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 520176e69fa..c54fb12a091 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4015,10 +4015,9 @@ async def stream(): _DONE = object() while True: if cancel_event.is_set(): - # Watcher set cancel_event between chunks. Reset here: - # closing the generator does not signal a subprocess backend, - # so it would keep decoding. The finally's reset is guarded on - # cancel_event being unset, so it will not double-run. + # Watcher set cancel_event between chunks. Reset here: closing + # the generator does not signal a subprocess backend, so it would + # keep decoding. The finally's reset is guarded, so no double-run. backend.reset_generation_state() break chunk = await asyncio.to_thread(next, gen, _DONE) @@ -9829,9 +9828,8 @@ async def _stream(): drop_until_tool_end = False gen = run_gen() - # Watcher to cancel on disconnect: the in-loop poll only fires between - # events, so a disconnect during a long prefill/step would otherwise hold - # the decode slot until the next event or a failed send. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # events, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) @@ -9923,9 +9921,8 @@ async def _stream(): captured_finish_reason = None gen = run_gen() - # Watcher to cancel on disconnect: the in-loop poll only fires between - # chunks, so a disconnect during a long prefill/step would otherwise hold - # the decode slot until the next chunk or a failed send. + # Watcher to cancel on disconnect: the in-loop poll fires only between + # chunks, so a mid-prefill disconnect would otherwise hold the decode slot. disconnect_watcher = asyncio.create_task( _await_disconnect_then_cancel(request, cancel_event) ) From a740531788de46eb866ff27fcef9fde7c786bf32 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 16:28:02 +0000 Subject: [PATCH 21/21] Fold marker-collection comment to two lines --- studio/backend/core/tool_healing.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index 706a84e578f..f423ca956bb 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -363,9 +363,8 @@ def parse_tool_calls_from_text( """ tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect JSON/Gemma markers; nesting is decided by _marker_coverage (a - # marker inside another call's coverage is data, not executed). Markers - # inside an open value are that parameter's data. + # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside + # another call's coverage, or an open value, is data not executed. markers = [] # (start, brace_end, kind, match); brace_end < 0 = unclosed for start_re, gemma, kind in ( (_TC_JSON_START_RE, False, "json"),