diff --git a/litellm/constants.py b/litellm/constants.py index 663dbf3d6d1..73f0d1e160e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1676,3 +1676,40 @@ ADVISOR_TOOL_DESCRIPTION: Final[str] = ( "want to verify your reasoning, or face a complex decision. " "Describe your question or challenge clearly in the 'question' field." ) + +# Headers that must be stripped from a provider exception before it's forwarded as +# the proxy's own HTTP response, or they conflict with the framing the proxy sets. +HTTP_FRAMING_HEADERS: Final[frozenset[str]] = frozenset( + { + "content-length", + "transfer-encoding", + "content-encoding", + "content-type", + "set-cookie", + "cookie", + "proxy-authenticate", + "proxy-authorization", + } +) + +# Browser-facing security headers that a malicious or misconfigured upstream +# provider must not be able to set on the proxy's own response. +BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset( + { + "access-control-allow-origin", + "access-control-allow-credentials", + "access-control-allow-methods", + "access-control-allow-headers", + "access-control-expose-headers", + "content-security-policy", + "content-security-policy-report-only", + "clear-site-data", + "strict-transport-security", + "x-frame-options", + "cross-origin-opener-policy", + "cross-origin-embedder-policy", + "cross-origin-resource-policy", + } +) + +UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 076eb9a278f..50eada8018d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -27,6 +27,7 @@ from litellm.constants import ( MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_DATA_PREFIX, + UNSAFE_PROXY_RESPONSE_HEADERS, ) from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer @@ -2689,6 +2690,7 @@ class ProxyBaseLLMRequestProcessing: _response_headers: Final = getattr(_response, "headers", None) if _response_headers: headers = get_response_headers(dict(_response_headers)) + headers = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} headers.update(custom_headers) # Call response headers hook for failure @@ -2704,13 +2706,16 @@ class ProxyBaseLLMRequestProcessing: except Exception: pass + headers = {k: v for k, v in headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} + self._apply_router_cooldown_retry_after(headers, e) if isinstance(e, ProxyException): - e.headers = { + merged_headers = { **e.headers, **{k: v if isinstance(v, str) else str(v) for k, v in headers.items()}, } + e.headers = {k: v for k, v in merged_headers.items() if k.lower() not in UNSAFE_PROXY_RESPONSE_HEADERS} raise e if isinstance(e, HTTPException): diff --git a/litellm/router.py b/litellm/router.py index b0fc33c8bb4..9cde292657c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -19,7 +19,7 @@ import threading import time import traceback from collections import defaultdict -from collections.abc import AsyncGenerator, Callable, Generator, Mapping +from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast @@ -300,6 +300,26 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") +def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: + for chunk in chunks: + if not chunk.choices: + continue + delta = chunk.choices[0].delta + if ( + delta.get("content") + or delta.get("tool_calls") + or delta.get("function_call") + or delta.get("reasoning_content") + or delta.get("thinking_blocks") + or delta.get("reasoning_items") + or delta.get("audio") + or delta.get("images") + or delta.get("annotations") + ): + return True + return False + + class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) @@ -2087,6 +2107,13 @@ class Router: async for item in model_response: yield item except MidStreamFallbackError as e: + if not e.is_pre_first_chunk and ( + e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) + ): + if e.original_exception is not None: + raise e.original_exception from e + raise + from litellm.main import stream_chunk_builder complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) @@ -2105,24 +2132,7 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - if e.is_pre_first_chunk or not e.generated_content: - # No content was generated before the error (e.g. a - # rate-limit 429 on the very first chunk). Retry with - # the original messages — adding a continuation prompt - # would waste tokens and confuse the model. - initial_kwargs["messages"] = messages - else: - initial_kwargs["messages"] = messages + [ - { - "role": "system", - "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", - }, - { - "role": "assistant", - "content": e.generated_content, - "prefix": True, - }, - ] + initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( e=e, @@ -2642,6 +2652,13 @@ class Router: for item in model_response: yield item except MidStreamFallbackError as e: + if not e.is_pre_first_chunk and ( + e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) + ): + if e.original_exception is not None: + raise e.original_exception from e + raise + from litellm.main import stream_chunk_builder complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) @@ -2661,20 +2678,7 @@ class Router: router_self.content_policy_fallbacks, ) initial_kwargs["original_function"] = router_self._completion - if e.is_pre_first_chunk or not e.generated_content: - initial_kwargs["messages"] = messages - else: - initial_kwargs["messages"] = messages + [ - { - "role": "system", - "content": "You are a helpful assistant. You are given a message and you need to respond to it. You are also given a generated content. You need to respond to the message in continuation of the generated content. Do not repeat the same content. Your response should be in continuation of this text: ", - }, - { - "role": "assistant", - "content": e.generated_content, - "prefix": True, - }, - ] + initial_kwargs["messages"] = messages router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = router_self.function_with_fallbacks( **initial_kwargs, @@ -2872,6 +2876,13 @@ class Router: llm_provider="", ) + if ( + isinstance(response, CustomStreamWrapper) + and response.completion_stream is None + and response.make_call is not None + ): + await response.fetch_stream() + self.success_calls[model_name] += 1 verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) # debug how often this deployment picked @@ -6109,7 +6120,7 @@ class Router: """ Common utilities for async_function_with_fallbacks """ - verbose_router_logger.debug("Traceback%s", traceback.format_exc()) + verbose_router_logger.debug("Traceback", exc_info=True) original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") # type: ignore @@ -6325,15 +6336,17 @@ class Router: except Exception as new_exception: parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) fallback_failure_exception_str = redact_string(str(new_exception)) + cooldown_info = await _async_get_cooldown_deployments_with_debug_info( + litellm_router_instance=self, + parent_otel_span=parent_otel_span, + ) verbose_router_logger.error( - "litellm.router.py::async_function_with_fallbacks() - Error occurred while trying to do fallbacks - {}\n{}\n\nDebug Information:\nCooldown Deployments={}".format( - fallback_failure_exception_str, - redact_string(traceback.format_exc()), - await _async_get_cooldown_deployments_with_debug_info( - litellm_router_instance=self, - parent_otel_span=parent_otel_span, - ), - ) + "litellm.router.py::async_function_with_fallbacks() - " + "Error occurred while trying to do fallbacks - %s\n" + "Debug Information:\nCooldown Deployments=%s", + fallback_failure_exception_str, + cooldown_info, + exc_info=True, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 4d98a05da8d..93735a84ef9 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2425,14 +2425,14 @@ class TestHandleLLMApiExceptionDictDetail: through ProxyException instead of being str()-mangled into a Python repr. """ - async def _invoke(self, exc: Exception): + async def _invoke(self, exc: Exception, callback_headers: Optional[dict] = None): from litellm.proxy._types import ProxyException, UserAPIKeyAuth processor = ProxyBaseLLMRequestProcessing(data={}) user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={}) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) try: await processor._handle_llm_api_exception( @@ -2952,6 +2952,112 @@ class TestHandleLLMApiExceptionRetryAfter: assert proxy_exc.headers["x-custom"] == "1" +class TestHandleLLMApiExceptionFramingHeaders: + """HTTP-framing headers on the provider exception must be stripped before the + proxy builds its own response, or they conflict with the framing the proxy + itself sets. Non-framing headers must survive unchanged.""" + + async def _invoke(self, exc: Exception, callback_headers: Optional[dict] = None): + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + + processor = ProxyBaseLLMRequestProcessing(data={}) + user_api_key_dict = UserAPIKeyAuth(api_key="sk-test") + proxy_logging_obj = MagicMock() + proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) + proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=callback_headers or {}) + + try: + await processor._handle_llm_api_exception( + e=exc, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + except ProxyException as raised: + return raised + raise AssertionError("ProxyException was not raised") + + async def test_strips_framing_headers_preserves_others(self): + exc = litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + exc.headers = { + "content-length": "42", + "transfer-encoding": "chunked", + "content-encoding": "gzip", + "content-type": "application/json", + "x-request-id": "abc-123", + } + proxy_exc = await self._invoke(exc) + assert "content-length" not in proxy_exc.headers + assert "transfer-encoding" not in proxy_exc.headers + assert "content-encoding" not in proxy_exc.headers + assert "content-type" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + + async def test_strips_framing_headers_on_existing_proxy_exception(self): + from litellm.proxy._types import ProxyException + + exc = ProxyException( + message="Resource exhausted", + type="rate_limit_error", + param=None, + code=429, + headers={ + "content-length": "42", + "transfer-encoding": "chunked", + "x-request-id": "abc-123", + }, + ) + proxy_exc = await self._invoke(exc) + assert "content-length" not in proxy_exc.headers + assert "transfer-encoding" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + + async def test_strips_browser_security_headers(self): + exc = litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + exc.headers = { + "access-control-allow-origin": "https://evil.example.com", + "content-security-policy": "default-src https://evil.example.com", + "clear-site-data": '"cache", "cookies", "storage"', + "strict-transport-security": "max-age=0", + "x-frame-options": "ALLOWALL", + "x-request-id": "abc-123", + } + proxy_exc = await self._invoke(exc) + assert "access-control-allow-origin" not in proxy_exc.headers + assert "content-security-policy" not in proxy_exc.headers + assert "clear-site-data" not in proxy_exc.headers + assert "strict-transport-security" not in proxy_exc.headers + assert "x-frame-options" not in proxy_exc.headers + assert proxy_exc.headers["x-request-id"] == "abc-123" + + async def test_strips_unsafe_headers_added_by_response_headers_hook(self): + exc = litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + exc.headers = {"x-request-id": "abc-123"} + proxy_exc = await self._invoke( + exc, + callback_headers={ + "x-frame-options": "ALLOWALL", + "content-length": "42", + "x-custom-safe": "1", + }, + ) + assert "x-frame-options" not in proxy_exc.headers + assert "content-length" not in proxy_exc.headers + assert proxy_exc.headers["x-custom-safe"] == "1" + assert proxy_exc.headers["x-request-id"] == "abc-123" + + class TestAsyncStreamingDataGeneratorFastPath: """Fast/slow path branching in async_streaming_data_generator.""" diff --git a/tests/test_litellm/test_redact_string_in_error_paths.py b/tests/test_litellm/test_redact_string_in_error_paths.py index acedb285dd4..4a624017ea2 100644 --- a/tests/test_litellm/test_redact_string_in_error_paths.py +++ b/tests/test_litellm/test_redact_string_in_error_paths.py @@ -5,8 +5,10 @@ Covers actual execution of redaction in: - WebSocket close reasons in realtime handlers (openai, azure, bedrock) - Gemini RAG ingestion x-goog-api-key header usage - Traceback redaction pattern used in proxy streaming +- Router fallback-failure traceback redaction """ +import logging import os import sys import traceback @@ -190,6 +192,55 @@ class TestProxyStreamingDataGeneratorRedaction: assert "RuntimeError" in redacted_tb +class TestRouterFallbackFailureTracebackRedaction: + """Test the fallback-failure error log in router.py's + async_function_with_fallbacks_common_utils. A prior version passed exc_info=True + alongside an already-redacted message, which bypasses redact_string() entirely + since the stdlib logging module renders exc_info separately from the message.""" + + @pytest.mark.asyncio + async def test_fallback_failure_does_not_leak_secret_via_exc_info(self, caplog): + import litellm + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, + }, + { + "model_name": "claude-3-haiku", + "litellm_params": {"model": "anthropic/claude-3-haiku-20240307", "api_key": "fake-key"}, + }, + ], + ) + + secret = "sk-testsecretvalue1234567890abcdef" + + with patch( + "litellm.router.run_async_fallback", + new=AsyncMock(side_effect=RuntimeError(f"boom api_key={secret}")), + ): + with caplog.at_level(logging.ERROR, logger="LiteLLM Router"): + with pytest.raises(Exception): + await router.async_function_with_fallbacks_common_utils( + e=Exception("original failure"), + disable_fallbacks=False, + fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}], + context_window_fallbacks=None, + content_policy_fallbacks=None, + model_group="gpt-3.5-turbo", + args=(), + kwargs={"model": "gpt-3.5-turbo"}, + ) + + error_records = [r for r in caplog.records if r.levelno == logging.ERROR] + assert error_records, "expected an error log for the fallback failure" + for record in error_records: + assert secret not in record.getMessage() + assert secret not in (record.exc_text or "") + + def _make_mock_ingest_options(): mock = MagicMock() mock.vector_store_config = {} diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 46b5ce65c3f..33aab1cf708 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -1782,10 +1782,12 @@ async def test_acompletion_streaming_iterator(): assert all(chunk in mock_chunks for chunk in collected_chunks) print("✓ Successfully streamed all chunks") - # Test 2: MidStreamFallbackError with fallback - print("\n=== Test 2: MidStreamFallbackError with fallback ===") + # Test 2: MidStreamFallbackError with generated content is re-raised, not silently continued + print("\n=== Test 2: MidStreamFallbackError re-raises when content already generated ===") - # Create error that should trigger after first chunk + # Error with generated content and is_pre_first_chunk=False (the default): + # the router must re-raise instead of attempting a continuation-prompt fallback, + # because partial content has already been sent to the client. error = MidStreamFallbackError( message="Connection lost", model="gpt-4", @@ -1812,66 +1814,109 @@ async def test_acompletion_streaming_iterator(): self.index += 1 return item - mock_error_response = AsyncIteratorWithError( - mock_chunks, 1 - ) # Error after first chunk + mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk setattr(mock_error_response, "model", "gpt-4") setattr(mock_error_response, "custom_llm_provider", "openai") setattr(mock_error_response, "logging_obj", MagicMock()) - # Mock the fallback response - fallback_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content=" world"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), - ] - - mock_fallback_response = AsyncIterator(fallback_chunks) - - # Mock the fallback function - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=mock_fallback_response, - ) as mock_fallback_utils: - collected_chunks = [] - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) + result = await router._acompletion_streaming_iterator( + model_response=mock_error_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + # Collect streamed chunks — the first chunk succeeds, then the error re-raises + collected_chunks = [] + with pytest.raises(MidStreamFallbackError): async for chunk in result: collected_chunks.append(chunk) - # Verify fallback was called - assert mock_fallback_utils.called - call_args = mock_fallback_utils.call_args - - # Check that generated content was added to messages - fallback_kwargs = call_args.kwargs["kwargs"] - modified_messages = fallback_kwargs["messages"] - - # Should have original message + system message + assistant message with prefix - assert len(modified_messages) == 3 - assert modified_messages[0] == {"role": "user", "content": "Hello"} - assert modified_messages[1]["role"] == "system" - assert "continuation" in modified_messages[1]["content"] - assert modified_messages[2]["role"] == "assistant" - assert modified_messages[2]["content"] == "Hello" - assert modified_messages[2]["prefix"] == True - - # Verify fallback parameters - assert call_args.kwargs["disable_fallbacks"] == False - assert call_args.kwargs["model_group"] == "gpt-4" - - # Should get original chunk + fallback chunks - assert len(collected_chunks) == 3 # 1 original + 2 fallback - print("✓ Fallback system called correctly with proper message modification") + assert len(collected_chunks) == 1, "one chunk yielded before the error" + print("✓ MidStreamFallbackError re-raised correctly when content was already generated") print("\n=== All tests passed! ===") +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): + """Async: when the mid-stream MidStreamFallbackError wraps a real provider + exception (original_exception), the router must re-raise that original + exception instead of the internal wrapper, so the client sees the + specific error type/code (e.g. RateLimitError) rather than a generic + MidStreamFallbackError.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError, RateLimitError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + set_verbose=True, + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + original_exception = RateLimitError( + message="rate limited", + llm_provider="vertex_ai", + model="gpt-4", + ) + error = MidStreamFallbackError( + message="rate limited", + model="gpt-4", + llm_provider="openai", + original_exception=original_exception, + generated_content="Hello", + ) + + mock_chunks = [ + MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), + MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), + ] + + class AsyncIteratorWithError: + def __init__(self, items, error_after_index): + self.items = items + self.index = 0 + self.error_after_index = error_after_index + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index >= len(self.items): + raise StopAsyncIteration + if self.index == self.error_after_index: + raise error + item = self.items[self.index] + self.index += 1 + return item + + mock_error_response = AsyncIteratorWithError(mock_chunks, 1) + setattr(mock_error_response, "model", "gpt-4") + setattr(mock_error_response, "custom_llm_provider", "openai") + setattr(mock_error_response, "logging_obj", MagicMock()) + + result = await router._acompletion_streaming_iterator( + model_response=mock_error_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + with pytest.raises(RateLimitError) as exc_info: + async for _ in result: + pass + assert exc_info.value is original_exception + assert exc_info.value.type == "throttling_error" + assert exc_info.value.code == "429" + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_edge_cases(): """Test edge cases for _acompletion_streaming_iterator.""" @@ -2113,6 +2158,196 @@ def test_completion_streaming_iterator_preserves_hidden_params(): assert result._hidden_params.get("litellm_call_id") == "test-sync-call" +def test_completion_streaming_iterator_reraises_mid_chunk_error(): + """Sync: MidStreamFallbackError with generated_content and is_pre_first_chunk=False + must be re-raised immediately; the router cannot recover after partial content + has already been sent to the client.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + mid_chunk_error = MidStreamFallbackError( + message="Connection reset", + model="gpt-4", + llm_provider="openai", + generated_content="Hello, I am", + is_pre_first_chunk=False, + ) + + class SyncIteratorMidChunkError: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + + def __iter__(self): + return self + + def __next__(self): + raise mid_chunk_error + + mock_response = SyncIteratorMidChunkError() + + result = router._completion_streaming_iterator( + model_response=mock_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + with pytest.raises(MidStreamFallbackError): + list(result) + + +def test_completion_streaming_iterator_reraises_original_exception_when_available(): + """Sync: when the mid-chunk MidStreamFallbackError wraps a real provider + exception (original_exception), the router must re-raise that original + exception instead of the internal wrapper, so the client sees the + specific error type/code (e.g. RateLimitError) rather than a generic + MidStreamFallbackError.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError, RateLimitError + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + original_exception = RateLimitError( + message="rate limited", + llm_provider="vertex_ai", + model="gpt-4", + ) + mid_chunk_error = MidStreamFallbackError( + message="rate limited", + model="gpt-4", + llm_provider="openai", + original_exception=original_exception, + generated_content="Hello, I am", + is_pre_first_chunk=False, + ) + + class SyncIteratorMidChunkError: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [] + + def __iter__(self): + return self + + def __next__(self): + raise mid_chunk_error + + mock_response = SyncIteratorMidChunkError() + + result = router._completion_streaming_iterator( + model_response=mock_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + with pytest.raises(RateLimitError) as exc_info: + list(result) + assert exc_info.value is original_exception + assert exc_info.value.type == "throttling_error" + assert exc_info.value.code == "429" + + +def test_completion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): + """Sync: a reasoning-only chunk sets is_pre_first_chunk=False without populating + generated_content (which only tracks text deltas). The re-raise guard must still + detect this via the raw chunks on the wrapper, or the router silently retries and + the client receives duplicated/inconsistent output.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError + from litellm.types.utils import Delta, StreamingChoices + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + mid_chunk_error = MidStreamFallbackError( + message="Connection reset", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=False, + ) + + reasoning_chunk = litellm.ModelResponseStream( + id="chatcmpl-partial-1", + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), + ) + ], + ) + + class SyncIteratorNoTextChunkError: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [reasoning_chunk] + + def __iter__(self): + return self + + def __next__(self): + raise mid_chunk_error + + mock_response = SyncIteratorNoTextChunkError() + + with patch.object(router, "function_with_fallbacks") as mock_fallback: + result = router._completion_streaming_iterator( + model_response=mock_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + with pytest.raises(MidStreamFallbackError): + list(result) + + assert not mock_fallback.called, ( + "fallback must not be attempted once any content, text or non-text, has already streamed" + ) + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation(): """When MidStreamFallbackError has is_pre_first_chunk=True, use original messages.""" @@ -2181,6 +2416,81 @@ async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation assert fallback_kwargs["messages"] == messages +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): + """Async: a reasoning-only chunk sets is_pre_first_chunk=False without populating + generated_content (which only tracks text deltas). The re-raise guard must still + detect this via the raw chunks on the wrapper, or the router silently retries and + the client receives duplicated/inconsistent output.""" + from unittest.mock import MagicMock + + from litellm.exceptions import MidStreamFallbackError + from litellm.types.utils import Delta, StreamingChoices + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, + } + ], + ) + + messages = [{"role": "user", "content": "Test"}] + initial_kwargs = {"model": "gpt-4", "stream": True} + + mid_chunk_error = MidStreamFallbackError( + message="Connection reset", + model="gpt-4", + llm_provider="openai", + generated_content="", + is_pre_first_chunk=False, + ) + + reasoning_chunk = litellm.ModelResponseStream( + id="chatcmpl-partial-1", + model="gpt-4", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), + ) + ], + ) + + class AsyncIteratorNoTextChunkError: + def __init__(self): + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + self.chunks = [reasoning_chunk] + + def __aiter__(self): + return self + + async def __anext__(self): + raise mid_chunk_error + + mock_response = AsyncIteratorNoTextChunkError() + + with patch.object(router, "async_function_with_fallbacks_common_utils") as mock_fallback_utils: + iterator = await router._acompletion_streaming_iterator( + model_response=mock_response, + messages=messages, + initial_kwargs=initial_kwargs, + ) + + with pytest.raises(MidStreamFallbackError): + async for _ in iterator: + pass + + assert not mock_fallback_utils.called, ( + "fallback must not be attempted once any content, text or non-text, has already streamed" + ) + + # --------------------------------------------------------------------------- # Shared helpers for the _aresponses_streaming_iterator test suite. # --------------------------------------------------------------------------- @@ -4683,9 +4993,7 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f StreamingChoices( finish_reason=None, index=0, - delta=Delta( - content="The Roman Empire began when", role="assistant" - ), + delta=Delta(content="The Roman Empire began when", role="assistant"), ) ], usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), @@ -4738,56 +5046,28 @@ async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_f assert len(collected) == 1 logging_obj.dispatch_success_handlers.assert_not_called() - # Fallback success: the fallback stream owns success accounting via - # _combine_fallback_usage, so this iterator must not dispatch its own. + # Mid-stream errors with generated content are now re-raised immediately; + # no continuation-prompt fallback is attempted. Success handlers must + # still not be dispatched in this path. model_response, logging_obj = _make_interrupted_model_response() - class _FallbackStream: - def __init__(self, items): - self.items = items - self.index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - item = self.items[self.index] - self.index += 1 - return item - - fallback_stream = _FallbackStream( - [ - litellm.ModelResponseStream( - id="chatcmpl-fallback-1", - model="gpt-3.5-turbo", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content=" continued", role="assistant"), - ) - ], - ) - ] - ) with patch.object( router, "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): + new=AsyncMock(), + ) as mock_fallback: result = await router._acompletion_streaming_iterator( model_response=model_response, messages=messages, initial_kwargs=dict(initial_kwargs), ) collected = [] - async for chunk in result: - collected.append(chunk) + with pytest.raises(MidStreamFallbackError): + async for chunk in result: + collected.append(chunk) - assert len(collected) == 2 + assert len(collected) == 1, "only the partial chunk before the error" + mock_fallback.assert_not_called() logging_obj.dispatch_success_handlers.assert_not_called() @@ -5906,6 +6186,198 @@ class TestRouterRequestTimeoutPropagation: ) +# --------------------------------------------------------------------------- +# Deferred-stream eager-fetch tests +# --------------------------------------------------------------------------- + + +def _make_deferred_stream_wrapper(make_call_fn): + """Return a CustomStreamWrapper with completion_stream=None and the given make_call.""" + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + return CustomStreamWrapper( + completion_stream=None, + model="vertex_ai/gemini-2.0-flash", + logging_obj=logging_obj, + custom_llm_provider="vertex_ai_beta", + make_call=make_call_fn, + ) + + +def _make_router_with_vertex_and_fallback(): + return litellm.Router( + model_list=[ + { + "model_name": "my-gemini", + "litellm_params": { + "model": "vertex_ai/gemini-2.0-flash", + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + }, + { + "model_name": "my-fallback", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + }, + }, + ], + fallbacks=[{"my-gemini": ["my-fallback"]}], + num_retries=0, + ) + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_error_propagates_through_acompletion(): + """Regression: a deferred-stream CustomStreamWrapper whose make_call raises a 429 + must propagate the exception from within _acompletion's except block so that + fail_calls is incremented (i.e., deployment cooldown fires) and the standard + router fallback chain can handle it. + + Before the fix, the HTTP call happened inside __anext__ (outside the except block), + so fail_calls was never incremented. + """ + import litellm as _litellm + + rate_limit_err = _litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + + async def failing_make_call(**kwargs): + raise rate_limit_err + + router = _make_router_with_vertex_and_fallback() + deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=deferred_wrapper, + ): + with pytest.raises(_litellm.RateLimitError): + await router._acompletion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + model_name = router.model_list[0]["litellm_params"]["model"] + assert router.fail_calls[model_name] == 1, ( + "fail_calls must be incremented when the deferred HTTP call fails; " + "without the eager fetch_stream() fix this stays at 0" + ) + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_preserves_original_headers_on_error(): + """Router is used both by the proxy and directly as an SDK. HTTP-framing headers + (Content-Length, Transfer-Encoding, ...) must NOT be stripped at this layer, or + direct SDK callers lose legitimate provider metadata (e.g. content-type, + proxy-authenticate) that only the proxy's own response construction needs to + worry about. Stripping happens in the proxy layer instead + (_handle_llm_api_exception).""" + import litellm as _litellm + + err = _litellm.RateLimitError( + message="Resource exhausted", + llm_provider="vertex_ai", + model="gemini-2.0-flash", + ) + err.headers = { + "content-length": "42", + "transfer-encoding": "chunked", + "content-encoding": "gzip", + "content-type": "application/json", + "x-request-id": "abc-123", + } + + async def failing_make_call(**kwargs): + raise err + + router = _make_router_with_vertex_and_fallback() + deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=deferred_wrapper, + ): + with pytest.raises(_litellm.RateLimitError) as exc_info: + await router._acompletion( + model="vertex_ai/gemini-2.0-flash", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + raised = exc_info.value + headers = getattr(raised, "headers", {}) + assert headers.get("content-length") == "42" + assert headers.get("transfer-encoding") == "chunked" + assert headers.get("content-encoding") == "gzip" + assert headers.get("content-type") == "application/json" + assert headers.get("x-request-id") == "abc-123" + + +@pytest.mark.asyncio +async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): + """When completion_stream is already populated (non-deferred provider), the eager + fetch_stream() call must be skipped entirely; no exception should be raised even + if make_call would fail. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + async def would_fail(**kwargs): + raise RuntimeError("should not be called") + + logging_obj = MagicMock() + logging_obj.model_call_details = {"litellm_params": {}} + + async def noop_aiter(): + return + yield + + already_set_wrapper = CustomStreamWrapper( + completion_stream=noop_aiter(), + model="openai/gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + make_call=would_fail, + ) + + router = litellm.Router( + model_list=[ + { + "model_name": "my-model", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + }, + } + ], + ) + + with patch( + "litellm.acompletion", + new_callable=AsyncMock, + return_value=already_set_wrapper, + ): + result = await router._acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": "Hello"}], + stream=True, + specific_deployment=router.model_list[0], + ) + + assert result is not None, "should return a streaming wrapper without errors" + + class TestAdvisorSubCallCooldown: """Regression for LIT-4565: an advisor orchestration failure must not cool down the selected (healthy) deployment, which would reject unrelated @@ -5976,6 +6448,70 @@ class TestAdvisorSubCallCooldown: assert "dep-1" not in self._cooled_down_ids(router) +def test_stream_chunks_have_generated_content_detects_text_and_non_text(): + from litellm.router import _stream_chunks_have_generated_content + from litellm.types.utils import ( + ChatCompletionDeltaToolCall, + Delta, + Function, + StreamingChoices, + ) + + def _chunk(delta): + return litellm.ModelResponseStream( + id="chatcmpl-1", + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=delta)], + ) + + assert _stream_chunks_have_generated_content([]) is False + + empty_chunk = _chunk(Delta(role="assistant")) + assert _stream_chunks_have_generated_content([empty_chunk]) is False + + text_chunk = _chunk(Delta(content="Hello")) + assert _stream_chunks_have_generated_content([text_chunk]) is True + + reasoning_chunk = _chunk(Delta(reasoning_content="Thinking")) + assert _stream_chunks_have_generated_content([reasoning_chunk]) is True + + tool_call_delta = Delta( + tool_calls=[ + ChatCompletionDeltaToolCall( + id="call_1", + function=Function(name="get_weather", arguments="{}"), + type="function", + index=0, + ) + ] + ) + tool_call_chunk = _chunk(tool_call_delta) + assert _stream_chunks_have_generated_content([tool_call_chunk]) is True + + thinking_delta = Delta(thinking_blocks=[{"type": "thinking", "thinking": "Let me think..."}]) + thinking_chunk = _chunk(thinking_delta) + assert _stream_chunks_have_generated_content([thinking_chunk]) is True + + reasoning_items_delta = Delta(reasoning_items=[{"type": "reasoning", "id": "rs_1"}]) + reasoning_items_chunk = _chunk(reasoning_items_delta) + assert _stream_chunks_have_generated_content([reasoning_items_chunk]) is True + + audio_delta = Delta(audio={"data": "abc123", "expires_at": 1234567890, "transcript": "hello"}) + audio_chunk = _chunk(audio_delta) + assert _stream_chunks_have_generated_content([audio_chunk]) is True + + images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_chunk = _chunk(images_delta) + assert _stream_chunks_have_generated_content([images_chunk]) is True + + annotations_delta = Delta( + annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] + ) + annotations_chunk = _chunk(annotations_delta) + assert _stream_chunks_have_generated_content([annotations_chunk]) is True + + def test_get_configured_token_limits_reads_deployment_model_info(): router = litellm.Router( model_list=[