From 35e18e7680f97eb3f067f278e5539b7ee37bb986 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 23 Jun 2026 13:28:49 +0000 Subject: [PATCH] Quote-aware Gemma strip, symmetric unstarted cleanup, ReDoS anchor Address review findings on the tool-strip and streaming paths: - strip_tool_call_markup stripped Gemma-native spans with a plain regex that stops at the first , so a literal close marker inside a <|"|>-quoted argument truncated the span and leaked its suffix into visible text. A brace/quote-aware _strip_gemma_native_spans now removes complete spans (keeping an incomplete one unless final), matching the parser's own balance logic. - The Gemma close pattern this PR added (<\|tool_call>.*?) had no \Z fallback, so a run of unclosed markers backtracked from every open position (quadratic, and the streaming stripper re-scans per token). It is now anchored to (?:|\Z) like routes/inference.py's _TOOL_XML_RE, linear with identical output on well-formed input. - _SameTaskStreamingResponse added unstarted_cleanup for the OpenAI passthrough, but the local GGUF/safetensors streams that enter _TrackedCancel before returning only unregister in the generator finally, which never runs if the client disconnects before the body iterator starts, leaking cancel-registry entries. Each such stream now passes unstarted_cleanup to exit its tracker. - __call__ reads _unstarted_cleanup via getattr so a response built through __new__ (the cancel-timing test) without __init__ does not raise AttributeError; the test also sets the attribute explicitly. - Document that the verbatim /v1/chat/completions passthrough delegates /<|tool_call> splitting to llama-server (--jinja, --reasoning-format auto) and is intentionally not re-parsed locally, noting the llama.cpp dependency. Adds a regression test for the close-marker-inside-quoted-argument strip. --- studio/backend/core/tool_healing.py | 56 ++++++++++++++++++- studio/backend/routes/inference.py | 40 ++++++++++++- .../tests/test_gemma_tool_parse_edge_cases.py | 13 +++++ .../test_stream_cancel_registration_timing.py | 1 + 4 files changed, 106 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index fe26d48c7f..fe8b94a659 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -13,15 +13,25 @@ import re # Pre-compiled patterns for tool XML stripping. The hyphen in the name # char-class lets dashed MCP tool/parameter names (mcp__srv__list-issues, # issue-number) parse alongside the built-ins. +# +# The Gemma close marker is anchored to ``(?:|\Z)`` (the safe form +# routes/inference.py's _TOOL_XML_RE uses): the plain ``<\|tool_call>.*?`` +# this PR introduced backtracks from every open position on a run of unclosed +# markers (quadratic, and strip_tool_markup_streaming re-scans the cumulative +# buffer per token), whereas the ``\Z`` alternative lets the first open consume +# to EOF in one linear pass. strip_tool_call_markup additionally strips Gemma +# spans via the brace/quote-aware _strip_gemma_native_spans, so a literal close +# marker inside a <|"|>-quoted argument cannot truncate the span and leak its +# suffix; the regex below is the streaming-stripper fallback. +_TC_GEMMA_CLOSED_PAT = re.compile(r"<\|tool_call>.*?(?:|\Z)", re.DOTALL) _TOOL_CLOSED_PATS = [ re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), + _TC_GEMMA_CLOSED_PAT, re.compile(r""), re.compile(r".*?", re.DOTALL), ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ re.compile(r".*$", re.DOTALL), - re.compile(r"<\|tool_call>.*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] @@ -443,6 +453,41 @@ def parse_tool_calls_from_text( return tool_calls +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """Remove complete Gemma-native ``<|tool_call>call:NAME{...}`` + spans, brace- and quote-balanced so a literal ```` inside a + ``<|"|>``-quoted argument does not truncate the span and leak its suffix + (which the plain ``.*?`` regex does). A span without a balanced closing + ``}`` or a trailing close marker is incomplete: dropped to EOF when + ``final`` (the response is over), otherwise kept verbatim so a call that is + still streaming is not stripped mid-token. + """ + out: list[str] = [] + cursor = 0 + for match in _TC_GEMMA_START_RE.finditer(text): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + if final: + out.append(text[cursor:start]) + cursor = len(text) + continue + tail = text[brace_end + 1 :] + leading_ws = len(tail) - len(tail.lstrip()) + close = _TC_GEMMA_END_TAG_RE.match(tail, leading_ws) + if close is None: + if final: + out.append(text[cursor:start]) + cursor = len(text) + continue + out.append(text[cursor:start]) + cursor = brace_end + 1 + close.end() + out.append(text[cursor:]) + return "".join(out) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -450,7 +495,14 @@ def strip_tool_call_markup(text: str, *, final: bool = False) -> str: When ``final`` is True, trailing incomplete tool-call blocks are removed too, and the result is stripped of surrounding whitespace. """ + # Gemma-native spans are stripped brace/quote-aware first; the regex form is + # not quote-aware and would truncate a span at a close marker inside a quoted + # argument. Skip that regex below and let the remaining patterns handle the + # JSON/XML formats and any orphan close marker. + text = _strip_gemma_native_spans(text, final = final) patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS for pat in patterns: + if pat is _TC_GEMMA_CLOSED_PAT: + continue text = pat.sub("", text) return text.strip() if final else text diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index a90201803a..01e60be873 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -863,9 +863,13 @@ class _SameTaskStreamingResponse(StreamingResponse): aclose = getattr(self.body_iterator, "aclose", None) if aclose is not None: await aclose() - if self._unstarted_cleanup is not None: + # getattr (not self._unstarted_cleanup) so a response built via + # __new__ (some tests, pickling) without __init__ does not raise + # AttributeError here. + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -873,6 +877,20 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """Build an ``unstarted_cleanup`` for a local stream that entered ``tracker`` + (a ``_TrackedCancel``) before returning the response. The generator exits the + tracker in its ``finally``, but that never runs if the client disconnects + before the body iterator starts, leaking the cancel-registry entry. This + exits the tracker on that pre-start path only (mutually exclusive with the + generator's finally, so it never double-exits).""" + + async def _cleanup() -> None: + tracker.__exit__(None, None, None) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -4953,6 +4971,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( audio_input_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5422,6 +5441,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5571,6 +5591,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( gguf_stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -5958,6 +5979,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( sf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_sf_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6150,6 +6172,7 @@ async def openai_chat_completions( return _SameTaskStreamingResponse( stream_chunks(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -9493,6 +9516,19 @@ async def _openai_passthrough_stream( response ``id``, ``finish_reason`` (including ``"tool_calls"``), ``delta.tool_calls``, and any client-requested trailing ``usage`` chunk so the client sees a standard OpenAI response. + + Reasoning/tool-call extraction here is delegated to llama-server: this path + forwards to its ``/v1/chat/completions`` (Studio launches with ``--jinja`` + and ``--reasoning-format auto``), which parses Gemma-native ```` into + ``reasoning_content`` and ``<|tool_call>`` into structured ``tool_calls`` + server-side, so the relayed ``delta.content`` carries no raw markup. This is + deliberately NOT re-parsed with the local reasoning extractor / Gemma parser + (verified end to end on the current llama.cpp build), unlike Studio's own + ``/completion``-level generation paths, which must parse the raw text + themselves. The dependency is on llama.cpp's chat parser: if a future build + or chat template stops splitting ````/``<|tool_call>``, raw markup + would relay into ``content`` and this path would need the local extractor as + a safety net. """ target_url = f"{llama_backend.base_url}/v1/chat/completions" body = _build_openai_passthrough_body( diff --git a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py index 8df8d37a52..d573522bcc 100644 --- a/studio/backend/tests/test_gemma_tool_parse_edge_cases.py +++ b/studio/backend/tests/test_gemma_tool_parse_edge_cases.py @@ -22,6 +22,7 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) from core.inference.tool_call_parser import parse_tool_calls_from_text +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -159,3 +160,15 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): ) calls = parse_tool_calls_from_text(content) assert [c["function"]["name"] for c in calls] == ["python"], calls + + +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # A literal inside a <|"|>-quoted argument must not truncate the + # span: the parser keeps it as data, and stripping must remove the whole span + # (brace/quote-aware), not stop at the inner marker and leak the suffix. + text = '<|tool_call>call:python{code:<|"|>print("")<|"|>}' + calls = parse_tool_calls_from_text(text) + assert len(calls) == 1, calls + assert _args(calls[0]) == {"code": 'print("")'} + assert strip_tool_call_markup("before " + text + " after") == "before after" + assert strip_tool_call_markup("before " + text + " after", final = True) == "before after" diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py index 33deb7af9d..73e60a5b0f 100644 --- a/tests/studio/test_stream_cancel_registration_timing.py +++ b/tests/studio/test_stream_cancel_registration_timing.py @@ -411,6 +411,7 @@ def test_same_task_response_closes_body_iterator_on_send_disconnect(): response = m["_SameTaskStreamingResponse"].__new__(m["_SameTaskStreamingResponse"]) response.body_iterator = agen response.background = None + response._unstarted_cleanup = None async def stream_response(_send): raise OSError("client disconnected")