From 1d3e26fd98b40d40f498e14b2470e8acc79fb9f6 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Wed, 2 Sep 2026 03:28:04 +0000
Subject: [PATCH] fix: preserve shared optional callbacks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/router.py                             | 19 +++++----
 .../test_router_retry_policy_update.py        | 41 ++++++++++++++++++-
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index e2865542e89..e7588d8ad5a 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -2099,16 +2099,19 @@ class Router:
         self.add_optional_pre_call_checks(optional_pre_call_checks)
 
     def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None:
-        if self.optional_callbacks is None:
+        if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks):
             return
-        removed: Final = [cb for cb in self.optional_callbacks if isinstance(cb, callback_cls)]
-        if not removed:
+        self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
+        if any(
+            router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or []))
+            for router in tuple(_live_routers)
+        ):
             return
-        self.optional_callbacks = [cb for cb in self.optional_callbacks if not isinstance(cb, callback_cls)]
-        for cb in removed:
-            litellm.logging_callback_manager.remove_callback_from_list_by_object(
-                litellm.callbacks, cb, require_self=False
-            )
+        for cb in tuple(litellm.callbacks):
+            if type(cb) is callback_cls:
+                litellm.logging_callback_manager.remove_callback_from_list_by_object(
+                    litellm.callbacks, cb, require_self=False
+                )
 
     def print_deployment(self, deployment: dict):
         """
diff --git a/tests/test_litellm/test_router_retry_policy_update.py b/tests/test_litellm/test_router_retry_policy_update.py
index 26251290176..be568134763 100644
--- a/tests/test_litellm/test_router_retry_policy_update.py
+++ b/tests/test_litellm/test_router_retry_policy_update.py
@@ -151,8 +151,45 @@ def test_remove_optional_pre_call_check_removes_local_and_global_callbacks():
     router.set_optional_pre_call_checks(["prompt_caching"])
     router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck)
 
-    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
-    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or []))
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
+
+
+def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router():
+    router_a = _build_router()
+    router_b = _build_router()
+
+    router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
+    router_b.update_settings(optional_pre_call_checks=["prompt_caching"])
+
+    router_a.update_settings(optional_pre_call_checks=[])
+
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
+    assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
+    assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
+
+    router_b.update_settings(optional_pre_call_checks=[])
+
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
+
+
+def test_remove_optional_pre_call_check_keeps_global_callback_when_second_router_clears_first():
+    router_a = _build_router()
+    router_b = _build_router()
+
+    router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
+    router_b.update_settings(optional_pre_call_checks=["prompt_caching"])
+
+    router_b.update_settings(optional_pre_call_checks=[])
+
+    assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
+    assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
+
+    router_a.update_settings(optional_pre_call_checks=[])
+
+    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)
 
 
 def test_update_settings_replaces_toggleable_pre_call_checks():