feat: Add streaming tool-call parse buffer limit to prevent excessive memory usage (#8811)

Author: pskiran1

python/openai/README.md

@@ -688,6 +688,17 @@ python3 openai_frontend/main.py \
   --tool-call-parser llama3
 ```
 
+During streaming, the tool-call parser re-parses the entire accumulated output with each new chunk. This means that very large tool-call arguments can significantly increase per-request CPU and memory usage. The `--max-tool-call-parse-bytes` flag sets a limit on the maximum number of bytes the streaming tool parser will process per request. If a streamed response exceeds this limit, the stream is truncated with finish_reason="length" and backend inference is cancelled.
+The default limit is 128 KiB (131,072 bytes). Increase this value if you expect larger tool-call payloads.
+Example usage:
+```
+python3 openai_frontend/main.py \
+  --model-repository tests/vllm_models \
+  --tokenizer meta-llama/Meta-Llama-3.1-8B-Instruct \
+  --tool-call-parser llama3 \
+  --max-tool-call-parse-bytes 131072
+```
+
 Example for making a tool calling request:
 
 ```python

python/openai/openai_frontend/engine/triton_engine.py

@@ -30,6 +30,7 @@
 import asyncio
 import base64
 import json
+import logging
 import time
 import uuid
 from dataclasses import dataclass
@@ -52,6 +53,7 @@
 from engine.utils.chat import load_chat_template, parse_chat_messages
 from engine.utils.tokenizer import get_tokenizer
 from engine.utils.tool_call_parsers import ToolCallParser, ToolParserManager
+from engine.utils.tool_call_parsers.utils import DEFAULT_MAX_TOOL_CALL_PARSE_BYTES
 from engine.utils.triton import (
     RequestKind,
     TritonLoraConfig,
@@ -96,6 +98,8 @@
 )
 from utils.utils import ClientError, ServerError
 
+logger = logging.getLogger(__name__)
+
 
 # TODO: Improve type hints
 @dataclass
@@ -129,6 +133,7 @@ def __init__(
         lora_separator: Optional[str] = None,
         tool_call_parser: Optional[str] = None,
         chat_template: Optional[str] = None,
+        max_tool_call_parse_bytes: int = DEFAULT_MAX_TOOL_CALL_PARSE_BYTES,
     ):
         # Assume an already configured and started server
         self.server = server
@@ -137,6 +142,7 @@ def __init__(
         self.backend = backend
         self.lora_separator = lora_separator
         self.default_max_tokens = default_max_tokens
+        self.max_tool_call_parse_bytes = max_tool_call_parse_bytes
 
         self.model_metadata = self._get_model_metadata()
         self._metadata_lock = asyncio.Lock()
@@ -147,6 +153,12 @@ def __init__(
         )
         self.chat_template = load_chat_template(chat_template)
 
+        if self.tool_call_parser is not None:
+            print(
+                f"[INFO] Streaming tool-call parse buffer limit set to "
+                f"{self.max_tool_call_parse_bytes} bytes"
+            )
+
     def ready(self) -> bool:
         return self.server.ready()
 
@@ -658,6 +670,7 @@ async def _streaming_chat_iterator(
         )
 
         previous_text = ""
+        tool_parse_truncated = False
         include_usage = request.stream_options and request.stream_options.include_usage
         usage_accumulator = _StreamingUsageAccumulator(backend)
 
@@ -671,6 +684,23 @@ async def _streaming_chat_iterator(
             if include_usage:
                 usage_accumulator.update(response)
 
+            # The auto-tool parser re-parses the full accumulated text on
+            # every chunk. Cap the buffer and cancel the backend to bound
+            # per-request CPU and memory.
+            if (
+                tool_choice_auto
+                and len(previous_text) + len(delta_text)
+                > self.max_tool_call_parse_bytes
+            ):
+                print(
+                    f"[WARNING] Streaming tool-call parse buffer exceeded "
+                    f"{self.max_tool_call_parse_bytes} bytes (request "
+                    f"{request_id}); truncating response and cancelling "
+                    f"backend inference."
+                )
+                tool_parse_truncated = True
+                break
+
             (
                 response_delta,
                 finish_reason,
@@ -717,6 +747,31 @@ async def _streaming_chat_iterator(
             )
             yield f"data: {chunk.model_dump_json(exclude_unset=True)}\n\n"
 
+        # On truncation, cancel the backend to stop it from populating the response
+        # queue. Otherwise, abandoned streams may continue to consume memory.
+        # The final chunk is sent with finish_reason="length" to explicitly
+        # indicate to the client that a cutoff has occurred.
+        if tool_parse_truncated:
+            try:
+                responses.cancel()
+            except Exception:
+                logger.debug(
+                    "Failed to cancel inference after tool-call parse "
+                    "truncation (request %s)",
+                    request_id,
+                    exc_info=True,
+                )
+            choice = ChatCompletionStreamingResponseChoice(
+                index=0,
+                delta=ChatCompletionStreamResponseDelta(content=""),
+                logprobs=None,
+                finish_reason=ChatCompletionFinishReason.length,
+            )
+            chunk = self._get_streaming_chat_response_chunk(
+                choice, request_id, created, model, usage=None
+            )
+            yield f"data: {chunk.model_dump_json(exclude_unset=True)}\n\n"
+
         # Send the final usage chunk if requested via stream_options.
         if include_usage:
             usage_payload = usage_accumulator.get_final_usage()

python/openai/openai_frontend/engine/utils/tool_call_parsers/utils.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -34,6 +34,13 @@
 import partial_json_parser
 from partial_json_parser.core.options import Allow
 
+# Default value for --max-tool-call-parse-bytes CLI flag.
+# Specifies the maximum accumulated output (in bytes) that the
+# streaming tool-call parser processes per request.
+# Since the parser re-parses the entire buffer with each new chunk,
+# this limit helps bound per-request CPU and memory usage.
+DEFAULT_MAX_TOOL_CALL_PARSE_BYTES: int = 128 * 1024  # 128 KiB
+
 
 # partial_json_parser doesn't support extra data and
 # JSONDecorder.raw_decode doesn't support partial JSON

python/openai/openai_frontend/main.py

@@ -33,6 +33,7 @@
 
 import tritonserver
 from engine.triton_engine import TritonLLMEngine
+from engine.utils.tool_call_parsers.utils import DEFAULT_MAX_TOOL_CALL_PARSE_BYTES
 from frontend.fastapi_frontend import FastApiFrontend
 from utils.utils import HTTP_DEFAULT_MAX_INPUT_SIZE, validate_positive_int
 
@@ -134,6 +135,15 @@ def parse_args():
         default=None,
         help="Specify the parser for handling tool calling related response text. Options include: 'llama3' and 'mistral'.",
     )
+    triton_group.add_argument(
+        "--max-tool-call-parse-bytes",
+        type=int,
+        default=DEFAULT_MAX_TOOL_CALL_PARSE_BYTES,
+        help="Maximum accumulated output (in bytes) that the streaming tool-call parser will process per request. "
+        "Once this limit is reached, the stream is truncated with finish_reason='length' and backend inference is cancelled. "
+        "This prevents unbounded memory growth caused by excessively large tool-call arguments. "
+        f"Default: {DEFAULT_MAX_TOOL_CALL_PARSE_BYTES}.",
+    )
     # Allows the user to try a different chat template to craft better prompts and receive more targeted tool-calling responses from the model.
     # Some Mistral models have a separate chat template file, in addition to the tokenizer_config.json,
     # such as mistralai/Mistral-Small-3.1-24B-Instruct-2503.
@@ -266,6 +276,7 @@ def main():
         backend=args.backend,
         lora_separator=args.lora_separator,
         tool_call_parser=args.tool_call_parser,
+        max_tool_call_parse_bytes=args.max_tool_call_parse_bytes,
         chat_template=args.chat_template,
         default_max_tokens=args.default_max_tokens,
     )

python/openai/tests/test_tool_calling.py

@@ -23,18 +23,22 @@
 # OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
 import json
-import os
+from pathlib import Path
 from typing import Dict, List, Optional
 
 import openai
 import pytest
+from fastapi.testclient import TestClient
 from openai.types.chat import (
     ChatCompletionMessageParam,
     ChatCompletionMessageToolCall,
     ChatCompletionNamedToolChoiceParam,
     ChatCompletionToolParam,
 )
+from tests.utils import setup_fastapi_app, setup_server
 
 # resources for testing the tool callings
 WEATHER_TOOL: ChatCompletionToolParam = {
@@ -592,3 +596,82 @@ async def test_inconsistent_tool_choice_and_tools(
                 tools=[],
                 logprobs=False,
             )
+
+
+class TestStreamingToolParseLimit:
+    """
+    Test the streaming tool-call parse-size limit. Utilizes the
+    tool_parser_models/tool_stream_py model, which streams a Mistral-format
+    tool call in small fragments. A prompt containing "big-tool-args" produces
+    an argument large enough to exceed the configured limit, and any other prompt
+    results in an argument that remains within the limit.
+    """
+
+    MODEL = "tool_stream_py"
+    # Small enough that a single tool call can deterministically exceed it.
+    LIMIT_BYTES = 512
+
+    @pytest.fixture(scope="class")
+    def client(self, tokenizer_model: str):
+        model_repository = str(Path(__file__).parent / "tool_parser_models")
+        server = setup_server(model_repository)
+        app = setup_fastapi_app(
+            tokenizer=tokenizer_model,
+            server=server,
+            backend=None,
+            tool_call_parser="mistral",
+            max_tool_call_parse_bytes=self.LIMIT_BYTES,
+        )
+        with TestClient(app) as test_client:
+            yield test_client
+        server.stop()
+
+    def _stream(self, client, content, with_tools):
+        """Streams a chat completion
+        returns (finish_reasons, streamed_text)."""
+        payload = {
+            "model": self.MODEL,
+            "messages": [{"role": "user", "content": content}],
+            "stream": True,
+            "max_completion_tokens": 4096,
+        }
+        if with_tools:
+            payload["tools"] = [WEATHER_TOOL]
+
+        finish_reasons: List[str] = []
+        text = ""
+        with client.stream("POST", "/v1/chat/completions", json=payload) as response:
+            assert response.status_code == 200
+            for line in response.iter_lines():
+                if not line.startswith("data: "):
+                    continue
+                data = line[len("data: ") :]
+                if data.strip() == "[DONE]":
+                    break
+                choice = json.loads(data)["choices"][0]
+                if choice.get("finish_reason"):
+                    finish_reasons.append(choice["finish_reason"])
+                delta = choice.get("delta") or {}
+                if delta.get("content"):
+                    text += delta["content"]
+                for tool_call in delta.get("tool_calls") or []:
+                    arguments = (tool_call.get("function") or {}).get("arguments")
+                    if arguments:
+                        text += arguments
+        return finish_reasons, text
+
+    def test_truncated_when_tool_call_exceeds_limit(self, client):
+        finish_reasons, _ = self._stream(client, "big-tool-args", with_tools=True)
+        # Once the accumulated tool-call text passes the cap, the stream will be terminated with finish_reason="length".
+        assert finish_reasons[-1] == "length"
+
+    def test_not_truncated_within_limit(self, client):
+        finish_reasons, _ = self._stream(client, "small tool call", with_tools=True)
+        # A short tool call stays under the cap and completes normally.
+        assert "length" not in finish_reasons
+
+    def test_limit_ignored_without_tools(self, client):
+        finish_reasons, text = self._stream(client, "big-tool-args", with_tools=False)
+        # Plain streaming does not enter the tool parser, so the cap never fires.
+        assert "length" not in finish_reasons
+        assert len(text) > self.LIMIT_BYTES

python/openai/tests/tool_parser_models/tool_stream_py/1/model.py

@@ -0,0 +1,79 @@
+# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+
+import numpy as np
+import triton_python_backend_utils as pb_utils
+
+# Output is streamed in small fragments to mimic token-by-token generation,
+# which exercises the streaming tool-call parser.
+_CHUNK_SIZE = 8
+
+
+def _decode_prompt(request):
+    tensor = pb_utils.get_input_tensor_by_name(request, "text_input").as_numpy()
+    value = tensor.reshape(-1)[0]
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    return str(value)
+
+
+def _response(text):
+    return pb_utils.InferenceResponse(
+        [
+            pb_utils.Tensor("text_output", np.array([[text]], dtype=object)),
+            pb_utils.Tensor("num_input_tokens", np.array([[1]], dtype=np.int32)),
+            pb_utils.Tensor("num_output_tokens", np.array([[1]], dtype=np.int32)),
+        ]
+    )
+
+
+class TritonPythonModel:
+    """Decoupled mock that streams a Mistral-format tool call.
+
+    The argument length is driven by the prompt so tests can exercise both
+    sides of the streaming tool-call parse-size limit:
+      * a prompt containing "big-tool-args" -> a large argument that exceeds
+        the limit;
+      * any other prompt -> a small argument that stays within it.
+    """
+
+    def execute(self, requests):
+        for request in requests:
+            sender = request.get_response_sender()
+            prompt = _decode_prompt(request)
+            arg = "A" * (2000 if "big-tool-args" in prompt else 5)
+            text = (
+                '[TOOL_CALLS][{"name": "get_current_weather", '
+                '"arguments": {"city": "' + arg + '"}}]'
+            )
+            for start in range(0, len(text), _CHUNK_SIZE):
+                sender.send(_response(text[start : start + _CHUNK_SIZE]))
+            sender.send(
+                _response(""),
+                flags=pb_utils.TRITONSERVER_RESPONSE_COMPLETE_FINAL,
+            )
+        return None