From e1e58f3f9818a054856cb159ab67bf488266b79e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 13 Jul 2026 19:46:15 -0700 Subject: [PATCH] fix(router): keep router_budget_limiting enforced when budgets are configured Router.__init__ auto-enables router_budget_limiting whenever a deployment has max_budget/budget_duration set or provider_budget_config is configured (RouterBudgetLimiting.should_init_router_budget_limiter), independent of what optional_pre_call_checks explicitly lists. _remove_optional_pre_call_checks didn't account for that: a save (via the Admin UI's new multi-select, or a config-sync payload) that simply omitted "router_budget_limiting" from the list would unregister the RouterBudgetLimiting callback and null out router_budget_logger, silently letting deployments keep serving requests after their configured budget is exhausted. _remove_optional_pre_call_checks now checks should_init_router_budget_limiter before actually removing the callback for this one check, and returns the checks it kept active despite being in removed_checks so _apply_optional_pre_call_checks_setting can fold them back into the tracked optional_pre_call_checks list - keeping the UI's displayed state honest about what's still enforced. --- litellm/router.py | 24 +++++++++++++++--- tests/test_litellm/test_router.py | 41 +++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 6ebb5962d7b..58427a1797f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9710,11 +9710,18 @@ class Router: if value is not None: self.default_litellm_params = {**self.default_litellm_params, **value} - def _remove_optional_pre_call_checks(self, removed_checks: OptionalPreCallChecks) -> None: + def _remove_optional_pre_call_checks(self, removed_checks: OptionalPreCallChecks) -> list[str]: """ Reverse of `add_optional_pre_call_checks` for the subset of checks that can be safely turned off at runtime: clears the corresponding flag(s) on shared affinity callbacks, or unregisters the dedicated callback instance entirely. + + Returns the subset of `removed_checks` that were kept active anyway because + they're still required by config (currently only `router_budget_limiting`, + which `Router.__init__` auto-enables whenever budgets are configured on the + deployments/provider, independent of `optional_pre_call_checks`) - callers + should fold these back into the tracked `optional_pre_call_checks` list so it + doesn't claim a check is off when it's actually still enforced. """ from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( EncryptedContentAffinityCheck, @@ -9748,13 +9755,23 @@ class Router: "enforce_model_rate_limits": ModelRateLimitingCheck, "router_budget_limiting": RouterBudgetLimiting, } + retained_checks: list[str] = [] for check, callback_type in removable_callback_types.items(): if check not in removed_checks: continue + if check == "router_budget_limiting" and RouterBudgetLimiting.should_init_router_budget_limiter( + model_list=self.model_list, provider_budget_config=self.provider_budget_config + ): + # Budgets are still configured on the deployments/provider - Router.__init__ + # would auto-enable this regardless of optional_pre_call_checks, so a save + # that omits it must not silently disable budget enforcement. + retained_checks.append(check) + continue litellm.logging_callback_manager.remove_callbacks_by_type(optional_callbacks, callback_type) litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, callback_type) if check == "router_budget_limiting": self.router_budget_logger = None + return retained_checks def _apply_optional_pre_call_checks_setting(self, value: OptionalPreCallChecks | None) -> None: if value is None: @@ -9763,9 +9780,8 @@ class Router: removed_checks = [check for check in self.optional_pre_call_checks if check not in value] if new_checks: self.add_optional_pre_call_checks(new_checks) - if removed_checks: - self._remove_optional_pre_call_checks(removed_checks) - self.optional_pre_call_checks = list(value) + retained_checks = self._remove_optional_pre_call_checks(removed_checks) if removed_checks else [] + self.optional_pre_call_checks = list(dict.fromkeys([*value, *retained_checks])) # Settings whose update logic doesn't fit `setattr(self, var, value)` (e.g. # merge-not-replace, or side effects beyond storing the value). Dispatched via diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 726a1218e03..fcd2725c185 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5167,6 +5167,47 @@ async def test_update_settings_optional_pre_call_checks_removes_unregistered_che assert not any(cb.__class__.__name__ == "RouterBudgetLimiting" for cb in litellm.callbacks) +@pytest.mark.asyncio +async def test_update_settings_cannot_remove_router_budget_limiting_while_budgets_configured(): + """ + Regression test: `Router.__init__` auto-enables `router_budget_limiting` whenever + a deployment has `max_budget`/`budget_duration` set or `provider_budget_config` is + configured, independent of what's in `optional_pre_call_checks` + (`RouterBudgetLimiting.should_init_router_budget_limiter`). If update_settings + honored a save that omits "router_budget_limiting" from the list while those + budgets are still configured, deployments would keep serving requests after their + budget is exhausted - a silent budget-enforcement bypass reachable via the Admin UI + or a config-sync payload that simply doesn't include the check. + """ + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "azure/gpt-4.1-mini", + "api_key": "fake-key", + "api_base": "https://fake.openai.azure.com", + "max_budget": 100, + "budget_duration": "1d", + }, + } + ], + ) + assert router.router_budget_logger is not None + assert "router_budget_limiting" in router.optional_pre_call_checks + + # A save that omits the (auto-enabled, config-required) check must not disable it. + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + + assert "router_budget_limiting" in router.optional_pre_call_checks + assert "prompt_caching" in router.optional_pre_call_checks + assert router.router_budget_logger is not None + assert any(isinstance(cb, RouterBudgetLimiting) for cb in (router.optional_callbacks or [])) + assert any(isinstance(cb, RouterBudgetLimiting) for cb in litellm.callbacks) + + def test_update_settings_optional_pre_call_checks_removes_affinity_flags(): """ Regression test: deployment_affinity/session_affinity/responses_api_deployment_check