diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index aa6ce6ea99e..06e586c0e2a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2316,8 +2316,8 @@ class CustomStreamWrapper: cache_hit = False if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response": cache_hit = True - self._check_max_streaming_duration() try: + self._check_max_streaming_duration() if self.completion_stream is None: self.fetch_sync_stream() diff --git a/tests/unit/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py index 234bc83e3c3..ddaf5868bec 100644 --- a/tests/unit/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/unit/litellm_core_utils/test_max_streaming_duration.py @@ -65,6 +65,22 @@ class TestCustomStreamWrapperMaxDuration: with pytest.raises(litellm.Timeout): wrapper.__next__() + def test_sync_next_timeout_dispatches_failure_handler_and_closes(self): + """The sync Timeout must flow through the failure pipeline (parity with __anext__).""" + wrapper = _make_custom_stream_wrapper() + wrapper.logging_obj.failure_handler = MagicMock() + 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__() + assert wrapper._closed is True + + # failure_handler is dispatched on a thread; give it a moment + deadline = time.time() + 2.0 + while time.time() < deadline and wrapper.logging_obj.failure_handler.call_count == 0: + time.sleep(0.01) + wrapper.logging_obj.failure_handler.assert_called_once() + @pytest.mark.asyncio async def test_should_raise_on_async_anext_when_exceeded(self): """__anext__ checks the limit before iterating and schedules failure logging."""