From 42388c3d689807f5e94de9311c40a09448bb488f Mon Sep 17 00:00:00 2001 From: Tin Date: Thu, 9 Jul 2026 12:52:18 -0700 Subject: [PATCH] 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 | 32 +++++++++----- .../mcp_server/mcp_server_manager.py | 5 ++- .../mcp_server/test_db_credentials.py | 28 +++++-------- .../mcp_server/test_mcp_server_manager.py | 42 ++++++++++--------- 4 files changed, 57 insertions(+), 50 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 135a9e055d6..96b28afc093 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/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; " diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cc4a9e63105..41a87a17d58 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/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, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 51641991ef9..1615d81fae9 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/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() 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 e11d78d07ae..97fdd186dda 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 @@ -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):