diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6927ba998f9..83c6677eb3c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2402,9 +2402,12 @@ class CustomStreamWrapper: (https://github.com/BerriAI/litellm/issues/13786). Re-raise the failure the stream died of, without logging it a second time. + The traceback is dropped first: raising one retained exception object over + and over would otherwise chain a frame per read onto it, holding every + consumer frame alive and growing the traceback without bound. """ if self._terminal_failure is not None: - raise self._terminal_failure + raise self._terminal_failure.with_traceback(None) def _raise_terminal_failure(self, exc: Exception) -> "NoReturn": self._terminal_failure = exc diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 7ba8b762f26..611edf8a4e1 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2270,6 +2270,43 @@ def test_sync_stream_stays_failed_after_raising_and_logs_failure_once(): assert len(failures) == 1 +def test_failed_stream_traceback_does_not_grow_per_read(): + """A consumer that keeps reading past the failure gets the same size traceback every time.""" + + def looping_stream(): + while True: + yield _make_chunk("the model is stuck on this") + + logging_obj = MagicMock() + logging_obj.completion_start_time = None + logging_obj.failure_handler = lambda exception, traceback_exception: None + + wrapper = CustomStreamWrapper( + completion_stream=looping_stream(), + model="test-model", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + def _traceback_length(error: BaseException) -> int: + length = 0 + tb = error.__traceback__ + while tb is not None: + length += 1 + tb = tb.tb_next + return length + + lengths = [] + for _ in range(litellm.REPEATED_STREAMING_CHUNK_LIMIT + 20): + try: + next(wrapper) + except Exception as e: + lengths.append(_traceback_length(e)) + + assert len(lengths) == 20 + assert len(set(lengths[1:])) == 1 + + def test_usage_chunk_after_finish_reason_updates_hidden_params(logging_obj): """ Test that provider-reported usage from a post-finish_reason chunk