diff --git a/studio/backend/core/inference/tool_call_parser.py b/studio/backend/core/inference/tool_call_parser.py index c31f4b272e..9e82e40de2 100644 --- a/studio/backend/core/inference/tool_call_parser.py +++ b/studio/backend/core/inference/tool_call_parser.py @@ -530,13 +530,23 @@ def parse_tool_calls_from_text( # Formats tool_healing does not cover: ```` (MiniCPM-5 / MiniMax-M2), # Llama-3 and Mistral. Run only after tool_healing found nothing, so a strict-rejected - # call is never re-healed here. + # call is never re-healed here. Blank any JSON/Gemma marker coverage first: markup inside + # a marker's span (even one that failed to parse) is that call's data, not a sibling, so + # a nested ```` / ``<|python_tag|>`` / ``[TOOL_CALLS]`` must not be promoted. + fallback_content = content + coverage = _tool_healing.marker_coverage(content) + if coverage: + chars = list(content) + for cov_start, cov_end in coverage: + for i in range(cov_start, min(cov_end, len(chars))): + chars[i] = " " + fallback_content = "".join(chars) for parser in ( _parse_function_xml, # attribute form _parse_llama3_python_tag, # Llama-3 <|python_tag|> _parse_mistral_tool_calls, # Mistral [TOOL_CALLS] ): - calls = parser(content, id_offset = id_offset, allow_incomplete = allow_incomplete) + calls = parser(fallback_content, id_offset = id_offset, allow_incomplete = allow_incomplete) if calls: return calls diff --git a/studio/backend/core/tool_healing.py b/studio/backend/core/tool_healing.py index ff8faf2308..b91403ed57 100644 --- a/studio/backend/core/tool_healing.py +++ b/studio/backend/core/tool_healing.py @@ -10,20 +10,45 @@ orchestrator, structlog, httpx, or the rest of the studio backend. 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. +# Strip patterns. The name-class hyphen matches dashed MCP names. Closed pairs +# 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) +_TC_GEMMA_END_PAT = re.compile(r"") _TOOL_CLOSED_PATS = [ - re.compile(r".*?", re.DOTALL), - re.compile(r"<\|tool_call>.*?", re.DOTALL), - re.compile(r""), - re.compile(r".*?", re.DOTALL), + _TC_JSON_CLOSED_PAT, + _TC_GEMMA_CLOSED_PAT, + _TC_FUNC_CLOSED_PAT, + _TC_GEMMA_END_PAT, ] _TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [ - re.compile(r".*$", re.DOTALL), re.compile(r"<\|tool_call>.*$", re.DOTALL), + re.compile(r".*$", re.DOTALL), re.compile(r".*$", re.DOTALL), ] +# 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 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: "", + _TC_FUNC_CLOSED_PAT: "", +} + + +def strip_tool_patterns(text: str, patterns) -> str: + """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: + continue + text = pat.sub("", text) + return text + # Pre-compiled patterns for tool-call XML parsing. _TC_JSON_START_RE = re.compile(r"\s*\{") @@ -40,13 +65,9 @@ _TC_PARAM_CLOSE_RE = re.compile(r"\s*\s*$") _GEMMA_QUOTE = '<|"|>' _PARAM_CLOSE_TAG = "" _FUNC_CLOSE_TAG = "" -# A bare (unquoted) Gemma value ends at `}` or at a comma that begins the next -# `key:` pair. A comma NOT followed by a key token is part of the value (e.g. -# `location:New York, NY`), so it must not terminate the value. 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. -# Dots match the key-quoting scanner: a dotted key after a bare value must end the value at the comma. +# A bare (unquoted) Gemma value ends at `}` or at a comma beginning the next +# identifier-shaped `key:` pair; a comma before a non-key (`New York, NY`, +# `10:00, 11:00`) stays in the value. Dots let a dotted key end the value. _GEMMA_NEXT_KEY_RE = re.compile(r"\s*[A-Za-z_][\w.\-]*\s*:") @@ -143,14 +164,8 @@ 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 (``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() @@ -158,11 +173,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]) + "]") @@ -241,15 +254,12 @@ 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 stays 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:]) @@ -259,9 +269,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 @@ -329,6 +337,68 @@ def _trim_param_value(val: str) -> str: return val +def _marker_coverage(content: str, markers) -> list[tuple[int, int]]: + """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 + 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 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())) + 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 _build_markers(content: str): + """JSON/Gemma tool markers as ``(start, brace_end, kind, match)`` in document + order; ``brace_end < 0`` marks an unbalanced (to-EOF) open.""" + markers = [] + for start_re, gemma, kind in ( + (_TC_JSON_START_RE, False, "json"), + (_TC_GEMMA_START_RE, True, "gemma"), + ): + 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) + markers.append((m.start(), brace_end, kind, m)) + markers.sort(key = lambda c: c[0]) + return markers + + +def marker_coverage(content: str) -> list[tuple[int, int]]: + """Coverage spans of JSON/Gemma tool markers so other parsers can treat markup + inside a marker's coverage (even a marker that failed to parse) as that call's + data rather than a sibling call.""" + return _marker_coverage(content, _build_markers(content)) + + def parse_tool_calls_from_text( content: str, *, @@ -350,37 +420,26 @@ def parse_tool_calls_from_text( """ tool_calls: list[dict] = [] call_spans: list[tuple] = [] - # Collect every supported call format with spans, then emit in document - # order. A marker inside another call's argument string is data, not a - # separate executable call. - parsed_items = [] # (start, span_end, name, arguments) - candidates = [] # (start, brace_end, kind, match) - 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)) - 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)) - candidates.sort(key = lambda c: c[0]) - - candidate_spans = [(s, e) for s, e, _kind, _m in candidates] - for idx, (start, end, kind, m) in enumerate(candidates): - if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx): + # Collect JSON/Gemma markers; _marker_coverage decides nesting. A marker inside + # another call's coverage, or an open value, is data not executed. + markers = _build_markers(content) + 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): + # 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: + continue # unclosed: not parseable; the fallback still excludes its XML 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", "") # Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside a Hermes ). arguments = obj.get("arguments") @@ -390,10 +449,11 @@ def parse_tool_calls_from_text( 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 - span_end = end + 1 + # 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()) close_m = close_re.match(content, span_end + ws) @@ -401,11 +461,15 @@ def parse_tool_calls_from_text( span_end = close_m.end() parsed_items.append((start, span_end, name, arguments)) + # 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) if not _inside_open_parameter(content, fm.start()) - and not any(s <= fm.start() <= e for s, e in candidate_spans) + and not any(s <= fm.start() < e for s, e in coverage) ] for idx, fm in enumerate(func_starts): func_name = fm.group(1) @@ -481,90 +545,106 @@ def parse_tool_calls_from_text( ) call_spans.append((start, span_end)) - if not tool_calls: - func_starts = [ - fm - for fm in _TC_FUNC_START_RE.finditer(content) - if not _inside_open_parameter(content, fm.start()) - ] - for idx, fm in enumerate(func_starts): - func_name = fm.group(1) - body_start = fm.end() - next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content) - end_tag = _TC_END_TAG_RE.search(content[body_start:]) - if end_tag: - body_end = body_start + end_tag.start() - else: - body_end = len(content) - body_end = min(body_end, next_func) - body = content[body_start:body_end] - # Span for with_spans callers: through the close if present, else body end. - span_end = body_end - if not allow_incomplete: - close_idx = _func_close_index(content, body_start, body) - if close_idx < 0: - continue - body = body[:close_idx] - span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) - else: - # Terminate at the real close so trailing prose doesn't leak in; no close -> whole body. - close_idx = _func_close_index(content, body_start, body) - if close_idx >= 0: - body = body[:close_idx] - span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG) - - arguments: dict = {} - param_starts = list(_TC_PARAM_START_RE.finditer(body)) - if len(param_starts) == 1: - pm = param_starts[0] - val = body[pm.end() :] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - continue - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[pm.group(1)] = _trim_param_value(val) - else: - valid_params = True - for pidx, pm in enumerate(param_starts): - param_name = pm.group(1) - val_start = pm.end() - next_param = ( - param_starts[pidx + 1].start() - if pidx + 1 < len(param_starts) - else len(body) - ) - val = body[val_start:next_param] - if not allow_incomplete: - stripped_val = val.rstrip() - if not stripped_val.endswith(_PARAM_CLOSE_TAG): - valid_params = False - break - val = stripped_val[: -len(_PARAM_CLOSE_TAG)] - else: - val = _TC_PARAM_CLOSE_RE.sub("", val) - arguments[param_name] = _trim_param_value(val) - if not valid_params: - continue - - tc = { - "id": f"call_{id_offset + len(tool_calls)}", - "type": "function", - "function": { - "name": func_name, - "arguments": json.dumps(arguments), - }, - } - tool_calls.append(tc) - call_spans.append((fm.start(), span_end)) - if with_spans: return tool_calls, call_spans return tool_calls +def _strip_gemma_native_spans(text: str, *, final: bool) -> str: + """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): + start = match.start() + if start < cursor: + continue + brace_end = _balanced_brace_end(text, match.end() - 1, gemma_quotes = True) + if brace_end < 0: + # 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 + # 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: + out.append(text[cursor:start]) + cursor = len(text) + break + out.append(text[cursor:start]) + cursor = close.end() + out.append(text[cursor:]) + 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.""" + 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 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) + 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 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) + + def strip_tool_call_markup(text: str, *, final: bool = False) -> str: """Strip tool-call XML markup from text. @@ -572,7 +652,9 @@ 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. """ - patterns = _TOOL_ALL_PATS if final else _TOOL_CLOSED_PATS - for pat in patterns: - text = pat.sub("", text) - return text.strip() if final else text + if final: + return strip_tool_markup_final(text).strip() + # 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 4393c1b304..1a1a934009 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -852,17 +852,14 @@ class _SameTaskStreamingResponse(StreamingResponse): **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: @@ -873,15 +870,11 @@ class _SameTaskStreamingResponse(StreamingResponse): 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: @@ -893,16 +886,16 @@ class _SameTaskStreamingResponse(StreamingResponse): 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() - if self._unstarted_cleanup is not None: + cleanup = getattr(self, "_unstarted_cleanup", None) + if cleanup is not None: try: - await self._unstarted_cleanup() + await cleanup() except Exception: pass raise ClientDisconnect() @@ -910,6 +903,16 @@ class _SameTaskStreamingResponse(StreamingResponse): await self.background() +def _tracked_cancel_unstarted_cleanup(tracker): + """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) + + return _cleanup + + async def _aclose_stream_resources( *, watchers = (), @@ -4069,12 +4072,9 @@ async def generate_stream( _DONE = object() while True: if cancel_event.is_set(): - # The disconnect watcher set cancel_event between chunks. - # Reset the backend here: closing the Python generator does - # not signal a subprocess backend, so without this it keeps - # decoding after the client is gone. The finally's reset is - # guarded on cancel_event being unset, so it will not run - # again for this path. + # 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) @@ -5684,6 +5684,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", @@ -6163,6 +6164,7 @@ async def openai_chat_completions( if payload.stream: return _SameTaskStreamingResponse( gguf_tool_stream(), + unstarted_cleanup = _tracked_cancel_unstarted_cleanup(_tracker), media_type = "text/event-stream", headers = { "Cache-Control": "no-cache", @@ -6419,6 +6421,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", @@ -6852,6 +6855,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", @@ -7067,6 +7071,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", @@ -9977,11 +9982,8 @@ async def _anthropic_tool_stream( drop_until_tool_end = False gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between events, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # event or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # 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) ) @@ -10073,11 +10075,8 @@ async def _anthropic_plain_stream( captured_finish_reason = None gen = run_gen() - # Concurrent disconnect watcher: the loop only polls is_disconnected() - # between chunks, so a client disconnect during a long prefill or - # generation step would otherwise hold the decode slot until the next - # chunk or a failed send. The watcher sets cancel_event so the backend - # stops promptly. + # 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) ) @@ -11030,6 +11029,10 @@ 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 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( @@ -11446,11 +11449,9 @@ async def _openai_passthrough_stream( delta = choice.get("delta") if isinstance(delta, dict) and delta.get("tool_calls"): saw_tool_call_delta = True - # Detect an upstream error chunk independently of API - # monitoring: when monitor_id is None (skip_api_monitor), - # _monitor_openai_sse_line returns before inspecting the - # error, so without this the synthetic-finish guard would - # emit a successful finish_reason after a failed stream. + # 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 # With healing active, a content-bearing line may be replaced by 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 63df86ec17..fff6b240c5 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 @@ -25,6 +18,7 @@ from core.inference.tool_call_parser import ( _gemma_parse_value, parse_tool_calls_from_text, ) +from core.tool_healing import strip_tool_call_markup def _args(call: dict) -> dict: @@ -43,8 +37,6 @@ 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"} @@ -60,8 +52,7 @@ def test_empty_bare_value_becomes_empty_string_not_dropped(): 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}" ) @@ -70,8 +61,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}"}}' @@ -89,8 +78,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"}}' @@ -100,8 +87,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"}})' @@ -112,18 +97,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"]} @@ -137,8 +118,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}]}" ) @@ -152,9 +131,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}" @@ -175,6 +151,165 @@ def test_json_marker_inside_xml_parameter_is_not_a_second_call(): assert [c["function"]["name"] for c in calls] == ["python"], calls +def test_gemma_close_marker_inside_quoted_arg_is_not_leaked_when_stripping(): + # 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 + 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(): + # The failed Gemma candidate's span still covers its nested . + 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_unbalanced_gemma_call_with_xml_does_not_execute(): + # Unclosed braces cover to EOF, so the trailing is excluded. + 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(): + 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(): + # Coverage runs to the close marker, so in the gap is data. + 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(): + 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(): + # 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(): + assert ( + strip_tool_call_markup('before <|tool_call>{"name":"x"} after') + == "before after" + ) + + +def test_valid_call_after_missing_close_is_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) + ] + 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(): + 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 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"}}' + "" + ) + 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 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) + 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(): + # The to-EOF Gemma sweep must not eat 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" + + +def test_strip_final_keeps_text_after_closed_block_with_call_form_gemma_opener(): + # 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{"}}' + ) + 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 + + +def test_function_sibling_after_close_less_gemma_marker_is_recovered(): + # The close-less marker covers only its braces; the XML sibling is recovered. + 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(): + # 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) + ] + 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 + + def test_gemma_parse_value_always_advances_on_stray_delimiter(): # A stray delimiter (`,`, `}`, `]`) at the primitive position must still advance the # parser, or a looping caller spins forever (DoS). diff --git a/studio/backend/tests/test_tool_call_parser_strict.py b/studio/backend/tests/test_tool_call_parser_strict.py index 7664126d91..fded2a8443 100644 --- a/studio/backend/tests/test_tool_call_parser_strict.py +++ b/studio/backend/tests/test_tool_call_parser_strict.py @@ -1063,3 +1063,44 @@ class TestBareJsonStripRequiresTopLevelName: def test_real_call_still_strips_name_agnostic(self): from core.inference.tool_call_parser import strip_leading_bare_json_call assert strip_leading_bare_json_call('{"name":"web_search","parameters":{"q":"x"}}') == "" + + +class TestGemmaAwareClosedBlockPrePass: + """The closed JSON/function strip pre-pass must not delete across a complete + 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 + 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" 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 0000000000..dfa3101882 --- /dev/null +++ b/studio/backend/tests/test_tool_strip_guard.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. + +"""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 +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(): + # 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 == ""