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

litellm/constants.py

@@ -53,7 +53,7 @@ DEFAULT_IMAGE_TOKEN_COUNT = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250))
 # 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)
-MAX_STREAMING_CHUNK_DURATION_S = (
+MAX_STREAMING_DURATION_S = (
     float(_max_stream_duration_env) if _max_stream_duration_env is not None else None
 )
 

litellm/litellm_core_utils/streaming_handler.py

@@ -163,15 +163,15 @@ class CustomStreamWrapper:
         self.created: Optional[int] = None
 
     def _check_max_streaming_duration(self) -> None:
-        """Raise litellm.Timeout if the stream has exceeded MAX_STREAMING_CHUNK_DURATION_S."""
-        from litellm.constants import MAX_STREAMING_CHUNK_DURATION_S
+        """Raise litellm.Timeout if the stream has exceeded MAX_STREAMING_DURATION_S."""
+        from litellm.constants import MAX_STREAMING_DURATION_S
 
-        if MAX_STREAMING_CHUNK_DURATION_S is None:
+        if MAX_STREAMING_DURATION_S is None:
             return
         elapsed = time.time() - self._stream_created_time
-        if elapsed > MAX_STREAMING_CHUNK_DURATION_S:
+        if elapsed > MAX_STREAMING_DURATION_S:
             raise litellm.Timeout(
-                message=f"Stream exceeded max streaming duration of {MAX_STREAMING_CHUNK_DURATION_S}s (elapsed {elapsed:.1f}s)",
+                message=f"Stream exceeded max streaming duration of {MAX_STREAMING_DURATION_S}s (elapsed {elapsed:.1f}s)",
                 model=self.model or "",
                 llm_provider=self.custom_llm_provider or "",
             )
@@ -1251,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"]
@@ -1771,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}"
@@ -1939,7 +1939,7 @@ class CustomStreamWrapper:
                 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
 
@@ -2021,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:

litellm/responses/streaming_iterator.py

@@ -8,7 +8,7 @@ from typing import Any, Dict, Optional
 import httpx
 
 import litellm
-from litellm.constants import MAX_STREAMING_CHUNK_DURATION_S, STREAM_SSE_DONE_STRING
+from litellm.constants import MAX_STREAMING_DURATION_S, 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
@@ -85,13 +85,13 @@ class BaseResponsesAPIStreamingIterator:
         )  # GUARANTEE OPENAI HEADERS IN RESPONSE
 
     def _check_max_streaming_duration(self) -> None:
-        """Raise litellm.Timeout if the stream has exceeded MAX_STREAMING_CHUNK_DURATION_S."""
-        if MAX_STREAMING_CHUNK_DURATION_S is None:
+        """Raise litellm.Timeout if the stream has exceeded MAX_STREAMING_DURATION_S."""
+        if MAX_STREAMING_DURATION_S is None:
             return
         elapsed = time.time() - self._stream_created_time
-        if elapsed > MAX_STREAMING_CHUNK_DURATION_S:
+        if elapsed > MAX_STREAMING_DURATION_S:
             raise litellm.Timeout(
-                message=f"Stream exceeded max streaming duration of {MAX_STREAMING_CHUNK_DURATION_S}s (elapsed {elapsed:.1f}s)",
+                message=f"Stream exceeded max streaming duration of {MAX_STREAMING_DURATION_S}s (elapsed {elapsed:.1f}s)",
                 model=self.model or "",
                 llm_provider=self.custom_llm_provider or "",
             )
@@ -476,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:
@@ -484,6 +485,7 @@ class SyncResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator):
                     self.finished = True
                     raise StopIteration
 
+                self._check_max_streaming_duration()
                 result = self._process_chunk(chunk)
 
                 if self.finished:

tests/test_litellm/litellm_core_utils/test_max_streaming_duration.py

@@ -0,0 +1,126 @@
+"""
+Tests for MAX_STREAMING_DURATION_S — 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 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.MAX_STREAMING_DURATION_S", 10.0
+        ):
+            with pytest.raises(litellm.Timeout, match="max streaming duration"):
+                it._check_max_streaming_duration()