refactor(mcp): align the invalidation code with the v2 DI and typing discipline

The purge takes an injectable invalidate_token_cache callable defaulting to the manager's shared
invalidation, and MCPServerManager takes an injectable per_user_token_cache alongside the existing
per_user_oauth_token_store, so tests inject fakes instead of monkeypatching the global manager and
the module-level cache. The new identity helpers drop Any for object throughout

litellm/proxy/_experimental/mcp_server/db.py

@@ -3,7 +3,7 @@ import binascii
 import hashlib
 import json
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Set, Union, cast
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterable, List, Optional, Set, Union, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
@@ -1070,7 +1070,7 @@ async def list_user_oauth_credentials(
     return results
 
 
-def _decrypted_credential_field(creds: Dict[str, Any], field: str) -> Any:
+def _decrypted_credential_field(creds: Dict[str, object], field: str) -> object:
     """Return one credential field decrypted with the global salt key; non-string and legacy
     plaintext values come back unchanged (decrypt_value_helper returns the original on failure)."""
     value = creds.get(field)
@@ -1084,7 +1084,7 @@ def _decrypted_credential_field(creds: Dict[str, Any], field: str) -> Any:
     )
 
 
-def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]:
+def mcp_oauth_token_identity(server: object) -> tuple[object, ...]:
     """The upstream-OAuth-token-determining fields of an MCP server: the resource/audience (url, or
     spec_path for OpenAPI servers), the OAuth mode/grant (auth_type, oauth2_flow), the
     authorization-server endpoints, and the OAuth client + scopes. Mirrors the dashboard's
@@ -1098,12 +1098,12 @@ def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]:
     creds = getattr(server, "credentials", None)
     if isinstance(creds, str):
         try:
-            parsed: Any = json.loads(creds)
+            parsed: object = json.loads(creds)
         except ValueError:
             parsed = None
     else:
         parsed = creds
-    creds_dict: Dict[str, Any] = parsed if isinstance(parsed, dict) else {}
+    creds_dict: Dict[str, object] = parsed if isinstance(parsed, dict) else {}
     return (
         getattr(server, "url", None),
         getattr(server, "spec_path", None),
@@ -1118,25 +1118,35 @@ def mcp_oauth_token_identity(server: Any) -> tuple[Any, ...]:
     )
 
 
-async def purge_user_oauth_credentials_for_server(prisma_client: PrismaClient, server_id: str) -> int:
+async def purge_user_oauth_credentials_for_server(
+    prisma_client: PrismaClient,
+    server_id: str,
+    invalidate_token_cache: Optional[Callable[[str, str], Awaitable[None]]] = None,
+) -> int:
     """Delete every stored per-user OAuth credential for a server and invalidate each user's cached
     token everywhere it can be served from (the legacy per-user token cache and the v2 per-user OAuth
     token store), so no user keeps a token minted for a superseded configuration. Called when a server
     update changes a mint-relevant field (see mcp_oauth_token_identity). Returns the number of rows
     removed. A row inserted between the find and the delete is removed from the DB but cannot be
     evicted from the caches (its user_id was never seen); that case is detected, logged, and bounded
-    by the cache TTL."""
+    by the cache TTL.
+
+    invalidate_token_cache is injectable for tests; it defaults to the manager's shared
+    invalidate_user_oauth_token_cache, the single invalidation point for per-user tokens."""
     repo = MCPUserCredentialsRepository(prisma_client)
     rows = await repo.table.find_many(where={"server_id": server_id})
     if not rows:
         return 0
     deleted_count = await repo.table.delete_many(where={"server_id": server_id})
-    from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
-        global_mcp_server_manager,
-    )
+    if invalidate_token_cache is None:
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
+        )
+
+        invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache
 
     for row in rows:
-        await global_mcp_server_manager.invalidate_user_oauth_token_cache(row.user_id, server_id)
+        await invalidate_token_cache(row.user_id, server_id)
     if deleted_count != len(rows):
         verbose_proxy_logger.warning(
             "MCP server %s: purge removed %d credential row(s) but %d were enumerated; "

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -58,6 +58,7 @@ from litellm.proxy._experimental.mcp_server.sampling_handler import (
     MCP_SAMPLING_AVAILABLE,
 )
 from litellm.proxy._experimental.mcp_server.oauth2_token_cache import (
+    MCPPerUserTokenCache,
     mcp_per_user_token_cache,
     resolve_mcp_auth,
 )
@@ -802,10 +803,12 @@ class MCPServerManager:
         self,
         cred_provider: Optional[UpstreamCredentialProvider] = None,
         per_user_oauth_token_store: Optional[InvalidatableOAuthTokenStore] = None,
+        per_user_token_cache: Optional[MCPPerUserTokenCache] = None,
     ):
         self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore(
             self.get_mcp_server_by_id
         )
