fix(router): retry an evicted router's budget writes and keep the adaptive header on user_config requests

A retired budget limiter now keeps retrying its queued Redis increments every
sync interval until Redis accepts them, instead of giving up after one failed
write.

The x-litellm-adaptive-router-model header now comes from one stateless hook in
litellm.callbacks that holds no router, so user_config adaptive routers, whose
per-router hooks are request-scoped, still answer with the header.

litellm/router.py

@@ -10611,6 +10611,7 @@ class Router:
         changes, otherwise a released router keeps recording turns through its hook."""
         from litellm.router_strategy.adaptive_router.hooks import (
             AdaptiveRouterPostCallHook,
+            AdaptiveRouterResponseHeaderHook,
         )
 
         for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook):
@@ -10620,6 +10621,8 @@ class Router:
                 litellm.logging_callback_manager.add_litellm_callback(
                     AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)
                 )
+        if self.adaptive_routers:
+            litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook())
 
     def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
         """

litellm/router_strategy/adaptive_router/__init__.py

@@ -1,6 +1,9 @@
 """Adaptive router strategy. See README.md for design overview."""
 
 from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
-from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
+from litellm.router_strategy.adaptive_router.hooks import (
+    AdaptiveRouterPostCallHook,
+    AdaptiveRouterResponseHeaderHook,
+)
 
-__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"]
+__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"]

litellm/router_strategy/adaptive_router/hooks.py

@@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen
     return content, tool_calls
 
 
-class AdaptiveRouterPostCallHook(CustomLogger):
-    """One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
-
-    def __init__(self, adaptive_router: AdaptiveRouter) -> None:
-        self.adaptive_router = adaptive_router
-
+class AdaptiveRouterResponseHeaderHook(CustomLogger):
     async def async_post_call_response_headers_hook(
         self,
         data: Mapping[str, object],
@@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger):
             return None
         return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen}
 
+
+class AdaptiveRouterPostCallHook(CustomLogger):
+    """One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
+
+    def __init__(self, adaptive_router: AdaptiveRouter) -> None:
+        self.adaptive_router = adaptive_router
+
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         await self._record(kwargs, response_obj, response_status=200)
 

litellm/router_strategy/budget_limiter.py

@@ -154,10 +154,9 @@ class RouterBudgetLimiting(CustomLogger):
         self._flush_after_retire = loop.create_task(self._drain_increments_after_retire())
 
     async def _drain_increments_after_retire(self) -> None:
-        if not self.redis_increment_operation_queue:
-            return
-        if await self._push_in_memory_increments_to_redis():
-            await self._drain_increments_after_retire()
+        while self.redis_increment_operation_queue:
+            if not await self._push_in_memory_increments_to_redis():
+                await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL)
 
     @with_service_target("router_budgets")
     async def async_filter_deployments(

tests/unit/proxy/test_route_llm_request.py

@@ -21,7 +21,9 @@ from unittest.mock import MagicMock
 from fakeredis import TcpFakeServer
 from fastapi import HTTPException
 
+from litellm.caching.caching import DualCache
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.route_llm_request import (
     JSON_OBJECT_ADAPTER,
     ProxyModelNotFoundError,
@@ -29,6 +31,11 @@ from litellm.proxy.route_llm_request import (
     route_request,
     user_config_routers,
 )
+from litellm.proxy.utils import ProxyLogging
+from litellm.router_strategy.adaptive_router.config import (
+    ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
+    ADAPTIVE_ROUTER_RESPONSE_HEADER,
+)
 from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError
 
 
@@ -2565,6 +2572,41 @@ async def _proxy_router_outcomes(
         router.discard()
 
 
+@pytest.mark.asyncio
+async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None:
+    data: Final = {
+        "model": "smart-cheap-router",
+        "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}],
+        "metadata": {},
+        "user_config": {
+            "model_list": [
+                {
+                    "model_name": "smart-cheap-router",
+                    "litellm_params": {
+                        "model": "auto_router/adaptive_router",
+                        "adaptive_router_config": {"available_models": ["fast", "smart"]},
+                    },
+                },
+                {
+                    **_mock_deployment("fast"),
+                    "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}},
+                },
+                {
+                    **_mock_deployment("smart"),
+                    "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}},
+                },
+            ]
+        },
+    }
+
+    await (await route_request(data, None, None, "acompletion"))
+    headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook(
+        data=data, user_api_key_dict=UserAPIKeyAuth(), response=None
+    )
+
+    assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]}
+
+
 def _global_router_callbacks() -> tuple[list[object], list[object]]:
     return list(litellm.callbacks), list(litellm.input_callback)
 

tests/unit/router_strategy/adaptive_router/test_hooks.py

@@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import (
 )
 from litellm.router_strategy.adaptive_router.hooks import (
     AdaptiveRouterPostCallHook,
+    AdaptiveRouterResponseHeaderHook,
     _recent_tool_results,
     _resolve_session_key,
 )
@@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header():
     """The header hook returns the `x-litellm-adaptive-router-model` header
     so proxy header construction picks it up (works for both streaming and
     non-streaming; `async_post_call_success_hook` is too late for streaming)."""
-    hook = _make_hook()
-    headers = await hook.async_post_call_response_headers_hook(
+    headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
         data={"metadata": {"adaptive_router_chosen_model": "smart"}},
         user_api_key_dict=MagicMock(),
         response=MagicMock(),
@@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header():
 
 @pytest.mark.asyncio
 async def test_post_call_response_headers_hook_noop_when_metadata_missing_key():
-    hook = _make_hook()
-    headers = await hook.async_post_call_response_headers_hook(
+    headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
         data={"metadata": {"litellm_session_id": "sess-A"}},
         user_api_key_dict=MagicMock(),
         response=MagicMock(),
@@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key():
 
 @pytest.mark.asyncio
 async def test_post_call_response_headers_hook_noop_when_no_metadata():
-    hook = _make_hook()
-    headers = await hook.async_post_call_response_headers_hook(
+    headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
         data={},
         user_api_key_dict=MagicMock(),
         response=MagicMock(),
@@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata():
 
 @pytest.mark.asyncio
 async def test_post_call_response_headers_hook_noop_when_metadata_not_dict():
-    hook = _make_hook()
-    headers = await hook.async_post_call_response_headers_hook(
+    headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook(
         data={"metadata": "not-a-dict"},
         user_api_key_dict=MagicMock(),
         response=MagicMock(),

tests/unit/router_strategy/test_budget_limiter_hotpath.py

@@ -786,3 +786,28 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations:
     await limiter._push_in_memory_increments_to_redis()
     assert redis_cache.values[_SPEND_KEY] == 10.0
     assert limiter.redis_increment_operation_queue == []
+
+
+@pytest.mark.asyncio
+async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None:
+    pipeline_started: Final = asyncio.Event()
+    pipeline_completed: Final = asyncio.Event()
+    redis_cache: Final = _MockRedisCache(
+        initial_values={},
+        pipeline_started=pipeline_started,
+        pipeline_completed=pipeline_completed,
+        should_fail_pipeline=True,
+    )
+    budget_limiter: Final = _new_router_budget_limiter(
+        redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)]
+    )
+    budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600))
+
+    budget_limiter.retire()
+    await asyncio.wait_for(pipeline_started.wait(), timeout=1)
+    assert redis_cache.values == {}
+    redis_cache.should_fail_pipeline = False
+    await asyncio.wait_for(pipeline_completed.wait(), timeout=5)
+
+    assert redis_cache.values == {_SPEND_KEY: 10.0}
+    assert budget_limiter.redis_increment_operation_queue == []