diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 89d645b6f8a..fa1f73cea77 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -448,6 +448,12 @@ async def _store_per_user_token_server_side( ) return # Don't warm Redis if DB write failed + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager, + ) + + await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server.server_id) + # Warm the Redis cache so the first subsequent MCP call is a cache hit ttl = _compute_per_user_token_ttl(server, expires_in) await mcp_per_user_token_cache.set( diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index b2128cb0553..e8c8d9c1f5a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -70,6 +70,9 @@ to_server_spec, to_subject, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InvalidatableOAuthTokenStore, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, ) @@ -552,9 +555,16 @@ def _obo_needs_endpoint_discovery( """ return auth_type == MCPAuth.oauth2_token_exchange and not (token_exchange_endpoint or token_url) - def __init__(self, cred_provider: Optional[UpstreamCredentialProvider] = None): + def __init__( + self, + cred_provider: Optional[UpstreamCredentialProvider] = None, + per_user_oauth_token_store: Optional[InvalidatableOAuthTokenStore] = None, + ): + self._per_user_oauth_token_store = per_user_oauth_token_store or LazyPerUserOAuthTokenStore( + self.get_mcp_server_by_id + ) self._cred_provider = cred_provider or UpstreamCredentialProvider( - oauth_token_store=LazyPerUserOAuthTokenStore(self.get_mcp_server_by_id), + oauth_token_store=self._per_user_oauth_token_store, token_exchanger=build_token_exchanger(), ) self.registry: dict[str, MCPServer] = {} @@ -3771,6 +3781,19 @@ async def has_user_oauth_token(self, server: MCPServer, user_api_key_auth: Optio return False return await self._cred_provider.has_user_token(to_subject(user_api_key_auth, None), spec) + async def invalidate_user_oauth_token_cache(self, user_id: str, server_id: str) -> None: + """Drop the v2 chain's cached token for ``(user_id, server_id)`` after the credential row + changes (re-auth, revoke), so the next resolve reads the new row instead of serving the + replaced token until its cache TTL. Best-effort: a cache-drop failure is logged, never + raised, because the DB write already succeeded and the TTL remains the backstop. + """ + try: + await self._per_user_oauth_token_store.invalidate(user_id, server_id) + except Exception as exc: # noqa: BLE001 - cache drop is best-effort; TTL is the backstop + verbose_logger.warning( + "Failed to invalidate cached MCP OAuth token for user=%s server=%s: %s", user_id, server_id, exc + ) + async def _resolve_oauth2_headers_for_tool_call( self, mcp_server: MCPServer, diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py index fd2cb2f3e06..c1c70cf9050 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/oauth_token_store.py @@ -69,6 +69,17 @@ class OAuthTokenStore(Protocol): async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: ... +class InvalidatableOAuthTokenStore(OAuthTokenStore, Protocol): + """An ``OAuthTokenStore`` whose cached entry for a ``(user, server)`` pair can be dropped. + + The write side calls ``invalidate`` after a (re)authorization or revocation changes the + credential row, so reads stop serving the replaced token immediately instead of until its + cache TTL. ``CachedOAuthTokenStore`` (the top of the per-user chain) satisfies this. + """ + + async def invalidate(self, user_id: str, server_id: str) -> None: ... + + class TokenRefresher(Protocol): """Mints a fresh token from an expired one and persists it, returning the new token. diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py index 3bc10f1a0eb..21001c09f25 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/per_user_oauth_store.py @@ -24,8 +24,8 @@ ) from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( CachedOAuthTokenStore, + InvalidatableOAuthTokenStore, OAuthToken, - OAuthTokenStore, RefreshCoordinator, RefreshingTokenStore, TokenCacheBackend, @@ -51,7 +51,7 @@ _DEFAULT_TTL_SECONDS = 300.0 ServerLookup = Callable[[str], "MCPServer | None"] -StoreBuilder = Callable[[ServerLookup], tuple[OAuthTokenStore, bool]] +StoreBuilder = Callable[[ServerLookup], tuple[InvalidatableOAuthTokenStore, bool]] async def _read_credential(user_id: str, server_id: str) -> dict[str, object] | None: @@ -185,7 +185,7 @@ def __init__( self._server_lookup = server_lookup self._store_builder = store_builder self._redis_available = redis_available - self._store: OAuthTokenStore | None = None + self._store: InvalidatableOAuthTokenStore | None = None self._uses_redis = False self._fetch_lock = asyncio.Condition() self._local_fetches = 0 @@ -203,7 +203,26 @@ async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: if not uses_redis: await self._finish_local_fetch() - async def _store_for_fetch(self) -> tuple[OAuthTokenStore, bool]: + async def invalidate(self, user_id: str, server_id: str) -> None: + """Drop the chain's cached entry for ``(user_id, server_id)`` after the credential row + changes (re-auth, revoke). Builds the chain if no fetch has run yet, so a shared (Redis) + cache entry written by another worker is dropped too; the in-process case is then a no-op + on an empty cache. + """ + if self._uses_redis: + store = self._store + if store is not None: + await store.invalidate(user_id, server_id) + return + + store, uses_redis = await self._store_for_fetch() + try: + await store.invalidate(user_id, server_id) + finally: + if not uses_redis: + await self._finish_local_fetch() + + async def _store_for_fetch(self) -> tuple[InvalidatableOAuthTokenStore, bool]: async with self._fetch_lock: while ( self._store is not None and not self._uses_redis and self._redis_available() and self._local_fetches > 0 diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 3597e75404a..14724a25705 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1874,6 +1874,11 @@ async def store_mcp_oauth_user_credential( expires_in=payload.expires_in, scopes=payload.scopes, ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager, + ) + + await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server_id) # Read back the persisted record so the response reflects the stored # expires_at rather than recomputing it here (which could diverge by # milliseconds or if the storage logic ever adds a grace period). @@ -1914,6 +1919,11 @@ async def delete_mcp_oauth_user_credential( await delete_user_credential(prisma_client, user_id, server_id) except RecordNotFoundError: pass # Already gone — treat as a successful delete + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 + global_mcp_server_manager, + ) + + await global_mcp_server_manager.invalidate_user_oauth_token_cache(user_id, server_id) return MCPOAuthUserCredentialStatus( server_id=server_id, has_credential=False, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py index f64ce594efa..ca32cf2bb8d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_per_user_oauth_store.py @@ -3,8 +3,8 @@ import pytest from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import ( + InvalidatableOAuthTokenStore, OAuthToken, - OAuthTokenStore, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.per_user_oauth_store import ( LazyPerUserOAuthTokenStore, @@ -16,11 +16,15 @@ class _RecordingStore: def __init__(self, access_token: str) -> None: self._access_token = access_token self.calls: list[tuple[str, str]] = [] + self.invalidations: list[tuple[str, str]] = [] async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: self.calls.append((user_id, server_id)) return OAuthToken(access_token=self._access_token) + async def invalidate(self, user_id: str, server_id: str) -> None: + self.invalidations.append((user_id, server_id)) + class _BlockingStore: def __init__(self, access_token: str) -> None: @@ -28,6 +32,7 @@ def __init__(self, access_token: str) -> None: self.started = asyncio.Event() self.release = asyncio.Event() self.calls: list[tuple[str, str]] = [] + self.invalidations: list[tuple[str, str]] = [] async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: self.calls.append((user_id, server_id)) @@ -35,6 +40,9 @@ async def fetch(self, user_id: str, server_id: str) -> OAuthToken | None: await self.release.wait() return OAuthToken(access_token=self._access_token) + async def invalidate(self, user_id: str, server_id: str) -> None: + self.invalidations.append((user_id, server_id)) + class _RedisAvailability: def __init__(self) -> None: @@ -59,7 +67,7 @@ async def test_lazy_store_rebuilds_when_redis_becomes_available() -> None: redis_available = _RedisAvailability() build_calls = 0 - def build_store(_server_lookup: ServerLookup) -> tuple[OAuthTokenStore, bool]: + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: nonlocal build_calls build_calls += 1 if redis_available.available: @@ -94,7 +102,7 @@ async def test_lazy_store_allows_concurrent_local_fetches_without_redis() -> Non redis_available = _RedisAvailability() build_calls = 0 - def build_store(_server_lookup: ServerLookup) -> tuple[OAuthTokenStore, bool]: + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: nonlocal build_calls build_calls += 1 return local_store, False @@ -127,7 +135,7 @@ async def test_lazy_store_waits_for_in_flight_local_fetch_before_redis_rebuild() redis_store = _RecordingStore("redis") redis_available = _RedisAvailability() - def build_store(_server_lookup: ServerLookup) -> tuple[OAuthTokenStore, bool]: + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: if redis_available.available: return redis_store, True return local_store, False @@ -158,3 +166,83 @@ def server_lookup(_server_id: str) -> None: assert second is not None and second.access_token == "redis" assert local_store.calls == [("u", "s")] assert redis_store.calls == [("u", "s")] + + +@pytest.mark.asyncio +async def test_lazy_store_invalidate_builds_chain_and_delegates() -> None: + local_store = _RecordingStore("local") + build_calls = 0 + + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: + nonlocal build_calls + build_calls += 1 + return local_store, False + + def server_lookup(_server_id: str) -> None: + return None + + store = LazyPerUserOAuthTokenStore( + server_lookup, + store_builder=build_store, + redis_available=_RedisAvailability(), + ) + + await store.invalidate("u", "s") + + assert build_calls == 1 + assert local_store.invalidations == [("u", "s")] + + +@pytest.mark.asyncio +async def test_lazy_store_invalidate_reaches_the_store_fetch_reads() -> None: + local_store = _RecordingStore("local") + build_calls = 0 + + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: + nonlocal build_calls + build_calls += 1 + return local_store, False + + def server_lookup(_server_id: str) -> None: + return None + + store = LazyPerUserOAuthTokenStore( + server_lookup, + store_builder=build_store, + redis_available=_RedisAvailability(), + ) + + await store.fetch("u", "s") + await store.invalidate("u", "s") + + assert build_calls == 1 + assert local_store.calls == [("u", "s")] + assert local_store.invalidations == [("u", "s")] + + +@pytest.mark.asyncio +async def test_lazy_store_invalidate_works_after_redis_chain_is_built() -> None: + redis_store = _RecordingStore("redis") + redis_available = _RedisAvailability() + redis_available.available = True + build_calls = 0 + + def build_store(_server_lookup: ServerLookup) -> tuple[InvalidatableOAuthTokenStore, bool]: + nonlocal build_calls + build_calls += 1 + return redis_store, True + + def server_lookup(_server_id: str) -> None: + return None + + store = LazyPerUserOAuthTokenStore( + server_lookup, + store_builder=build_store, + redis_available=redis_available, + ) + + await store.fetch("u", "s") + await store.invalidate("u", "s") + + assert build_calls == 1 + assert redis_store.invalidations == [("u", "s")] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index fe90fd45856..c808b17678a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4055,3 +4055,104 @@ async def test_oauth_authorization_server_404_for_unknown_server_name(): mcp_server_name="does_not_exist", ) assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_store_per_user_token_server_side_invalidates_v2_token_cache(): + """A token stored by the OAuth callback (code exchange or refresh) drops the v2 per-user + token cache entry, so egress stops serving the replaced token immediately instead of + until its TTL.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _store_per_user_token_server_side, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv-cb-1", + name="cb_server", + url="https://upstream.example/mcp", + transport="http", + auth_type=MCPAuth.oauth2, + ) + invalidate_mock = AsyncMock(return_value=None) + cache_set_mock = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.mcp_per_user_token_cache.set", + new=cache_set_mock, + ), + patch.object( + manager_module.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + new=invalidate_mock, + ), + ): + await _store_per_user_token_server_side( + server=server, + user_id="user-cb-1", + token_response={"access_token": "fresh-tok", "expires_in": 3600}, + ) + + invalidate_mock.assert_awaited_once_with("user-cb-1", "srv-cb-1") + cache_set_mock.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_store_per_user_token_server_side_skips_invalidate_when_db_write_fails(): + """A failed DB write neither warms the v1 cache nor drops the v2 cache entry; the + previously stored token is still the truth.""" + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _store_per_user_token_server_side, + ) + from litellm.types.mcp import MCPAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + server = MCPServer( + server_id="srv-cb-2", + name="cb_server_2", + url="https://upstream.example/mcp", + transport="http", + auth_type=MCPAuth.oauth2, + ) + invalidate_mock = AsyncMock(return_value=None) + cache_set_mock = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.utils.get_prisma_client_or_throw", + return_value=MagicMock(), + ), + patch( + "litellm.proxy._experimental.mcp_server.db.store_user_oauth_credential", + new=AsyncMock(side_effect=RuntimeError("db down")), + ), + patch( + "litellm.proxy._experimental.mcp_server.oauth2_token_cache.mcp_per_user_token_cache.set", + new=cache_set_mock, + ), + patch.object( + manager_module.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + new=invalidate_mock, + ), + ): + await _store_per_user_token_server_side( + server=server, + user_id="user-cb-2", + token_response={"access_token": "fresh-tok", "expires_in": 3600}, + ) + + invalidate_mock.assert_not_awaited() + cache_set_mock.assert_not_awaited() 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 6d3e09c6b75..c25f853b901 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 @@ -2798,6 +2798,39 @@ async def has_user_token(self, subject, spec): assert await manager.has_user_oauth_token(server, user_auth) is False assert calls == [] # short-circuited on the None spec, never hit the resolver + @pytest.mark.asyncio + async def test_invalidate_user_oauth_token_cache_delegates_to_store(self): + """The write side's cache drop reaches the same per-user store the resolver reads.""" + + class _Store: + def __init__(self) -> None: + self.invalidations: list[tuple[str, str]] = [] + + async def fetch(self, user_id: str, server_id: str): + return None + + async def invalidate(self, user_id: str, server_id: str) -> None: + self.invalidations.append((user_id, server_id)) + + store = _Store() + manager = MCPServerManager(per_user_oauth_token_store=store) + await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + assert store.invalidations == [("alice", "srv-1")] + + @pytest.mark.asyncio + 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.""" + + class _Store: + async def fetch(self, user_id: str, server_id: str): + return None + + async def invalidate(self, user_id: str, server_id: str) -> None: + raise RuntimeError("redis down") + + manager = MCPServerManager(per_user_oauth_token_store=_Store()) + await manager.invalidate_user_oauth_token_cache("alice", "srv-1") + @pytest.mark.asyncio async def test_resolve_oauth2_headers_no_user_id(self): """Skip lookup entirely when user_api_key_auth has no user_id.""" diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index e223140b573..1dcf78295fb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -3742,6 +3742,148 @@ async def test_delete_mcp_oauth_user_credential_only_deletes_oauth(): assert result.has_credential is False +@pytest.mark.asyncio +async def test_store_mcp_oauth_user_credential_invalidates_cached_token(): + """Re-authorizing via the Tools-tab persist drops the v2 per-user token cache entry, so + egress stops serving the replaced token immediately instead of until its TTL.""" + from litellm.proxy._types import MCPOAuthUserCredentialRequest + + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + store_mcp_oauth_user_credential, + ) + + server_id = "srv-inv-1" + user_id = "user-inv-1" + invalidate_mock = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", + new=AsyncMock(return_value=generate_mock_mcp_server_db_record(server_id=server_id)), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.store_user_oauth_credential", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value={"type": "oauth2", "access_token": "new-tok"}), + ), + patch.object( + manager_module.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + new=invalidate_mock, + ), + ): + await store_mcp_oauth_user_credential( + server_id=server_id, + payload=MCPOAuthUserCredentialRequest(access_token="new-tok", expires_in=3600), + user_api_key_dict=_make_user_auth(user_id), + ) + + invalidate_mock.assert_awaited_once_with(user_id, server_id) + + +@pytest.mark.asyncio +async def test_delete_mcp_oauth_user_credential_invalidates_cached_token(): + """Revoking a stored OAuth credential drops the v2 per-user token cache entry, so the + revoked token stops flowing upstream immediately instead of until its TTL.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + delete_mcp_oauth_user_credential, + ) + + server_id = "srv-inv-2" + user_id = "user-inv-2" + invalidate_mock = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value={"type": "oauth2", "access_token": "revoked-tok"}), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential", + new=AsyncMock(return_value=None), + ), + patch.object( + manager_module.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + new=invalidate_mock, + ), + ): + result = await delete_mcp_oauth_user_credential( + server_id=server_id, + user_api_key_dict=_make_user_auth(user_id), + ) + + invalidate_mock.assert_awaited_once_with(user_id, server_id) + assert result.has_credential is False + + +@pytest.mark.asyncio +async def test_delete_mcp_oauth_user_credential_invalidates_when_record_already_gone(): + """A concurrent delete can remove the row between the read and the delete; the cache may + still hold the revoked token, so the invalidate must fire even on RecordNotFoundError.""" + if not mgmt_endpoints.MCP_AVAILABLE: + pytest.skip("MCP module not installed") + + from litellm.proxy._experimental.mcp_server import mcp_server_manager as manager_module + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + delete_mcp_oauth_user_credential, + ) + + server_id = "srv-inv-3" + user_id = "user-inv-3" + invalidate_mock = AsyncMock(return_value=None) + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", + return_value=_make_prisma_client(), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.get_user_oauth_credential", + new=AsyncMock(return_value={"type": "oauth2", "access_token": "revoked-tok"}), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.delete_user_credential", + new=AsyncMock(side_effect=mgmt_endpoints.RecordNotFoundError({}, message="already gone")), + ), + patch.object( + manager_module.global_mcp_server_manager, + "invalidate_user_oauth_token_cache", + new=invalidate_mock, + ), + ): + result = await delete_mcp_oauth_user_credential( + server_id=server_id, + user_api_key_dict=_make_user_auth(user_id), + ) + + invalidate_mock.assert_awaited_once_with(user_id, server_id) + assert result.has_credential is False + + @pytest.mark.asyncio async def test_list_mcp_user_credentials_batch_server_fetch(): """list_mcp_user_credentials uses a single batch DB call, not N+1 queries."""