fix(proxy): clear MCP OpenAPI mappings on server eviction; widen budget cache invalidation

Evict OpenAPI tools from global_mcp_tool_registry and strip tool_name_to_mcp_server_name_mapping entries when a server leaves the runtime registry (remove_server and approval-status eviction). Invalidate user_api_key_cache for keys, orgs, and team members on budget-tier spend resets alongside tags.

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

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -599,16 +599,57 @@ class MCPServerManager:
             )
             raise e
 
+    def _cleanup_server_tool_routing_artifacts(self, server: MCPServer) -> None:
+        """Drop OpenAPI global tools and name-mapping rows owned by ``server``.
+
+        When a server leaves ``self.registry`` (eviction, ``remove_server``, etc.),
+        OpenAPI tools remain in ``global_mcp_tool_registry`` and
+        ``tool_name_to_mcp_server_name_mapping`` unless removed here. Stale
+        mappings make ``_get_mcp_server_from_tool_name`` resolve to a prefix that
+        no longer exists in the live registry.
+        """
+        from litellm.proxy._experimental.mcp_server.tool_registry import (
+            global_mcp_tool_registry,
+        )
+
+        prefix_root = normalize_server_name(get_server_prefix(server))
+        if server.spec_path and prefix_root:
+            openapi_key_prefix = prefix_root + MCP_TOOL_PREFIX_SEPARATOR
+            global_mcp_tool_registry.unregister_tools_with_prefix(openapi_key_prefix)
+
+        owned_raw: Set[str] = set()
+        for p in iter_known_server_prefixes(server):
+            if p:
+                owned_raw.add(p)
+        if server.name:
+            owned_raw.add(server.name)
+
+        owned_normalized = {normalize_server_name(x) for x in owned_raw}
+
+        stale_mapping_keys: List[str] = []
+        for tool_name, mapped_server in list(
+            self.tool_name_to_mcp_server_name_mapping.items()
+        ):
+            if mapped_server in owned_raw:
+                stale_mapping_keys.append(tool_name)
+            elif normalize_server_name(str(mapped_server)) in owned_normalized:
+                stale_mapping_keys.append(tool_name)
+
+        for key in stale_mapping_keys:
+            del self.tool_name_to_mcp_server_name_mapping[key]
+
     def remove_server(self, mcp_server: LiteLLM_MCPServerTable):
         """
         Remove a server from the registry
         """
-        if mcp_server.server_name in self.get_registry():
-            del self.registry[mcp_server.server_name]
-            verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_name}")
-        elif mcp_server.server_id in self.get_registry():
-            del self.registry[mcp_server.server_id]
-            verbose_logger.debug(f"Removed MCP Server: {mcp_server.server_id}")
+        evicted: Optional[MCPServer] = self.registry.pop(mcp_server.server_id, None)
+        if evicted is None and mcp_server.server_name:
+            evicted = self.registry.pop(mcp_server.server_name, None)
+        if evicted is not None:
+            verbose_logger.debug(
+                "Removed MCP Server: %s", mcp_server.server_id or mcp_server.server_name
+            )
+            self._cleanup_server_tool_routing_artifacts(evicted)
         else:
             verbose_logger.warning(
                 f"Server ID {mcp_server.server_id} not found in registry"
@@ -831,7 +872,8 @@ class MCPServerManager:
         # health probes can't reach it.
         if mcp_server.approval_status not in (None, "active", "approved"):
             if mcp_server.server_id in self.registry:
-                del self.registry[mcp_server.server_id]
+                evicted = self.registry.pop(mcp_server.server_id)
+                self._cleanup_server_tool_routing_artifacts(evicted)
             return
         try:
             if mcp_server.server_id in self.registry:

litellm/proxy/_experimental/mcp_server/tool_registry.py

@@ -59,6 +59,22 @@ class MCPToolRegistry:
             ]
         return list(self.tools.values())
 
+    def unregister_tools_with_prefix(self, prefix: str) -> int:
+        """Remove tools whose registered name starts with ``prefix``.
+
+        Used when an OpenAPI-backed MCP server leaves the runtime registry so
+        stale tool handlers cannot be invoked after eviction.
+        """
+        if not prefix:
+            return 0
+        removed = 0
+        for name in list(self.tools.keys()):
+            if name.startswith(prefix):
+                del self.tools[name]
+                removed += 1
+                verbose_logger.debug("Unregistered MCP tool %s", name)
+        return removed
+
     def convert_tools_to_mcp_sdk_tool_type(
         self, tools: List[MCPTool]
     ) -> List["MCPToolSDKTool"]:

