fix(router): eagerly fetch Vertex AI deferred stream to surface HTTP errors in _acompletion fallback path (#34627)

* fix(router): eagerly fetch deferred stream to surface HTTP errors in fallback path

Providers like Vertex AI and Bedrock defer their HTTP call until the first
__anext__ on the returned CustomStreamWrapper (completion_stream=None,
make_call set). Errors raised inside __anext__ (e.g. 429, 503) escape the
_acompletion try/except block, so fail_calls is never incremented, deployment
cooldown does not fire, and the standard fallback chain is bypassed.

Call fetch_stream() on the wrapper before delegating to
_acompletion_streaming_iterator when completion_stream is None and make_call
is set. Any HTTP error now propagates through _acompletion's except block,
increments fail_calls, and enters the normal retry/fallback chain.

Strip Content-Length, Transfer-Encoding, Content-Encoding, and Content-Type
from exception headers at the same point to prevent HTTP framing mismatches
when LiteLLM builds its own error response body.

Add a re-raise guard in _acompletion_streaming_iterator (async and sync paths)
so MidStreamFallbackError with already-generated content re-raises to the
caller instead of silently injecting a continuation prompt into a fresh request
to a fallback model.

Apply logging cleanup in async_function_with_fallbacks_common_utils: use
%s-style formatting and exc_info=True instead of f-strings with
traceback.format_exc().

* fix(router): undo success_calls on deferred-stream fetch failure; broaden header strip

* fix(router): extract header-strip helper to keep _acompletion under strict C901 threshold

* test(router): add unit tests for _strip_http_framing_headers to satisfy router coverage gate

* test(router): add sync _completion_streaming_iterator re-raise test for mid-chunk MidStreamFallbackError

* fix(router): restore Fallbacks context in no-fallback log; document update_team mcp_rpm_limit

The log and debug message when no fallback model group is found was missing
the Fallbacks list, making it hard to understand why routing failed.

Also adds the missing mcp_rpm_limit documentation to update_team to fix
the documentation_test_api_docs CI check.

* fix(router): preserve original traceback in deferred stream fetch error re-raise

Using bare `raise` instead of `raise fetch_err` keeps the full inner
traceback from fetch_stream() intact so the error origin is visible in
logs and debuggers without being anchored to this line.

* style(test): restore black-style formatting in test_router.py

An earlier commit on this branch collapsed the file's pre-existing
multi-line formatting into single lines while adding the deferred-stream
tests, producing a diff full of unrelated reformatting noise. Restores
the untouched code to its original formatting; the actual new/changed
test content is unaffected (verified via AST comparison).

* fix(router): re-raise mid-stream fallback on any generated content, not just text

The re-raise guard added for MidStreamFallbackError only checked
generated_content, which tracks text deltas alone. A stream that emitted a
tool-call or reasoning-only chunk before failing had generated_content=""
despite already streaming to the client, so the router silently retried
and the client saw duplicated/inconsistent output. The guard now also
inspects the wrapper's raw chunks for tool_calls/reasoning_content.

Also moves the deferred-stream HTTP-framing-header stripping out of
Router._acompletion into the proxy's _handle_llm_api_exception: Router is
used directly as an SDK as well as by the proxy, and stripping headers
there dropped legitimate provider metadata (content-type,
proxy-authenticate) for direct SDK callers who never see the proxy's own
response construction.

schema.d.ts regenerated via make pre-commit; unrelated to this change.

* test(router): add direct coverage for _stream_chunks_have_generated_content

CI's router_code_coverage check flags any router.py function never referenced
by name in a test file; the new helper was only exercised indirectly through
the mid-stream re-raise guard tests.

* revert(ui): drop incidental schema.d.ts regeneration

Committing router.py/common_request_processing.py touched
pre_commit_lint.sh's litellm/proxy trigger for the API-type-sync check,
which force-regenerated schema.d.ts even though neither file changes any
route or model. The regenerated ordering of two unrelated Union/enum
fields (stream_timeout, user_role) isn't stable across process
invocations even against completely unmodified backend code (confirmed
by regenerating twice against the pre-existing committed code and getting
the same diff both times), so this reverts to the original committed
file rather than chase non-deterministic output.

* fix(proxy): strip framing headers on the pre-existing ProxyException branch too

_handle_llm_api_exception filtered framing headers into a local `headers`
dict, but for an exception that's already a ProxyException, it merged
{**e.headers, **headers}: the original e.headers came first, so a framing
header present there but absent from the filtered `headers` (because it
was just stripped) was never overwritten and survived into the response
unfiltered. Filters the merged result instead of relying on the merge
order to do it implicitly.

* chore: retrigger CI (no GitHub Actions check-suite was created for the previous two pushes)

* fix(router): detect thinking_blocks as generated content in mid-stream guard

Greptile flagged that a thinking-only delta (Anthropic extended thinking,
Delta.thinking_blocks) wasn't recognized as already-streamed content, so
a stream that emitted only thinking blocks before failing could still
restart via fallback and append an unrelated response after content the
client already received.

* fix(proxy): strip browser-facing security headers from provider exceptions too

veria-ai flagged that the framing-header denylist still let a malicious or
misconfigured provider set browser-facing headers (Access-Control-Allow-Origin,
Content-Security-Policy, Clear-Site-Data, etc.) on the proxy's own error
response. Adds a dedicated _BROWSER_SECURITY_HEADERS set alongside the
existing framing one and strips both wherever provider exception headers
reach the client response.

* refactor(router): address maintainer review mechanicals

- List[ModelResponseStream] -> list[ModelResponseStream] in
  _stream_chunks_have_generated_content (ruff UP006 strict-budget gate)
- drop _strip_http_framing_headers and its 3 tests: the proxy inlines the
  filter directly now, so the helper has had no production caller since
  the header-stripping was moved out of Router
- move HTTP_FRAMING_HEADERS/BROWSER_SECURITY_HEADERS/
  UNSAFE_PROXY_RESPONSE_HEADERS from router.py into litellm/constants.py,
  removing the router.py <-> proxy import path the two CodeQL
  cyclic-import alerts were pointing at
- move the eager fetch_stream() call before success_calls/logging/
  _track_deployment_metrics instead of incrementing then compensating
  with a manual decrement on failure
- fix a dead assert message: `mock_fallback.assert_not_called(), "..."`
  built a tuple, not an assert-with-message; assert_not_called() already
  raises on its own so this just drops the inert string

* revert(router): pull mid-stream continuation-removal out of this PR

Removing the continuation-prompt fallback (retrying with the partial
response as a prefixed assistant message) so a stream failing after
partial content always re-raises instead was a scope decision beyond
what this PR's title/issue (#31874) describe, and it directly conflicts
with #30242/#30743, which are already fixing the same code path for
Anthropic's removal of assistant-message prefill on Sonnet 4.6+/Opus
4.6+. Landing this PR's version first would delete the branch those PRs
are patching; landing theirs first would have this PR undo their fix on
rebase.

Restores the original prefill-based continuation-resume behavior
(including the is_pre_first_chunk guard already in litellm_internal_staging)
in both _acompletion_streaming_iterator and _completion_streaming_iterator,
and removes _stream_chunks_have_generated_content along with the tests
that only existed to cover the guard. This PR now only touches the
deferred-stream eager-fetch fix and the header-stripping fixes; the
non-text-content re-raise idea becomes a follow-up PR built on top of
whichever of #30242/#30743 lands.

* fix(proxy): re-filter unsafe headers after the response-headers hook merge

_handle_llm_api_exception filtered provider/framing headers once, then
merged in post_call_response_headers_hook's return value afterward
without re-filtering. The ProxyException branch happened to re-filter
after its own header merge, but the HTTPException/httpx.HTTPStatusError/
generic-exception branches passed the post-hook headers straight through
unfiltered, so a callback hook (any custom guardrail/logging plugin)
returning an unsafe header would bypass the strip entirely for those
paths. Filters once, right after the hook merge, so every branch gets
the same guarantee.

* Revert "revert(router): pull mid-stream continuation-removal out of this PR"

This reverts commit c5ca101f61.

* fix(router): detect reasoning_items as generated content in mid-stream guard

Greptile flagged that a structured reasoning-only delta (Delta.reasoning_items,
the OpenAI Responses-API-style reasoning item) wasn't recognized as
already-streamed content by _stream_chunks_have_generated_content, alongside
the existing thinking_blocks/tool_calls checks, so a stream that emitted only
reasoning_items before failing could still restart via fallback.

* fix(router): annotate _stream_chunks_have_generated_content with Sequence, not list

The type_discipline_gate LIT001 check flags mutable-collection parameter
annotations. chunks is only iterated, never mutated, so Sequence is the
correct read-only annotation and clears the ratcheted budget ceiling.

* fix(router): surface original provider exception, not the internal wrapper, when mid-stream fallback gives up

When content has already streamed and MidStreamFallbackError carries
original_exception (e.g. RateLimitError), both the async and sync
streaming iterators bare-re-raised the wrapper itself, so the client
lost the specific error type/code/provider_specific_fields instead of
seeing the real provider error. The fallback-failure path a few lines
below already unwraps to original_exception for the same reason; apply
the same pattern here.

Also extend _stream_chunks_have_generated_content to recognize audio,
images, and annotations deltas as generated content, matching
is_chunk_non_empty's existing annotations check and Delta's treatment
of audio/images as first-class content fields — a stream carrying only
one of these before failing was not recognized as already-streamed,
so the router could still restart it via fallback after the client had
received real content.

* chore: retrigger CI (frontend-lint cancelled, schema.d.ts flake)

frontend-lint's check-run shows conclusion=cancelled on 70e47f4897 with
no superseding run, and this PR touches no UI files. Verify schema.d.ts
matches the proxy OpenAPI spec is on the previously diagnosed
stream_timeout/user_role Union-ordering nondeterminism (e9fc5e5063).
Empty commit to force a fresh CI run for both rather than a manual
rerun, which requires repo admin rights this fork PR doesn't have.

---------

Co-authored-by: Deepanshu <deepanshu.lulla@alpha-sense.com>
This commit is contained in:
Deepanshu Lulla 2026-08-04 18:44:43 -04:00 committed by GitHub
parent 9ea5cfce0e
commit e2950a8995
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 884 additions and 136 deletions

View file

@ -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

View file

@ -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):

View file

@ -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:

View file

@ -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."""

View file

@ -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 = {}

View file

@ -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=[