diff --git a/litellm/router.py b/litellm/router.py index 845c388ecd6..1ef6327b975 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9710,12 +9710,63 @@ 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: + """ + 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. + """ + from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( + EncryptedContentAffinityCheck, + ) + + if self.optional_callbacks is None: + self.optional_callbacks = [] + + if any( + check in removed_checks + for check in ("deployment_affinity", "responses_api_deployment_check", "session_affinity") + ): + for callback in self.optional_callbacks: + if not isinstance(callback, DeploymentAffinityCheck): + continue + if "deployment_affinity" in removed_checks: + callback.enable_user_key_affinity = False + if "responses_api_deployment_check" in removed_checks: + callback.enable_responses_api_affinity = False + if "session_affinity" in removed_checks: + callback.enable_session_id_affinity = False + break + + if "encrypted_content_affinity" in removed_checks: + for callback in self.optional_callbacks: + if isinstance(callback, EncryptedContentAffinityCheck): + callback.enable_global_affinity = False + break + + removable_callback_types = { + "prompt_caching": PromptCachingDeploymentCheck, + "enforce_model_rate_limits": ModelRateLimitingCheck, + "router_budget_limiting": RouterBudgetLimiting, + } + for check, callback_type in removable_callback_types.items(): + if check not in removed_checks: + continue + litellm.logging_callback_manager.remove_callbacks_by_type(self.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 + def _apply_optional_pre_call_checks_setting(self, value: OptionalPreCallChecks | None) -> None: if value is None: return new_checks = [check for check in value if check not in self.optional_pre_call_checks] + 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) # 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/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index b23b2167150..9ee210c0bb2 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -82,6 +82,7 @@ ignored_function_names = [ "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "_merge_default_litellm_params_setting", # Tested indirectly via update_settings(default_litellm_params=...) "_apply_optional_pre_call_checks_setting", # Tested indirectly via update_settings(optional_pre_call_checks=...) + "_remove_optional_pre_call_checks", # Tested indirectly via update_settings(optional_pre_call_checks=...) ] diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 132d6597149..d14a1fe0886 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5386,6 +5386,77 @@ def test_update_settings_optional_pre_call_checks_is_idempotent(): assert len(prompt_caching_callbacks) == 1 +@pytest.mark.asyncio +async def test_update_settings_optional_pre_call_checks_removes_unregistered_checks(): + """ + Regression test: `update_settings(optional_pre_call_checks=...)` only ever + unioned new checks into `self.optional_pre_call_checks` - it never removed + anything absent from the incoming list. Unchecking "prompt_caching" in the + Admin UI's new multi-select and clicking Save would silently do nothing: + the DB got the smaller list, but the live router kept both + `optional_pre_call_checks` and the registered `PromptCachingDeploymentCheck` + callback unchanged, showing "saved" while the router kept enforcing the + removed check until a restart. + """ + from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( + PromptCachingDeploymentCheck, + ) + + router = _make_router_for_settings_tests() + router.update_settings( + optional_pre_call_checks=["prompt_caching", "router_budget_limiting"] + ) + assert router.optional_pre_call_checks == [ + "prompt_caching", + "router_budget_limiting", + ] + assert router.router_budget_logger is not None + + router.update_settings(optional_pre_call_checks=["prompt_caching"]) + + assert router.optional_pre_call_checks == ["prompt_caching"] + assert router.router_budget_logger is None + callbacks = router.optional_callbacks or [] + assert any(isinstance(cb, PromptCachingDeploymentCheck) for cb in callbacks) + assert not any(cb.__class__.__name__ == "RouterBudgetLimiting" for cb in callbacks) + assert not any( + cb.__class__.__name__ == "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 + share one `DeploymentAffinityCheck` callback instance keyed by boolean flags, not a + dedicated callback per check. Removing one of these three from + optional_pre_call_checks must clear only its flag, leaving the shared callback (and + any other still-enabled flag) intact rather than removing the whole callback. + """ + from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( + DeploymentAffinityCheck, + ) + + router = _make_router_for_settings_tests() + router.update_settings( + optional_pre_call_checks=["deployment_affinity", "session_affinity"] + ) + affinity_callback = next( + cb + for cb in (router.optional_callbacks or []) + if isinstance(cb, DeploymentAffinityCheck) + ) + assert affinity_callback.enable_user_key_affinity is True + assert affinity_callback.enable_session_id_affinity is True + + router.update_settings(optional_pre_call_checks=["session_affinity"]) + + assert router.optional_pre_call_checks == ["session_affinity"] + assert affinity_callback.enable_user_key_affinity is False + assert affinity_callback.enable_session_id_affinity is True + # The shared callback stays registered - session_affinity is still enabled. + assert affinity_callback in (router.optional_callbacks or []) + + def test_get_settings_includes_default_litellm_params_and_optional_pre_call_checks(): """ Regression test: the Admin UI's Router Settings page reads its current