+        self._per_user_token_cache = per_user_token_cache or mcp_per_user_token_cache
         self._cred_provider = cred_provider or UpstreamCredentialProvider(
             oauth_token_store=self._per_user_oauth_token_store,
             token_exchanger=build_token_exchanger(),
@@ -4070,7 +4073,7 @@ class MCPServerManager:
             verbose_logger.warning(
                 "Failed to invalidate cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc
             )
-        await mcp_per_user_token_cache.delete(user_id, server_id)
+        await self._per_user_token_cache.delete(user_id, server_id)
 
     async def _resolve_oauth2_headers_for_tool_call(
         self,

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

@@ -149,11 +149,11 @@ def test_mcp_oauth_token_identity_detects_change_under_encryption():
 
 
 @pytest.mark.asyncio
-async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(monkeypatch):
-    """The purge must route each (user, server) through the manager's shared invalidation, which is
-    the single point covering both the legacy per-user token cache and the v2 per-user OAuth token
-    store; evicting only one cache lets the other keep serving a token minted for the old config."""
-    from litellm.proxy._experimental.mcp_server import mcp_server_manager
+async def test_purge_user_oauth_credentials_for_server_invalidates_every_store():
+    """The purge must route each (user, server) through the injected invalidator (defaulting to the
+    manager's shared invalidation, the single point covering both the legacy per-user token cache and
+    the v2 per-user OAuth token store); evicting only one cache lets the other keep serving a token
+    minted for the old config."""
     from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server
 
     r1 = MagicMock(user_id="alice", server_id="srv-1")
@@ -163,13 +163,11 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(m
     prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2)
 
     invalidations = []
-    monkeypatch.setattr(
-        mcp_server_manager.global_mcp_server_manager,
-        "invalidate_user_oauth_token_cache",
-        AsyncMock(side_effect=lambda uid, sid: invalidations.append((uid, sid))),
-    )
 
-    purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1")
+    async def record_invalidation(user_id: str, server_id: str) -> None:
+        invalidations.append((user_id, server_id))
+
+    purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=record_invalidation)
 
     assert purged == 2
     prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once()
@@ -179,7 +177,6 @@ async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(m
 @pytest.mark.asyncio
 async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypatch):
     from litellm.proxy._experimental.mcp_server import db as db_module
-    from litellm.proxy._experimental.mcp_server import mcp_server_manager
     from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server
 
     prisma = MagicMock()
@@ -187,15 +184,10 @@ async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypat
         return_value=[MagicMock(user_id="alice", server_id="srv-1")]
     )
     prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2)
-    monkeypatch.setattr(
-        mcp_server_manager.global_mcp_server_manager,
-        "invalidate_user_oauth_token_cache",
-        AsyncMock(),
-    )
     warning = MagicMock()
     monkeypatch.setattr(db_module.verbose_proxy_logger, "warning", warning)
 
-    purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1")
+    purged = await purge_user_oauth_credentials_for_server(prisma, "srv-1", invalidate_token_cache=AsyncMock())
 
     assert purged == 2
     warning.assert_called_once()

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

@@ -3328,11 +3328,10 @@ class TestMCPServerManager:
         assert store.invalidations == [("alice", "srv-1")]
 
     @pytest.mark.asyncio
-    async def test_invalidate_user_oauth_token_cache_drops_legacy_cache_too(self, monkeypatch):
+    async def test_invalidate_user_oauth_token_cache_drops_legacy_cache_too(self):
         """A per-user token can be served from the legacy per-user token cache as well as the v2
         store; the shared invalidation must evict both, or the path not evicted keeps serving a
         token minted for a replaced credential row until its TTL."""
-        from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module
 
         class _Store:
             async def fetch(self, user_id: str, server_id: str):
@@ -3341,21 +3340,22 @@ class TestMCPServerManager:
             async def invalidate(self, user_id: str, server_id: str) -> None:
                 return None
 
-        legacy_deletes: list[tuple[str, str]] = []
-        monkeypatch.setattr(
-            manager_module.mcp_per_user_token_cache,
-            "delete",
-            AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))),
-        )
-        manager = MCPServerManager(per_user_oauth_token_store=_Store())
+        class _LegacyCache:
+            def __init__(self) -> None:
+                self.deletes: list[tuple[str, str]] = []
+
+            async def delete(self, user_id: str, server_id: str) -> None:
+                self.deletes.append((user_id, server_id))
+
+        legacy_cache = _LegacyCache()
+        manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=legacy_cache)
         await manager.invalidate_user_oauth_token_cache("alice", "srv-1")
-        assert legacy_deletes == [("alice", "srv-1")]
+        assert legacy_cache.deletes == [("alice", "srv-1")]
 
     @pytest.mark.asyncio
-    async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self, monkeypatch):
+    async def test_invalidate_user_oauth_token_cache_swallows_store_errors(self):
         """A cache-drop failure must not fail the credential write that triggered it, and the
         legacy cache must still be evicted after the v2 store drop fails."""
-        from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module
 
         class _Store:
             async def fetch(self, user_id: str, server_id: str):
@@ -3364,15 +3364,17 @@ class TestMCPServerManager:
             async def invalidate(self, user_id: str, server_id: str) -> None:
                 raise RuntimeError("redis down")
 
-        legacy_deletes: list[tuple[str, str]] = []
-        monkeypatch.setattr(
-            manager_module.mcp_per_user_token_cache,
-            "delete",
-            AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))),
-        )
-        manager = MCPServerManager(per_user_oauth_token_store=_Store())
+        class _LegacyCache:
+            def __init__(self) -> None:
+                self.deletes: list[tuple[str, str]] = []
+
+            async def delete(self, user_id: str, server_id: str) -> None:
+                self.deletes.append((user_id, server_id))
+
+        legacy_cache = _LegacyCache()
+        manager = MCPServerManager(per_user_oauth_token_store=_Store(), per_user_token_cache=legacy_cache)
         await manager.invalidate_user_oauth_token_cache("alice", "srv-1")
-        assert legacy_deletes == [("alice", "srv-1")]
+        assert legacy_cache.deletes == [("alice", "srv-1")]
 
     @pytest.mark.asyncio
     async def test_resolve_oauth2_headers_no_user_id(self):