From bbf2c4fd0c4e4b2c95b706e7e529554fc6b10a7a 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 Made-with: Cursor --- litellm/constants.py | 8 ++ .../litellm_core_utils/streaming_handler.py | 35 +++-- litellm/responses/streaming_iterator.py | 23 +++- .../test_max_streaming_duration.py | 126 ++++++++++++++++++ ui/litellm-dashboard/tsconfig.json | 29 +++- 5 files changed, 205 insertions(+), 16 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 3c84547d7ce..dd4feab0c91 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -46,6 +46,14 @@ DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = int( os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1) ) 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 +) DEFAULT_IMAGE_WIDTH = int(os.getenv("DEFAULT_IMAGE_WIDTH", 300)) DEFAULT_IMAGE_HEIGHT = int(os.getenv("DEFAULT_IMAGE_HEIGHT", 300)) # Maximum size for image URL downloads in MB (default 50MB, set to 0 to disable limit) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 890f10f9329..bc8dc082600 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -85,6 +85,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", {}) @@ -150,6 +151,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): return self @@ -1204,27 +1219,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"] @@ -1711,6 +1726,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() @@ -1723,7 +1739,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}; custom_llm_provider: {self.custom_llm_provider}" @@ -1885,12 +1901,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 @@ -1965,7 +1982,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: Optional[ ModelResponseStream diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6d0c4abac81..2f4aae0e9a4 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,10 @@ 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 +60,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 +87,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: @@ -363,6 +380,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: @@ -371,6 +389,7 @@ class ResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): self.finished = True raise StopAsyncIteration + self._check_max_streaming_duration() result = self._process_chunk(chunk) if self.finished: @@ -461,6 +480,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): def __next__(self): try: + self._check_max_streaming_duration() while True: # Get the next chunk from the stream try: @@ -469,6 +489,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 00000000000..f09bdfae649 --- /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() diff --git a/ui/litellm-dashboard/tsconfig.json b/ui/litellm-dashboard/tsconfig.json index c73661d32ea..d24bdd340f7 100644 --- a/ui/litellm-dashboard/tsconfig.json +++ b/ui/litellm-dashboard/tsconfig.json @@ -1,6 +1,10 @@ { "compilerOptions": { - "lib": ["dom", "dom.iterable", "esnext"], + "lib": [ + "dom", + "dom.iterable", + "esnext" + ], "allowJs": true, "skipLibCheck": true, "strict": true, @@ -10,7 +14,7 @@ "moduleResolution": "bundler", "resolveJsonModule": true, "isolatedModules": true, - "jsx": "preserve", + "jsx": "react-jsx", "incremental": true, "plugins": [ { @@ -18,9 +22,22 @@ } ], "paths": { - "@/*": ["./src/*"] - } + "@/*": [ + "./src/*" + ] + }, + "target": "ES2017" }, - "include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"], - "exclude": ["node_modules", "e2e_tests", "scripts"] + "include": [ + "next-env.d.ts", + "**/*.ts", + "**/*.tsx", + ".next/types/**/*.ts", + ".next/dev/types/**/*.ts" + ], + "exclude": [ + "node_modules", + "e2e_tests", + "scripts" + ] }