fix(router): enforce deployment budgets for dynamically added models (#29273)

* fix(router): enforce deployment budgets for dynamically added models

Register deployment max_budget/budget_duration when models are added via
upsert_deployment (e.g. /model/new) so RouterBudgetLimiting matches startup
model_list behavior.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(router): address CI lint and router coverage for budget sync helpers

Remove unused RouterBudgetLimiting import and add router unit tests for
deployment budget helper methods required by router_code_coverage.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(router): clear stale deployment budget on upsert without limits

Unregister deployment budget config when max_budget/budget_duration are
removed, including upsert replace paths. Hoist provider budget logger
lookup outside the provider loop.

Co-authored-by: Cursor <cursoragent@cursor.com>

* Fix mypy

* fix black

---------

Co-authored-by: Cursor <cursoragent@cursor.com>

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -28,6 +28,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
     def __init__(self, dual_cache: DualCache):
         self.dual_cache = dual_cache
         self.redis_increment_operation_queue = []
+        self.deployment_budget_config = None
 
     async def is_key_within_model_budget(
         self,

litellm/proxy/spend_tracking/spend_management_endpoints.py

@@ -21,7 +21,6 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
     get_spend_by_team_and_customer,
 )
 from litellm.proxy.utils import handle_exception_on_proxy
-from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
 
 if TYPE_CHECKING:
     from litellm.proxy.proxy_server import PrismaClient
@@ -3149,18 +3148,12 @@ async def provider_budgets() -> ProviderBudgetResponse:
                 "No provider budget config found. Please set a provider budget config in the router settings. https://docs.litellm.ai/docs/proxy/provider_budget_routing"
             )
 
+        router_budget_logger = llm_router._get_router_deployment_budget_limiter()
+        if router_budget_logger is None:
+            raise ValueError("No router budget logger found")
+
         provider_budget_response_dict: Dict[str, ProviderBudgetResponseObject] = {}
         for _provider, _budget_info in provider_budget_config.items():
-            router_budget_logger = next(
-                (
-                    cb
-                    for cb in (llm_router.optional_callbacks or [])
-                    if isinstance(cb, RouterBudgetLimiting)
-                ),
-                None,
-            )
-            if router_budget_logger is None:
-                raise ValueError("No router budget logger found")
             _provider_spend = (
                 await router_budget_logger._get_current_provider_spend(_provider) or 0.0
             )

litellm/router.py

@@ -1753,11 +1753,14 @@ class Router:
             if pre_call_check == "prompt_caching":
                 _callback = PromptCachingDeploymentCheck(cache=self.cache)
             elif pre_call_check == "router_budget_limiting":
+                if self._get_router_deployment_budget_limiter() is not None:
+                    continue
                 _callback = RouterBudgetLimiting(
                     dual_cache=self.cache,
                     provider_budget_config=self.provider_budget_config,
                     model_list=self.model_list,
                 )
+                self.router_budget_logger = _callback
             elif pre_call_check == "enforce_model_rate_limits":
                 _callback = ModelRateLimitingCheck(dual_cache=self.cache)
 
@@ -8529,6 +8532,7 @@ class Router:
             model=_deployment, model_id=deployment.model_info.id
         )
         self.model_names.add(deployment.model_name)