litellm/proxy/common_utils/reset_budget_job.py

@@ -87,13 +87,13 @@ class ResetBudgetJob:
     async def _invalidate_user_api_key_cache_entry(cache_key: str) -> None:
         """Drop a stale management-cache entry so the next read fetches from DB.
 
-        Some entity types (notably tags and end-users) are not handled by
-        SpendCounterReseed.from_db, so when a spend counter expires the
-        budget check falls back to ``cached_obj.spend``. If that cached
-        object lingers in ``user_api_key_cache`` past a budget reset, the
-        stale ``.spend`` keeps the entity blocked indefinitely. Deleting
-        the cache entry forces the next auth-time fetch to reload the
-        zeroed row from Postgres.
+        Tags and end-users are not reseeded by ``SpendCounterReseed.from_db``;
+        for those, when the spend counter expires the budget check falls back
+        to ``cached_obj.spend``. Keys, orgs, and team memberships are reseeded
+        from the DB, but auth still may consult ``user_api_key_cache`` objects
+        whose ``.spend`` field can lag a cross-pod DB reset. Deleting the cache
+        entry forces the next auth-time fetch to reload the zeroed row from
+        Postgres.
         """
         try:
             from litellm.proxy.proxy_server import user_api_key_cache
@@ -119,11 +119,8 @@ class ResetBudgetJob:
         Generic cascade: zero spend on rows whose budget_id is in the reset set.
 
         ``cache_key_fn`` is optional: when provided, after the DB update each
