[Bugfix][Parser] Stream Gemma4 string tool arguments incrementally (#61494)

Signed-off-by: Ben Browning <56071+bbrowning@users.noreply.github.com>
Co-authored-by: Flora Feng <4florafeng@gmail.com>

tests/parser/engine/test_gemma4_streaming_reasoning.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for the unified Gemma4 parser engine."""
 
+import dataclasses
 import json
 from unittest.mock import MagicMock
 
@@ -17,8 +18,14 @@ from tests.parser.engine.streaming_helpers import (
 from vllm.entrypoints.generate.base.protocol import DeltaMessage
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
+    ChatCompletionToolsParam,
+)
+from vllm.parser.gemma4 import (
+    STRING_DELIM,
+    Gemma4Parser,
+    _gemma4_arg_value_extender,
+    gemma4_config,
 )
-from vllm.parser.gemma4 import Gemma4Parser
 
 # ── Special token IDs (arbitrary but consistent) ─────────────────────
 CHANNEL_START_ID = 50  # <|channel>
@@ -1649,3 +1656,219 @@ class TestGemma4IsReasoningEnd:
     )
     def test_is_reasoning_end(self, ids, thinking, ended):
         assert self._parser(thinking).is_reasoning_end(ids) is ended
+
+
+# ── Argument value extender ──────────────────────────────────────────
+
+_D = STRING_DELIM
+
+_EXTENDER_TOOL = ChatCompletionToolsParam(
+    type="function",
+    function={
+        "name": "f",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"type": "string"},
+                "location": {"type": "string"},
+                "n": {"type": "integer"},
+            },
+        },
+    },
+)
+
+_EXTENDER_BODIES = [
+    f"location:{_D}Paris, France{_D}",
+    f"a:{_D}x{_D},b:{_D}hello }} world{_D}",
+    f"a:{_D}smile :-) ok{_D},n:42",
+    f"a:{_D}trailing spaces   {_D},n:42",
+    f'a:{_D}decoys <|x <|"y <|"|z{_D}',
+    f'a:{_D}quote " backslash \\ newline\nunicode 杭州{_D}',
+    f"a:{_D}x{_D},{_D}quoted key{_D}:{_D}val{_D}",
+    f"a:{_D}x{_D},nested:{{inner:{_D}deep value{_D},n:1}}",
+    f"items:[{_D}a{_D},{_D}b, c{_D}],a:{_D}last{_D}",
+    f"a:{_D}call:x{_D},b:{_D}brace {{ colon:{_D}",
+    f"a:{_D}{_D},b:{_D}}}}}{_D}",
+]
+
+_TRUNCATED_BODIES = [
+    f"a:{_D}hello }}",
+    f"a:{_D}hello )",
+    f"a:{_D}hello   ",
+    f'a:{_D}hello<|"',
+    f"a:{_D}x{_D},b:{_D}open",
+    f"a:{_D}",
+]
+
+
+_TOOL_MODES = pytest.mark.parametrize(
+    "tools", [[], [_EXTENDER_TOOL]], ids=["no_schema", "schema"]
+)
+
+
+def _pieces(text: str, size: int) -> list[str]:
+    return [text[i : i + size] for i in range(0, len(text), size)]
+
+
+def _stream_args(parser, request, chunks: list[str], finished: bool = True) -> str:
+    args = []
+    for i, chunk in enumerate(chunks):
+        last = finished and i == len(chunks) - 1
+        delta = parser.parse_delta(chunk, [], request, finished=last)
+        for tool_call in (delta and delta.tool_calls) or []:
+            if tool_call.function and tool_call.function.arguments:
+                args.append(tool_call.function.arguments)
+    return "".join(args)
+
+
+class TestGemma4ArgValueExtender:
+    """The extender streams exactly what ``_gemma4_arg_converter`` would, so
+    the converter-only path, which only differs in when it emits, must agree
+    with it byte for byte on the accumulated arguments."""
+
+    @staticmethod
+    def _parser(tokenizer, tools, extender: bool):
+        config = gemma4_config()
+        if not extender:
+            config = dataclasses.replace(config, arg_value_extender=None)
+        return Gemma4Parser(tokenizer, tools, parser_engine_config=config)
+
+    @_TOOL_MODES
+    @pytest.mark.parametrize("size", [1, 3, 7, 64])
+    @pytest.mark.parametrize("body", _EXTENDER_BODIES)
+    def test_matches_converter_at_every_chunking(
+        self, tool_call_tokenizer, mock_request, body, size, tools
+    ):
+        mock_request.tools = tools
+        text = f"<|tool_call>call:f{{{body}}}<tool_call|>"
+        chunks = _pieces(text, size)
+
+        reference = self._parser(tool_call_tokenizer, tools, extender=True)
+        expected = reference.extract_tool_calls(text, mock_request)
+        expected = expected.tool_calls[0].function.arguments
+
+        with_extender = _stream_args(
+            self._parser(tool_call_tokenizer, tools, extender=True),
+            mock_request,
+            chunks,
+        )
+        converter_only = _stream_args(
+            self._parser(tool_call_tokenizer, tools, extender=False),
+            mock_request,
+            chunks,
+        )
+
+        assert with_extender == converter_only == expected
+
+    @_TOOL_MODES
+    @pytest.mark.parametrize("size", [1, 3, 7])
+    @pytest.mark.parametrize("body", _TRUNCATED_BODIES)
+    def test_truncated_stream_matches_converter(
+        self, tool_call_tokenizer, mock_request, body, size, tools
+    ):
+        """A stream cut inside the value ends with the converter's final
+        flush, so the extender must hold back the whitespace, closer or
+        partial delimiter that flush would drop."""
+        mock_request.tools = tools
+        chunks = _pieces(f"<|tool_call>call:f{{{body}", size)
+
+        with_extender = _stream_args(
+            self._parser(tool_call_tokenizer, tools, extender=True),
+            mock_request,
+            chunks,
+        )
+        converter_only = _stream_args(
+            self._parser(tool_call_tokenizer, tools, extender=False),
+            mock_request,
+            chunks,
+        )
+
+        assert with_extender == converter_only
+        json.loads(with_extender)
+
+    @pytest.mark.parametrize("offset", range(1, len(STRING_DELIM)))
+    def test_split_closing_delimiter(self, tool_call_parser, mock_request, offset):
+        """``STRING_DELIM`` is not a lexer terminal, so the extender may see
+        the close split across deltas and must window across them."""
+        head = f"<|tool_call>call:f{{a:{_D}Paris"
+        tail = ",n:1}<tool_call|>"
+        chunks = [head, _D[:offset], _D[offset:] + tail]
+
+        streamed = _stream_args(tool_call_parser, mock_request, chunks)
+
+        assert streamed == '{"a": "Paris", "n": "1"}'
+
+    def test_string_value_streams_before_it_closes(
+        self, tool_call_tokenizer, mock_request
+    ):
+        """Body deltas carry no structural char, so without the extender the
+        gate holds the whole value until the close."""
+        head = [f"<|tool_call>call:f{{location:{_D}", "Paris"]
+
+        def streamed_before_close(extender: bool) -> str:
+            parser = self._parser(tool_call_tokenizer, [], extender)
+            return _stream_args(parser, mock_request, head, finished=False)
+
+        assert streamed_before_close(extender=True) == '{"location": "Paris'
+        assert streamed_before_close(extender=False) == '{"location": "'
+
+    def test_conversions_do_not_grow_with_value_length(
+        self, tool_call_tokenizer, mock_request
+    ):
+        def streaming_conversions(words: int) -> int:
+            calls = 0
+            base = gemma4_config()
+
+            def counting_converter(raw_args: str, partial: bool) -> str:
+                nonlocal calls
+                if partial:
+                    calls += 1
+                return base.arg_converter(raw_args, partial)
+
+            parser = Gemma4Parser(
+                tool_call_tokenizer,
+                parser_engine_config=dataclasses.replace(
+                    base, arg_converter=counting_converter
+                ),
+            )
+            value = "word, " * words
+            text = f"<|tool_call>call:f{{a:{_D}{value}{_D},n:1}}<tool_call|>"
+            _stream_args(parser, mock_request, _pieces(text, 4))
+            return calls
+
+        assert streaming_conversions(words=40) == streaming_conversions(words=400)
+
+    @pytest.mark.parametrize(
+        "raw_args",
+        [
+            f"a:{_D}closed{_D}",
+            f"a:{_D}x{_D},{_D}open ke",
+            f"nested:{{inner:{_D}deep",
+            f"a:{_D}x{_D},b:{_D}y{_D},a:{_D}dup",
+            "a:tr",
+            'a:<|"',
+        ],
+        ids=["closed", "open_key", "nested", "duplicate_key", "bare", "opener"],
+    )
+    def test_factory_declines_without_open_last_string_value(self, raw_args):
+        assert _gemma4_arg_value_extender(raw_args) is None
+
+    def test_extender_holds_back_what_the_converter_hides(self):
+        extend = _gemma4_arg_value_extender(f"a:{_D}hello")
+        assert extend is not None
+        assert [extend(d) for d in [" }", "x", ")", "(", '<|"', "y", '<|"', "|>"]] == [
+            " ",
+            "}x",
+            "",
+            ")(",
+            "",
+            '<|\\"y',
+            "",
+            None,
+        ]
+
+    def test_factory_holds_back_the_tail_already_hidden(self):
+        extend = _gemma4_arg_value_extender(f"a:{_D}hello }} ")
+        assert extend is not None
+        assert extend("x") == "} x"

tests/parser/engine/test_parser_engine.py

@@ -1685,6 +1685,57 @@ class TestArgDeltaWithConverter:
         ]
 
 
+# ── TestArgValueExtender ──────────────────────────────────────────
+
+
+def _extender_config(structural: frozenset[str] | None = None) -> ParserEngineConfig:
+    """``_kv_converter`` plus an extender for its trailing ``k=v`` record,
+    which is open until whitespace follows it; mirrors
+    ``_glm47_arg_value_extender``."""
+
+    def factory(raw_args: str):
+        if re.search(r"\w+=\S+$", raw_args) is None:
+            return None
+
+        def extend(delta: str) -> str | None:
+            if re.search(r"\s", delta):
+                return None
+            return json.dumps(delta, ensure_ascii=False)[1:-1]
+
+        return extend
+
+    return dataclasses.replace(
+        _converter_config(),
+        arg_value_extender=factory,
+        arg_structural_chars=structural,
+    )
+
+
+class TestArgValueExtender:
+    """Engine contract for ``ParserEngineConfig.arg_value_extender``."""
+
+    def test_live_extender_sees_deltas_without_structural_chars(self):
+        """A delta without structural chars must still reach the live
+        extender: gating it out drops ``el`` from the streamed value and the
+        next accepted delta ``l=`` lands on the gap, so the converted JSON
+        no longer extends what was streamed."""
+        engine = _make_engine(_extender_config(structural=frozenset("= ")))
+        assert _run_streaming_tool(engine, "f", ["a", "=h", "el", "l=", "o "]) == {
+            "a": "hell=o"
+        }
+
+    @pytest.mark.parametrize(
+        "overrides",
+        [
+            {"arg_converter": None},
+            {"stream_arg_deltas": False},
+        ],
+    )
+    def test_extender_requires_streaming_conversion(self, overrides):
+        with pytest.raises(ValueError, match="arg_value_extender"):
+            dataclasses.replace(_extender_config(), **overrides)
+
+
 # ── TestSafeArgPrefix ────────────────────────────────────────────
 
 

vllm/parser/engine/parser_engine.py

@@ -342,8 +342,10 @@ class ParserEngine(Parser):
         their unterminated content instead of buffering long arguments until
         the closing tag arrives.
 
-        Also returns whether the prefix ends inside that streamed trailing
-        string value.
+        Also returns whether the returned prefix stops inside that trailing
+        string, so streamed text may be appended to it.  This holds even when
+        *json_str* closes the string: converters serialise an open raw value
+        as a closed JSON string.
         """
         last_colon = -1
         last_key: str | None = None
@@ -1053,11 +1055,8 @@ class ParserEngine(Parser):
         if not self._stream_arg_deltas:
             return None
 
-        structural = self._arg_structural_chars
-        if structural is not None and structural.isdisjoint(raw_delta):
-            return None
-
         slot = self._tool_slots[idx]
+        structural = self._arg_structural_chars
         if slot.extend_open_value is not None:
             diff = slot.extend_open_value(raw_delta)
             if diff is not None:
@@ -1066,6 +1065,8 @@ class ParserEngine(Parser):
                 slot.append_streamed(diff)
                 return diff
             slot.extend_open_value = None
+        elif structural is not None and structural.isdisjoint(raw_delta):
+            return None
 
         try:
             current_json = converter(slot.args, True)

vllm/parser/engine/parser_engine_config.py

@@ -94,9 +94,17 @@ class ParserEngineConfig:
 
     arg_structural_chars: frozenset[str] | None = None
 
-    # Called with the raw argument text when its converted JSON ends in an
-    # open value; returns an extender for that value, or None. Streaming
-    # skips ``arg_converter`` for every delta the extender accepts.
+    # Called with the raw argument text after each streaming conversion whose
+    # prefix-safe JSON extends the JSON streamed so far; returns an extender
+    # when the text ends in an open value that is the last value in the
+    # converted JSON, else None. While live, the extender receives every delta
+    # in place of ``arg_converter`` until it returns None, meaning the delta
+    # closed the value. Returned text is streamed as is, so it must match what
+    # ``arg_converter`` would serialise for those raw characters. Deltas are
+    # lexer output: a terminal is never split across deltas, but the lexer
+    # flushes a held partial terminal as text at the end of the stream or
+    # before a terminal matched by token ID, so a truncated value may end in
+    # a terminal prefix. Requires ``arg_converter`` and ``stream_arg_deltas``.
     arg_value_extender: Callable[[str], ArgValueExtender | None] | None = None
 
     # Special tokens exempt from auto-drop but not state-machine terminals.
@@ -118,6 +126,15 @@ class ParserEngineConfig:
     # Reject tool calls whose names are absent from the request tools.
     validate_tool_names: bool = False
 
+    def __post_init__(self) -> None:
+        if self.arg_value_extender is not None and (
+            self.arg_converter is None or not self.stream_arg_deltas
+        ):
+            raise ValueError(
+                f"{self.name}: arg_value_extender requires arg_converter "
+                "and stream_arg_deltas=True"
+            )
+
     def terminal_literal(self, name: str) -> str | None:
         """Canonical spelling of terminal *name*, or ``None`` if undeclared."""
         value = self.terminals.get(name)

vllm/parser/gemma4.py

@@ -22,6 +22,7 @@ from vllm.logger import init_logger
 from vllm.parser.engine.events import EventType, SemanticEvent
 from vllm.parser.engine.parser_engine import ParserEngine
 from vllm.parser.engine.parser_engine_config import (
+    ArgValueExtender,
     ParserEngineConfig,
     ParserState,
     Transition,
@@ -66,6 +67,12 @@ def _strip_partial_delim(value: str) -> str:
 
 
 def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
+    return _parse_gemma4_fields(args_str, partial=partial)[0]
+
+
+def _parse_gemma4_fields(
+    args_str: str, *, partial: bool = False
+) -> tuple[dict, str | None]:
     """Parse Gemma4's custom key:value format into a Python dict.
 
     Format examples::
@@ -82,13 +89,15 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
             omitted because they may be incomplete and type-unstable
             (e.g. partial boolean parsed as bare string).
 
-    Returns a dict ready for ``json.dumps()``.
+    Returns a dict ready for ``json.dumps()`` and the key of the trailing
+    top-level ``STRING_DELIM`` value left unterminated, if any.
 
     """
     if not args_str or not args_str.strip():
-        return {}
+        return {}, None
 
     result: dict = {}
+    open_key: str | None = None
     i = 0
     n = len(args_str)
 
@@ -130,6 +139,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
                 if partial:
                     value = _strip_partial_delim(value)
                 result[key] = value
+                open_key = key
                 break
             result[key] = args_str[val_start:end_pos]
             i = end_pos + _DELIM_LEN
@@ -199,7 +209,7 @@ def _parse_gemma4_args(args_str: str, *, partial: bool = False) -> dict:
                 break
             result[key] = raw_val
 
-    return result
+    return result, open_key
 
 
 def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list:
@@ -283,16 +293,64 @@ def _parse_gemma4_array(arr_str: str, *, partial: bool = False) -> list:
     return items
 
 
+def _drop_call_closer(text: str, paren_excess: bool) -> str:
+    """Drop a trailing ``}`` or, with more ``)`` than ``(``, a trailing ``)``.
+
+    The closer of the call body is dropped before parsing, so a value that
+    itself ends in one loses that character until more text arrives.
+    """
+    if text.endswith("}") or (text.endswith(")") and paren_excess):
+        return text[:-1]
+    return text
+
+
 def _gemma4_arg_converter(raw_args: str, partial: bool) -> str:
     """Convert Gemma4 custom arg format to a JSON string."""
     text = raw_args.strip()
-    if text.endswith("}") or text.endswith(")") and text.count("(") < text.count(")"):
-        text = text[:-1]
-
+    text = _drop_call_closer(text, text.count("(") < text.count(")"))
     parsed = _parse_gemma4_args(text, partial=partial)
     return json.dumps(parsed, ensure_ascii=False)
 
 
+def _gemma4_arg_value_extender(raw_args: str) -> ArgValueExtender | None:
+    """Extend the open top-level ``STRING_DELIM`` value at the end of
+    *raw_args*, if it is the last value in the converted JSON.
+
+    The extender streams exactly what ``_gemma4_arg_converter`` would, so it
+    holds back the same tail: trailing whitespace, a dropped closer and a
+    partial ``STRING_DELIM``. ``STRING_DELIM`` is a special token but not a
+    lexer terminal, so unlike a tag it may reach the extender split across
+    deltas; the held tail doubles as the window that detects the close.
+    """
+    text = raw_args.strip()
+    opens, closes = text.count("("), text.count(")")
+    body = _drop_call_closer(text, opens < closes)
+    params, open_key = _parse_gemma4_fields(body, partial=True)
+    if open_key is None or next(reversed(params)) != open_key:
+        return None
+
+    held_len = (
+        len(raw_args)
+        - len(raw_args.rstrip())
+        + len(text)
+        - len(_strip_partial_delim(body))
+    )
+    held = raw_args[len(raw_args) - held_len :] if held_len else ""
+
+    def extend(delta: str) -> str | None:
+        nonlocal held, opens, closes
+        text = held + delta
+        if STRING_DELIM in text:
+            return None
+        opens += delta.count("(")
+        closes += delta.count(")")
+        visible = _strip_partial_delim(_drop_call_closer(text.rstrip(), opens < closes))
+        held = text[len(visible) :]
+        return json.dumps(visible, ensure_ascii=False)[1:-1]
+
+    return extend
+
+
 @functools.cache
 def gemma4_config(thinking: bool = False) -> ParserEngineConfig:
     return ParserEngineConfig(
@@ -396,6 +454,7 @@ def gemma4_config(thinking: bool = False) -> ParserEngineConfig:
             ParserState.TOOL_ARGS: EventType.ARG_VALUE_CHUNK,
         },
         arg_converter=_gemma4_arg_converter,
+        arg_value_extender=_gemma4_arg_value_extender,
         tool_args_json=False,
         arg_structural_chars=frozenset(",:{}[]<"),
         preserve_tokens=frozenset({STRING_DELIM}),
@@ -423,12 +482,8 @@ class Gemma4Parser(ParserEngine):
     ) -> None:
         chat_kwargs = kwargs.get("chat_template_kwargs", {}) or {}
         thinking = chat_kwargs.get("enable_thinking", False)
-        super().__init__(
-            tokenizer,
-            tools,
-            parser_engine_config=gemma4_config(thinking=thinking),
-            **kwargs,
-        )
+        kwargs.setdefault("parser_engine_config", gemma4_config(thinking=thinking))
+        super().__init__(tokenizer, tools, **kwargs)
         vocab = self.vocab
         self._tool_call_token_id: int | None = vocab.get("<|tool_call>")
         self._new_turn_token_id: int | None = vocab.get("<|turn>")