fix(proxy): walk fallback groups breadth first and release reservations on a rejected WebSocket

Each fallback group is looked up once however often the request repeats it, and a WebSocket closed by the endpoint gate runs the failure hook so rate-limit reservations are released

litellm/proxy/response_api_endpoints/endpoints.py

@@ -1571,6 +1571,11 @@ async def _responses_websocket_session(
             team_id=user_api_key_dict.team_id,
         )
         if unsupported_endpoint is not None:
+            await proxy_logging_obj.post_call_failure_hook(
+                user_api_key_dict=user_api_key_dict,
+                original_exception=unsupported_endpoint,
+                request_data=data,
+            )
             raise unsupported_endpoint
     except Exception as e:
         verbose_proxy_logger.exception("Responses WebSocket pre-call error")

litellm/proxy/route_llm_request.py

@@ -1,6 +1,7 @@
 import asyncio
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
+from itertools import chain
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol
 
@@ -491,19 +492,35 @@ def _fallback_target_group(target: object) -> str | None:
     return named_model if isinstance(named_model, str) else None
 
 
-def _fallback_target_answers(
-    group: str | None,
+def _next_fallback_groups(
+    model: str,
     data: Mapping[str, object],
+    fallbacks: list[object],  # mutable-ok: mirrors get_fallback_model_group's contract
+) -> tuple[str, ...]:
+    return tuple(
+        group
+        for target in _fallback_targets(model, data, fallbacks)
+        if has_unattempted_fallback_target((target,), data) and (group := _fallback_target_group(target)) is not None
+    )
+
+
+def _frontier_reaches_chat(
+    frontier: tuple[str, ...],
+    seen: frozenset[str],
+    data: Mapping[str, object],
+    fallbacks: list[object],  # mutable-ok: mirrors get_fallback_model_group's contract
     llm_router: LitellmRouter,
     team_id: str | None,
-    walked: frozenset[str],
 ) -> bool:
-    if group is None or group in walked:
+    reached: Final = dict.fromkeys(chain.from_iterable(_next_fallback_groups(m, data, fallbacks) for m in frontier))
+    fresh: Final = tuple(group for group in reached if group not in seen)
+    if not fresh:
         return False
-    deployments: Final = _group_deployments(group, llm_router, team_id)
-    if not deployments:
-        return False
-    return _any_serves_chat(deployments) or _fallback_would_answer(group, data, llm_router, team_id, walked)
+    deployments_by_group: Final = tuple((group, _group_deployments(group, llm_router, team_id)) for group in fresh)
+    if any(deployments and _any_serves_chat(deployments) for _, deployments in deployments_by_group):
+        return True
+    expandable: Final = tuple(group for group, deployments in deployments_by_group if deployments)
+    return _frontier_reaches_chat(expandable, seen | frozenset(fresh), data, fallbacks, llm_router, team_id)
 
 
 def _fallback_would_answer(
@@ -511,7 +528,6 @@ def _fallback_would_answer(
     data: Mapping[str, object],
     llm_router: LitellmRouter,
     team_id: str | None,
-    walked: frozenset[str] = frozenset(),
 ) -> bool:
     if data.get("disable_fallbacks") is True:
         return False
@@ -522,12 +538,7 @@ def _fallback_would_answer(
         fallbacks: Final = _FALLBACK_LIST_ADAPTER.validate_python(raw_fallbacks)
     except ValidationError:
         return True
-    walked_here: Final = walked | {model}
-    return any(
-        has_unattempted_fallback_target((target,), data)
-        and _fallback_target_answers(_fallback_target_group(target), data, llm_router, team_id, walked_here)
-        for target in _fallback_targets(model, data, fallbacks)
-    )
+    return _frontier_reaches_chat((model,), frozenset({model}), data, fallbacks, llm_router, team_id)
 
 
 def endpoint_unsupported_by_model(

litellm/router.py

@@ -1076,7 +1076,7 @@ class Router:
         self.assistants_config = assistants_config
         self.search_tools = search_tools or []
         self.guardrail_list = guardrail_list or []
-        self.deployment_names: list = []  # names of models under litellm_params. ex. azure/chatgpt-v-2
+        self.deployment_names: list[str] = []  # names of models under litellm_params. ex. azure/chatgpt-v-2
         self.deployment_latency_map = {}
         ### CACHING ###
         cache_type: Literal["local", "redis", "redis-semantic", "s3", "disk"] = "local"  # default to an in-memory cache

tests/unit/proxy/response_api_endpoints/test_endpoints.py

@@ -15,6 +15,7 @@ from httpx import Response
 
 import litellm
 from litellm.proxy.proxy_server import app
+from litellm.proxy.route_llm_request import ProxyUnsupportedEndpointError
 from litellm.types.llms.openai import ResponsesAPIResponse
 from tests._master_key import MASTER_KEY
 
@@ -1069,6 +1070,10 @@ class TestResponsesWSDeclaredNonChatModel:
                 return_value=fake_llm_call(),
             ) as mock_route_request,
             patch("litellm.proxy.proxy_server.llm_router", router),  # test-quality-ok: the session reads the proxy's router global
+            patch(  # test-quality-ok: the session reads the proxy's logging global; the failure hook releases reservations
+                "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook",
+                new_callable=AsyncMock,
+            ) as mock_failure_hook,
         ):
             await responses_websocket_endpoint(
                 websocket=ws,
@@ -1079,7 +1084,11 @@ class TestResponsesWSDeclaredNonChatModel:
         if not rejected:
             mock_route_request.assert_awaited_once()
             ws.close.assert_not_awaited()
+            mock_failure_hook.assert_not_awaited()
             return
+        mock_failure_hook.assert_awaited_once()
+        assert isinstance(mock_failure_hook.await_args.kwargs["original_exception"], ProxyUnsupportedEndpointError)
+        assert mock_failure_hook.await_args.kwargs["request_data"]["model"] == "embed-only"
         frame = json.loads(ws.send_text.await_args.args[0])
         assert frame["type"] == "error"
         assert frame["error"]["type"] == "invalid_request_error"

tests/unit/proxy/test_route_llm_request.py

@@ -2060,6 +2060,29 @@ def test_endpoint_unsupported_ignores_endpoints_copied_from_the_loaded_catalog(d
     assert isinstance(_unsupported("acompletion", "catalog-row", router), ProxyUnsupportedEndpointError) is rejected
 
 
+class _LookupCountingRouter(Router):
+    def __init__(self, **router_kwargs: object) -> None:
+        self.looked_up: list[str | None] = []
+        super().__init__(**router_kwargs)
+        self.looked_up.clear()
+
+    def get_model_list(self, model_name: str | None = None, team_id: str | None = None):
+        self.looked_up.append(model_name)
+        return super().get_model_list(model_name=model_name, team_id=team_id)
+
+
+def test_endpoint_unsupported_looks_up_each_repeated_fallback_group_once():
+    fallback_groups: Final = ("embed-a", "embed-b", "embed-c")
+    router: Final = _LookupCountingRouter(
+        model_list=[_embedding_deployment(name, EMBEDDING_ONLY) for name in ("embed-only", *fallback_groups)]
+    )
+
+    error: Final = _unsupported("acompletion", "embed-only", router, fallbacks=list(fallback_groups) * 5)
+
+    assert isinstance(error, ProxyUnsupportedEndpointError)
+    assert sorted(name for name in router.looked_up if name in fallback_groups) == sorted(fallback_groups)
+
+
 def _router_with_defaults(default_litellm_params: dict[str, object]):
     import litellm