+        self._sync_deployment_budget_config(deployment=deployment)
         return deployment
 
     def _update_deployment_indices_after_removal(
@@ -8717,12 +8721,64 @@ class Router:
                 self._update_deployment_indices_after_removal(
                     model_id=id, removal_idx=deployment_idx
                 )
+                _budget_limiter = self._get_router_deployment_budget_limiter()
+                if _budget_limiter is not None:
+                    _budget_limiter.unregister_deployment_budget(model_id=id)
                 return item
             else:
                 return None
         except Exception:
             return None
 
+    def _get_router_deployment_budget_limiter(
+        self,
+    ) -> Optional[RouterBudgetLimiting]:
+        """
+        Return the router's deployment-budget callback.
+
+        Uses exact-type matching so proxy subclasses (e.g. virtual-key model budgets)
+        registered on litellm.callbacks are not mistaken for router deployment budgets.
+        """
+        if self.router_budget_logger is not None:
+            return self.router_budget_logger
+
+        if self.optional_callbacks:
+            for _cb in self.optional_callbacks:
+                if type(_cb) is RouterBudgetLimiting:
+                    self.router_budget_logger = _cb
+                    return _cb
+        return None
+
+    def _deployment_has_budget_limits(self, deployment: Deployment) -> bool:
+        return (
+            deployment.litellm_params.get("max_budget") is not None
+            and deployment.litellm_params.get("budget_duration") is not None
+            and deployment.model_info.id is not None
+        )
+
+    def _sync_deployment_budget_config(self, deployment: Deployment) -> None:
+        model_id = deployment.model_info.id
+        if model_id is None:
+            return
+
+        _budget_limiter = self._get_router_deployment_budget_limiter()
+
+        if not self._deployment_has_budget_limits(deployment=deployment):
+            if _budget_limiter is not None:
+                _budget_limiter.unregister_deployment_budget(model_id=model_id)
+            return
+
+        if _budget_limiter is None:
+            self.add_optional_pre_call_checks(
+                optional_pre_call_checks=["router_budget_limiting"]
+            )
+            _budget_limiter = self._get_router_deployment_budget_limiter()
+
+        if _budget_limiter is not None:
+            _budget_limiter.register_deployment_budget(
+                deployment=deployment.to_json(exclude_none=True)
+            )
+
     def get_deployment(self, model_id: str) -> Optional[Deployment]:
         """
         Returns -> Deployment or None

litellm/router_strategy/budget_limiter.py

@@ -96,9 +96,7 @@ class RouterBudgetLimiting(CustomLogger):
         self,
         dual_cache: DualCache,
         provider_budget_config: Optional[dict],
-        model_list: Optional[
-            Union[List[DeploymentTypedDict], List[Dict[str, Any]]]
-        ] = None,
+        model_list: Optional[List[Union[DeploymentTypedDict, Dict[str, Any]]]] = None,
     ):
         self.dual_cache = dual_cache
         self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = []
@@ -854,9 +852,7 @@ class RouterBudgetLimiting(CustomLogger):
 
     def _init_deployment_budgets(
         self,
-        model_list: Optional[
-            Union[List[DeploymentTypedDict], List[Dict[str, Any]]]
-        ] = None,
+        model_list: Optional[List[Union[DeploymentTypedDict, Dict[str, Any]]]] = None,
     ):
         if model_list is None:
             return
@@ -887,6 +883,22 @@ class RouterBudgetLimiting(CustomLogger):
             f"Initialized Deployment Budget Config: {self.deployment_budget_config}"
         )
 
+    def register_deployment_budget(
+        self,
+        deployment: Union[Dict[str, Any], DeploymentTypedDict],
+    ) -> None:
+        """
+        Register or refresh deployment-level budget config for a runtime-added deployment.
+        """
+        self._init_deployment_budgets(model_list=[deployment])
+
+    def unregister_deployment_budget(self, model_id: str) -> None:
+        if self.deployment_budget_config is None:
+            return
+        self.deployment_budget_config.pop(model_id, None)
+        if len(self.deployment_budget_config) == 0:
+            self.deployment_budget_config = None
+
     def _init_tag_budgets(self):
         if litellm.tag_budget_config is None:
             return

tests/router_unit_tests/test_router_helper_utils.py

@@ -14,7 +14,7 @@ import litellm
 from unittest.mock import patch, MagicMock, AsyncMock
 from create_mock_standard_logging_payload import create_standard_logging_payload
 from litellm.types.utils import StandardLoggingPayload
-from litellm.types.router import Deployment, LiteLLM_Params
+from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
 
 
 @pytest.fixture
@@ -997,9 +997,7 @@ def test_filter_cooldown_deployments(model_list):
         healthy_deployments=router._get_all_deployments(model_name="gpt-5-mini"),  # type: ignore
         cooldown_deployments=[],
     )
-    assert len(deployments) == len(
-        router._get_all_deployments(model_name="gpt-5-mini")
-    )
+    assert len(deployments) == len(router._get_all_deployments(model_name="gpt-5-mini"))
 
 
 def test_track_deployment_metrics(model_list):
@@ -2379,3 +2377,123 @@ def test_get_router_model_info_with_deployment_object():
     # Verify we got valid model info back
     assert model_info is not None
     assert isinstance(model_info, dict)
+
+
+def test_deployment_has_budget_limits():
+    router = Router(model_list=[])
+
+    with_budget = Deployment(
+        model_name="budgeted-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            max_budget=0.001,
+            budget_duration="1d",
+        ),
+        model_info=ModelInfo(id="budget-deployment-id"),
+    )
+    without_budget = Deployment(
+        model_name="unbudgeted-model",
+        litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
+        model_info=ModelInfo(id="no-budget-deployment-id"),
+    )
+
+    assert router._deployment_has_budget_limits(deployment=with_budget) is True
+    assert router._deployment_has_budget_limits(deployment=without_budget) is False
+
+
+def test_sync_deployment_budget_config(monkeypatch):
+    import asyncio
+
+    monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
+
+    router = Router(model_list=[], optional_pre_call_checks=[])
+    deployment = Deployment(
+        model_name="dynamic-budget-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            api_key="fake-key",
+            max_budget=0.000000000001,
+            budget_duration="1d",
+        ),
+        model_info=ModelInfo(id="runtime-budget-deployment"),
+    )
+
+    router._sync_deployment_budget_config(deployment=deployment)
+
+    budget_limiter = router._get_router_deployment_budget_limiter()
+    assert budget_limiter is not None
+    config = budget_limiter._get_budget_config_for_deployment(
+        "runtime-budget-deployment"
+    )
+    assert config is not None
+    assert config.max_budget == 0.000000000001
+
+
+def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch):
+    import asyncio
+
+    monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
+
+    router = Router(model_list=[], optional_pre_call_checks=[])
+    model_id = "runtime-budget-deployment"
+    budgeted = Deployment(
+        model_name="dynamic-budget-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            api_key="fake-key",
+            max_budget=0.000000000001,
+            budget_duration="1d",
+        ),
+        model_info=ModelInfo(id=model_id),
+    )
+    unbudgeted = Deployment(
+        model_name="dynamic-budget-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            api_key="fake-key",
+        ),
+        model_info=ModelInfo(id=model_id),
+    )
+
+    router._sync_deployment_budget_config(deployment=budgeted)
+    budget_limiter = router._get_router_deployment_budget_limiter()
+    assert budget_limiter is not None
+    assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
+
+    router._sync_deployment_budget_config(deployment=unbudgeted)
+    assert budget_limiter._get_budget_config_for_deployment(model_id) is None
+
+
+def test_upsert_deployment_clears_stale_budget_config(monkeypatch):
+    import asyncio
+
+    monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
+
+    router = Router(model_list=[], optional_pre_call_checks=[])
+    model_id = "upsert-budget-deployment"
+    budgeted = Deployment(
+        model_name="dynamic-budget-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            api_key="fake-key",
+            max_budget=0.000000000001,
+            budget_duration="1d",
+        ),
+        model_info=ModelInfo(id=model_id),
+    )
+    unbudgeted = Deployment(
+        model_name="dynamic-budget-model",
+        litellm_params=LiteLLM_Params(
+            model="openai/gpt-4o-mini",
+            api_key="fake-key",
+        ),
+        model_info=ModelInfo(id=model_id),
+    )
+
+    router.upsert_deployment(deployment=budgeted)
+    budget_limiter = router._get_router_deployment_budget_limiter()
+    assert budget_limiter is not None
+    assert budget_limiter._get_budget_config_for_deployment(model_id) is not None
+
+    router.upsert_deployment(deployment=unbudgeted)
+    assert budget_limiter._get_budget_config_for_deployment(model_id) is None

tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py

@@ -234,3 +234,72 @@ async def test_get_llm_provider_for_deployment_matches_legacy_behavior(
     legacy_provider = _legacy_provider_resolution(deployment)
 
     assert current_provider == legacy_provider
+
+
+def test_register_deployment_budget_for_runtime_added_deployment(
+    disable_budget_sync, monkeypatch
+):
+    import asyncio
+
+    monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
+    budget_limiter = RouterBudgetLimiting(
+        dual_cache=DualCache(),
+        provider_budget_config={},
+    )
+    model_id = "dynamic-deployment-id"
+    budget_limiter.register_deployment_budget(
+        deployment={
+            "model_name": "dynamic-budget-model",
+            "litellm_params": {
+                "model": "openai/gpt-4o-mini",
+                "max_budget": 0.000000000001,
+                "budget_duration": "1d",
+            },
+            "model_info": {"id": model_id},
+        }
+    )
+
+    config = budget_limiter._get_budget_config_for_deployment(model_id)
+    assert config is not None
+    assert config.max_budget == 0.000000000001
+    assert config.budget_duration == "1d"
+
+    budget_limiter.unregister_deployment_budget(model_id=model_id)
+    assert budget_limiter._get_budget_config_for_deployment(model_id) is None
+
+
+def test_router_add_deployment_registers_deployment_budget(
+    disable_budget_sync, monkeypatch
+):
+    import asyncio
+
+    from litellm import Router
+    from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
+
+    monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
+
+    router = Router(
+        model_list=[],
+        optional_pre_call_checks=[],
+    )
+
+    router.add_deployment(
+        deployment=Deployment(
+            model_name="dynamic-budget-model",
+            litellm_params=LiteLLM_Params(
+                model="openai/gpt-4o-mini",
+                api_key="fake-key",
+                max_budget=0.000000000001,
+                budget_duration="1d",
+            ),
+            model_info=ModelInfo(id="runtime-budget-deployment"),
+        )
+    )
+
+    budget_limiter = router._get_router_deployment_budget_limiter()
+    assert budget_limiter is not None
+    config = budget_limiter._get_budget_config_for_deployment(
+        "runtime-budget-deployment"
+    )
+    assert config is not None
+    assert config.max_budget == 0.000000000001