From 5fbae0a1e4954a694e8817dfb241671e07d80eb9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 18 Jun 2026 13:45:32 +0000
Subject: [PATCH] fix(lint): use builtin generics in FallbackEntry and cover
 defensive branches

The FallbackEntry alias used typing.Dict/List, tripping the UP006 strict-rule
budget; switch to builtin generics. Drop the now-unreachable None guard in
_prefill_explicitly_unsupported (callers filter None before the lookup) and add
tests for the malformed-fallback and resolver-raises paths so the defensive
branches are exercised.
---
 .../router_utils/fallback_event_handlers.py   |  6 ++--
 .../test_fallback_event_handlers.py           | 32 +++++++++++++++++++
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py
index a3cbb7ef4fa..17a8aba2d0c 100644
--- a/litellm/router_utils/fallback_event_handlers.py
+++ b/litellm/router_utils/fallback_event_handlers.py
@@ -20,10 +20,10 @@ else:
 
 MID_STREAM_CONTINUATION_SYSTEM_PROMPT = "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: "
 
-FallbackEntry = Union[str, Dict[str, List[str]]]
+FallbackEntry = str | dict[str, list[str]]
 
 
-def _prefill_explicitly_unsupported(model: str | None) -> bool:
+def _prefill_explicitly_unsupported(model: str) -> bool:
     """True only when the model registry explicitly marks the model as NOT
     supporting assistant prefill (``supports_assistant_prefill: false``).
 
@@ -31,8 +31,6 @@ def _prefill_explicitly_unsupported(model: str | None) -> bool:
     the legacy prefill behavior — only models that would reject the prefill with
     a 400 anyway are routed to the user-message continuation.
     """
-    if model is None:
-        return False
     try:
         from litellm.utils import get_model_info
 
diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py
index 0933b4d410e..6a2b9b92363 100644
--- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py
+++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py
@@ -241,3 +241,35 @@ def test_resolver_returning_prefill_supporting_model_keeps_legacy():
         resolve_underlying_model=aliases.get,
     )
     _assert_legacy_prefill(result)
+
+
+def test_malformed_fallback_entry_does_not_crash():
+    """A malformed fallback entry (e.g. an empty dict) makes the underlying
+    fallback-group resolution raise; the capability check must swallow it and
+    fall back to the legacy prefill rather than break the mid-stream retry."""
+    result = build_mid_stream_continuation_messages(
+        messages=MESSAGES,
+        generated_content=PARTIAL,
+        model_group="gpt-4",
+        fallbacks=[{}],
+    )
+    _assert_legacy_prefill(result)
+
+
+def test_resolver_raising_is_swallowed():
+    """The injected resolver wraps router lookups that can raise; an exception
+    must not propagate out of the capability check, and the raw group name is
+    still evaluated."""
+
+    def boom(_model_group):
+        raise RuntimeError("deployment lookup failed")
+
+    result = build_mid_stream_continuation_messages(
+        messages=MESSAGES,
+        generated_content=PARTIAL,
+        model_group="claude-sonnet-4-6",
+        resolve_underlying_model=boom,
+    )
+    # raw name still classified as prefill-rejecting -> user continuation
+    assert len(result) == 2
+    assert result[1]["role"] == "user"