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
This commit is contained in:
Tin 2026-07-09 12:52:18 -07:00
parent 48124734a0
commit 42388c3d68
4 changed files with 57 additions and 50 deletions

View file

@ -3,7 +3,7 @@ import binascii
import hashlib import hashlib
import json import json
from datetime import datetime, timedelta, timezone 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._logging import verbose_proxy_logger
from litellm._uuid import uuid from litellm._uuid import uuid
@ -1070,7 +1070,7 @@ async def list_user_oauth_credentials(
return results 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 """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).""" plaintext values come back unchanged (decrypt_value_helper returns the original on failure)."""
value = creds.get(field) 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 """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 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 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) creds = getattr(server, "credentials", None)
if isinstance(creds, str): if isinstance(creds, str):
try: try:
parsed: Any = json.loads(creds) parsed: object = json.loads(creds)
except ValueError: except ValueError:
parsed = None parsed = None
else: else:
parsed = creds 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 ( return (
getattr(server, "url", None), getattr(server, "url", None),
getattr(server, "spec_path", 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 """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 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 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 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 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 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) repo = MCPUserCredentialsRepository(prisma_client)
rows = await repo.table.find_many(where={"server_id": server_id}) rows = await repo.table.find_many(where={"server_id": server_id})
if not rows: if not rows:
return 0 return 0
deleted_count = await repo.table.delete_many(where={"server_id": server_id}) deleted_count = await repo.table.delete_many(where={"server_id": server_id})
from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( if invalidate_token_cache is None:
global_mcp_server_manager, 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: 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): if deleted_count != len(rows):
verbose_proxy_logger.warning( verbose_proxy_logger.warning(
"MCP server %s: purge removed %d credential row(s) but %d were enumerated; " "MCP server %s: purge removed %d credential row(s) but %d were enumerated; "

View file

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

View file

@ -149,11 +149,11 @@ def test_mcp_oauth_token_identity_detects_change_under_encryption():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_purge_user_oauth_credentials_for_server_invalidates_every_store(monkeypatch): async def test_purge_user_oauth_credentials_for_server_invalidates_every_store():
"""The purge must route each (user, server) through the manager's shared invalidation, which is """The purge must route each (user, server) through the injected invalidator (defaulting to the
the single point covering both the legacy per-user token cache and the v2 per-user OAuth token manager's shared invalidation, the single point covering both the legacy per-user token cache and
store; evicting only one cache lets the other keep serving a token minted for the old config.""" the v2 per-user OAuth token store); evicting only one cache lets the other keep serving a token
from litellm.proxy._experimental.mcp_server import mcp_server_manager minted for the old config."""
from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server
r1 = MagicMock(user_id="alice", server_id="srv-1") 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) prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2)
invalidations = [] 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 assert purged == 2
prisma.db.litellm_mcpusercredentials.delete_many.assert_awaited_once() 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 @pytest.mark.asyncio
async def test_purge_user_oauth_credentials_for_server_logs_raced_rows(monkeypatch): 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 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 from litellm.proxy._experimental.mcp_server.db import purge_user_oauth_credentials_for_server
prisma = MagicMock() 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")] return_value=[MagicMock(user_id="alice", server_id="srv-1")]
) )
prisma.db.litellm_mcpusercredentials.delete_many = AsyncMock(return_value=2) 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() warning = MagicMock()
monkeypatch.setattr(db_module.verbose_proxy_logger, "warning", warning) 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 assert purged == 2
warning.assert_called_once() warning.assert_called_once()

View file

@ -3328,11 +3328,10 @@ class TestMCPServerManager:
assert store.invalidations == [("alice", "srv-1")] assert store.invalidations == [("alice", "srv-1")]
@pytest.mark.asyncio @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 """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 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.""" 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: class _Store:
async def fetch(self, user_id: str, server_id: str): 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: async def invalidate(self, user_id: str, server_id: str) -> None:
return None return None
legacy_deletes: list[tuple[str, str]] = [] class _LegacyCache:
monkeypatch.setattr( def __init__(self) -> None:
manager_module.mcp_per_user_token_cache, self.deletes: list[tuple[str, str]] = []
"delete",
AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), async def delete(self, user_id: str, server_id: str) -> None:
) self.deletes.append((user_id, server_id))
manager = MCPServerManager(per_user_oauth_token_store=_Store())
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") 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 @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 """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.""" 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: class _Store:
async def fetch(self, user_id: str, server_id: str): 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: async def invalidate(self, user_id: str, server_id: str) -> None:
raise RuntimeError("redis down") raise RuntimeError("redis down")
legacy_deletes: list[tuple[str, str]] = [] class _LegacyCache:
monkeypatch.setattr( def __init__(self) -> None:
manager_module.mcp_per_user_token_cache, self.deletes: list[tuple[str, str]] = []
"delete",
AsyncMock(side_effect=lambda uid, sid: legacy_deletes.append((uid, sid))), async def delete(self, user_id: str, server_id: str) -> None:
) self.deletes.append((user_id, server_id))
manager = MCPServerManager(per_user_oauth_token_store=_Store())
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") 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 @pytest.mark.asyncio
async def test_resolve_oauth2_headers_no_user_id(self): async def test_resolve_oauth2_headers_no_user_id(self):