From cbdaaaeba4ff5dc555533c74cb08c1f4b9456954 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 26 Feb 2026 11:39:09 -0800 Subject: [PATCH] [Feat] Add control for setting upperbound on chunk processing time (#22209) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add LITELLM_MAX_STREAMING_DURATION_SECONDS * add add LITELLM_MAX_STREAMING_DURATION_SECONDS * fix: address Greptile review - rename constant, add sync check, add tests - Rename MAX_STREAMING_CHUNK_DURATION_S → MAX_STREAMING_DURATION_S (misleading "CHUNK") - Add _check_max_streaming_duration to SyncResponsesAPIStreamingIterator.__next__ - Add 8 unit tests covering both CustomStreamWrapper and ResponsesAPI paths - Fix pre-existing pyright errors in streaming_handler.py Made-with: Cursor * add add LITELLM_MAX_STREAMING_DURATION_SECONDS --- litellm/constants.py | 8 ++ .../litellm_core_utils/streaming_handler.py | 35 +++-- litellm/responses/streaming_iterator.py | 20 ++- .../test_max_streaming_duration.py | 126 ++++++++++++++++++ 4 files changed, 179 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py diff --git a/litellm/constants.py b/litellm/constants.py index b1a0021bcc..3d2cebf222 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -49,6 +49,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int( ) DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) +# Maximum wall-clock seconds a streaming response is allowed to run. +# Streams exceeding this duration are terminated with a Timeout error. +# None (default) = no limit. Set env var to a number of seconds to enable globally. +_max_stream_duration_env = os.getenv("LITELLM_MAX_STREAMING_DURATION_SECONDS", None) +LITELLM_MAX_STREAMING_DURATION_SECONDS = ( + float(_max_stream_duration_env) if _max_stream_duration_env is not None else None +) + # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. # Set to 0 to disable truncation. diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index ccd5c1dd8f..baf274f2c6 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -96,6 +96,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._stream_created_time: float = time.time() litellm_params: GenericLiteLLMParams = GenericLiteLLMParams( **self.logging_obj.model_call_details.get("litellm_params", {}) @@ -161,6 +162,20 @@ class CustomStreamWrapper: self.is_function_call = self.check_is_function_call(logging_obj=logging_obj) self.created: Optional[int] = None + def _check_max_streaming_duration(self) -> None: + """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" + from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS + + if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)", + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + ) + def __iter__(self) -> Iterator["ModelResponseStream"]: return self @@ -1236,27 +1251,27 @@ class CustomStreamWrapper: else: completion_obj["content"] = str(chunk) elif self.custom_llm_provider == "petals": - if len(self.completion_stream) == 0: + if self.completion_stream is None or len(self.completion_stream) == 0: if self.received_finish_reason is not None: raise StopIteration else: self.received_finish_reason = "stop" chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] + new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] + self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] elif self.custom_llm_provider == "palm": # fake streaming response_obj = {} - if len(self.completion_stream) == 0: + if self.completion_stream is None or len(self.completion_stream) == 0: if self.received_finish_reason is not None: raise StopIteration else: self.received_finish_reason = "stop" chunk_size = 30 - new_chunk = self.completion_stream[:chunk_size] + new_chunk = self.completion_stream[:chunk_size] # type: ignore[index] completion_obj["content"] = new_chunk - self.completion_stream = self.completion_stream[chunk_size:] + self.completion_stream = self.completion_stream[chunk_size:] # type: ignore[index] elif self.custom_llm_provider == "triton": response_obj = self.handle_triton_stream(chunk) completion_obj["content"] = response_obj["text"] @@ -1743,6 +1758,7 @@ class CustomStreamWrapper: and self.custom_llm_provider == "cached_response" ): cache_hit = True + self._check_max_streaming_duration() try: if self.completion_stream is None: self.fetch_sync_stream() @@ -1755,7 +1771,7 @@ class CustomStreamWrapper: ): chunk = self.completion_stream else: - chunk = next(self.completion_stream) + chunk = next(self.completion_stream) # type: ignore[arg-type] if chunk is not None and chunk != b"": print_verbose( f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk.decode('utf-8', errors='replace') if isinstance(chunk, bytes) else chunk}; custom_llm_provider: {self.custom_llm_provider}" @@ -1917,12 +1933,13 @@ class CustomStreamWrapper: and self.custom_llm_provider == "cached_response" ): cache_hit = True + self._check_max_streaming_duration() try: if self.completion_stream is None: await self.fetch_stream() if is_async_iterable(self.completion_stream): - async for chunk in self.completion_stream: + async for chunk in self.completion_stream: # type: ignore[union-attr] if chunk == "None" or chunk is None: continue # skip None chunks @@ -2004,7 +2021,7 @@ class CustomStreamWrapper: ): chunk = self.completion_stream else: - chunk = next(self.completion_stream) + chunk = next(self.completion_stream) # type: ignore[arg-type] if chunk is not None and chunk != b"": processed_chunk = self.chunk_creator(chunk=chunk) if processed_chunk is None: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index edcbb0d11b..43ef4610b4 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1,5 +1,6 @@ import asyncio import json +import time import traceback from datetime import datetime from typing import Any, Dict, Optional @@ -7,7 +8,7 @@ from typing import Any, Dict, Optional import httpx import litellm -from litellm.constants import STREAM_SSE_DONE_STRING +from litellm.constants import LITELLM_MAX_STREAMING_DURATION_SECONDS, STREAM_SSE_DONE_STRING from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -56,6 +57,7 @@ class BaseResponsesAPIStreamingIterator: self.completed_response: Optional[ResponsesAPIStreamingResponse] = None self.start_time = getattr(logging_obj, "start_time", datetime.now()) self._failure_handled = False # Track if failure handler has been called + self._stream_created_time: float = time.time() # track request context for hooks self.litellm_metadata = litellm_metadata @@ -82,6 +84,18 @@ class BaseResponsesAPIStreamingIterator: self.response.headers or {} ) # GUARANTEE OPENAI HEADERS IN RESPONSE + def _check_max_streaming_duration(self) -> None: + """Raise litellm.Timeout if the stream has exceeded LITELLM_MAX_STREAMING_DURATION_SECONDS.""" + if LITELLM_MAX_STREAMING_DURATION_SECONDS is None: + return + elapsed = time.time() - self._stream_created_time + if elapsed > LITELLM_MAX_STREAMING_DURATION_SECONDS: + raise litellm.Timeout( + message=f"Stream exceeded max streaming duration of {LITELLM_MAX_STREAMING_DURATION_SECONDS}s (elapsed {elapsed:.1f}s)", + model=self.model or "", + llm_provider=self.custom_llm_provider or "", + ) + def _process_chunk(self, chunk) -> Optional[ResponsesAPIStreamingResponse]: """Process a single chunk of data from the stream""" if not chunk: @@ -357,6 +371,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): async def __anext__(self) -> ResponsesAPIStreamingResponse: try: + self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: @@ -365,6 +380,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = True raise StopAsyncIteration + self._check_max_streaming_duration() result = self._process_chunk(chunk) if self.finished: @@ -460,6 +476,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __next__(self): try: + self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: @@ -468,6 +485,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = True raise StopIteration + self._check_max_streaming_duration() result = self._process_chunk(chunk) if self.finished: diff --git a/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py new file mode 100644 index 0000000000..f09bdfae64 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py @@ -0,0 +1,126 @@ +""" +Tests for LITELLM_MAX_STREAMING_DURATION_SECONDS — the global cap on streaming response wall-clock time. + +Covers: + - CustomStreamWrapper (chat/completions) sync + async + - BaseResponsesAPIStreamingIterator (responses) sync + async +""" + +import os +import sys +import time +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +import litellm +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_custom_stream_wrapper() -> CustomStreamWrapper: + """Build a minimal CustomStreamWrapper for testing.""" + return CustomStreamWrapper( + completion_stream=None, + model="test-model", + logging_obj=MagicMock(), + custom_llm_provider="openai", + ) + + +# --------------------------------------------------------------------------- +# CustomStreamWrapper (chat/completions) +# --------------------------------------------------------------------------- + + +class TestCustomStreamWrapperMaxDuration: + def test_should_not_raise_when_duration_is_none(self): + """No limit configured → never raises.""" + wrapper = _make_custom_stream_wrapper() + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", None): + wrapper._check_max_streaming_duration() # should not raise + + def test_should_not_raise_when_under_limit(self): + """Stream is under the limit → no error.""" + wrapper = _make_custom_stream_wrapper() + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0): + wrapper._check_max_streaming_duration() # should not raise + + def test_should_raise_timeout_when_exceeded(self): + """Stream exceeded the limit → litellm.Timeout.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 # simulate 20s elapsed + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout, match="max streaming duration"): + wrapper._check_max_streaming_duration() + + def test_should_raise_on_sync_next_when_exceeded(self): + """__next__ should check the limit before iterating.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout): + wrapper.__next__() + + @pytest.mark.asyncio + async def test_should_raise_on_async_anext_when_exceeded(self): + """__anext__ should check the limit before iterating.""" + wrapper = _make_custom_stream_wrapper() + wrapper._stream_created_time = time.time() - 20 + with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): + with pytest.raises(litellm.Timeout): + await wrapper.__anext__() + + +# --------------------------------------------------------------------------- +# BaseResponsesAPIStreamingIterator (responses) +# --------------------------------------------------------------------------- + +class TestResponsesStreamingIteratorMaxDuration: + def _make_base_iterator(self): + """Build a minimal BaseResponsesAPIStreamingIterator for testing.""" + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + mock_response = MagicMock() + mock_response.headers = {} + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {"litellm_params": {}} + mock_logging_obj.start_time = time.time() + + mock_provider_config = MagicMock() + return BaseResponsesAPIStreamingIterator( + response=mock_response, + model="test-model", + responses_api_provider_config=mock_provider_config, + logging_obj=mock_logging_obj, + custom_llm_provider="openai", + ) + + def test_should_not_raise_when_duration_is_none(self): + it = self._make_base_iterator() + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", None + ): + it._check_max_streaming_duration() + + def test_should_not_raise_when_under_limit(self): + it = self._make_base_iterator() + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 60.0 + ): + it._check_max_streaming_duration() + + def test_should_raise_timeout_when_exceeded(self): + it = self._make_base_iterator() + it._stream_created_time = time.time() - 20 + with patch( + "litellm.responses.streaming_iterator.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0 + ): + with pytest.raises(litellm.Timeout, match="max streaming duration"): + it._check_max_streaming_duration()