-        matching row's entry in ``user_api_key_cache`` is also dropped. This
-        is required for entities whose spend counter is read with the cached
-        object's ``.spend`` as fallback (tags, end-users) — otherwise the
-        stale cached object pins enforcement to the pre-reset spend until
-        its TTL expires.
+        matching row's entry in ``user_api_key_cache`` is also dropped so
+        cached spend cannot stay pinned above the zeroed DB row after a reset.
         """
         budget_ids = [b.budget_id for b in budgets_to_reset if b.budget_id is not None]
         if not budget_ids:
@@ -161,6 +158,7 @@ class ResetBudgetJob:
             table=self.prisma_client.db.litellm_teammembership,
             counter_key_fn=lambda m: f"spend:team_member:{m.user_id}:{m.team_id}",
             log_subject="team memberships",
+            cache_key_fn=lambda m: f"{m.team_id}_{m.user_id}",
         )
 
     async def reset_budget_for_keys_linked_to_budgets(
@@ -178,6 +176,7 @@ class ResetBudgetJob:
             counter_key_fn=lambda k: f"spend:key:{k.token}",
             log_subject="keys",
             extra_where={"budget_duration": None, "spend": {"gt": 0}},
+            cache_key_fn=lambda k: k.token,
         )
 
     async def reset_budget_for_orgs_linked_to_budgets(
@@ -192,6 +191,7 @@ class ResetBudgetJob:
             counter_key_fn=lambda o: f"spend:org:{o.organization_id}",
             log_subject="orgs",
             extra_where={"spend": {"gt": 0}},
+            cache_key_fn=lambda o: f"org_id:{o.organization_id}",
         )
 
     async def reset_budget_for_tags_linked_to_budgets(

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -3380,6 +3380,50 @@ class TestApprovalStatusGate:
         )
         assert "evict-me" not in manager.registry
 
+    async def test_update_server_eviction_clears_openapi_routing_artifacts(
+        self, tmp_path
+    ):
+        """Rejecting a server must remove its OpenAPI tools and name mappings."""
+        from litellm.proxy._experimental.mcp_server.tool_registry import (
+            global_mcp_tool_registry,
+        )
+        from litellm.proxy._experimental.mcp_server.utils import (
+            add_server_prefix_to_name,
+            get_server_prefix,
+        )
+
+        manager = MCPServerManager()
+        await manager.add_server(
+            self._make_server("evict-openapi", MCPApprovalStatus.active)
+        )
+        assert "evict-openapi" in manager.registry
+
+        server = manager.registry["evict-openapi"]
+        server.spec_path = str(tmp_path / "unused.yaml")
+        prefix = get_server_prefix(server)
+        prefixed = add_server_prefix_to_name("demo_tool", prefix)
+
+        async def _noop_handler(**kwargs):
+            return None
+
+        global_mcp_tool_registry.register_tool(
+            name=prefixed,
+            description="demo",
+            input_schema={"type": "object"},
+            handler=_noop_handler,
+        )
+        manager.tool_name_to_mcp_server_name_mapping["demo_tool"] = prefix
+        manager.tool_name_to_mcp_server_name_mapping[prefixed] = prefix
+
+        await manager.update_server(
+            self._make_server("evict-openapi", MCPApprovalStatus.rejected)
+        )
+
+        assert "evict-openapi" not in manager.registry
+        assert prefixed not in global_mcp_tool_registry.tools
+        assert "demo_tool" not in manager.tool_name_to_mcp_server_name_mapping
+        assert prefixed not in manager.tool_name_to_mcp_server_name_mapping
+
     async def test_update_server_noop_for_unregistered_pending(self):
         # update_server called with a pending row that was never registered
         # should silently return without adding it. Locks in the early-return

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -1527,15 +1527,10 @@ def test_reset_budget_for_tags_linked_to_budgets_invalidates_each_tag_management
     assert deleted_keys == {"tag:tenant-a", "tag:tenant-b", "tag:tenant-c"}
 
 
-def test_reset_budget_for_keys_linked_to_budgets_does_not_touch_management_cache(
+def test_reset_budget_for_keys_linked_to_budgets_invalidates_management_cache(
     monkeypatch,
 ):
-    """Cache invalidation is opt-in: keys / orgs / team-members rely on
-    ``SpendCounterReseed.from_db`` (which DOES handle their counter keys),
-    so the cache_key_fn hook is intentionally not wired for them. This test
-    locks in that no-op so a future refactor doesn't accidentally start
-    clobbering the key cache (which would cost an extra DB round-trip per
-    reset cycle without fixing anything)."""
+    """Budget-tier key resets must drop the cached key object (hashed token key)."""
     counter_cache = _make_counter_invalidation_job(monkeypatch)
 
     expired_budget = type("B", (), {"budget_id": "budget-1"})
@@ -1552,7 +1547,61 @@ def test_reset_budget_for_keys_linked_to_budgets_does_not_touch_management_cache
     job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
     asyncio.run(job.reset_budget_for_keys_linked_to_budgets([expired_budget]))
 
-    counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited()
+    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
+        key="sk-linked"
+    )
+
+
+def test_reset_budget_for_orgs_linked_to_budgets_invalidates_management_cache(
+    monkeypatch,
+):
+    """Org rows use ``org_id:{id}`` in ``user_api_key_cache`` — evict on reset."""
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    expired_budget = type("B", (), {"budget_id": "budget-1"})
+    linked_org = type("Org", (), {"organization_id": "org-acme"})
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_organizationtable.find_many = AsyncMock(
+        return_value=[linked_org]
+    )
+    prisma_client.db.litellm_organizationtable.update_many = AsyncMock(
+        return_value={"count": 1}
+    )
+
+    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
+    asyncio.run(job.reset_budget_for_orgs_linked_to_budgets([expired_budget]))
+
+    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
+        key="org_id:org-acme"
+    )
+
+
+def test_reset_budget_for_team_members_invalidates_management_cache(monkeypatch):
+    """Team membership cache key matches auth: ``{team_id}_{user_id}``."""
+    counter_cache = _make_counter_invalidation_job(monkeypatch)
+
+    expired_budget = type("B", (), {"budget_id": "budget-1"})
+    membership = type(
+        "Membership",
+        (),
+        {"user_id": "alice", "team_id": "team-x", "budget_id": "budget-1"},
+    )
+
+    prisma_client = MagicMock()
+    prisma_client.db.litellm_teammembership.find_many = AsyncMock(
+        return_value=[membership]
+    )
+    prisma_client.db.litellm_teammembership.update_many = AsyncMock(
+        return_value={"count": 1}
+    )
+
+    job = ResetBudgetJob(proxy_logging_obj=MagicMock(), prisma_client=prisma_client)
+    asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
+
+    counter_cache.user_api_key_cache.async_delete_cache.assert_any_await(
+        key="team-x_alice"
+    )
 
 
 def test_reset_budget_for_tags_linked_to_budgets_management_cache_delete_failure_still_resets(