diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index f510dc16ad1..03c2f6fd73d 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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: diff --git a/litellm/proxy/_experimental/mcp_server/tool_registry.py b/litellm/proxy/_experimental/mcp_server/tool_registry.py index 58570aafadf..829be5be979 100644 --- a/litellm/proxy/_experimental/mcp_server/tool_registry.py +++ b/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"]: diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 71537cc62e6..f4aa0971ea9 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/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( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index c862de35c7e..b53420f0000 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/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 diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index fe45c52d41f..499b68b17ea 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/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(