mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(mcp): drop the cached per-user OAuth token when the credential row changes (#32302)
* fix(mcp): drop the cached per-user OAuth token when the credential row changes The v2 authorization_code chain Cached(Refreshing(V2PerUserTokenStore)) caches a positive token until its expires_at (or 300s without one), and CachedOAuthTokenStore.invalidate had no callers, so a re-authorization or revocation wrote the DB while egress kept serving the replaced token from the in-process cache until its TTL. LazyPerUserOAuthTokenStore now exposes invalidate, MCPServerManager threads it to the write side, and the three credential write sites (the OAuth callback, the Tools-tab persist endpoint, and the revoke endpoint) drop the cache entry after the row changes. The v2 refresher's own persist stays untouched; RefreshingTokenStore already feeds the rotated token back into the cache in the same fetch * test(mcp): pin cache invalidation on the revoke already-gone branch Greptile's review flagged that only the happy-path delete asserted the invalidate; a refactor moving the call inside the try block would silently skip the cache drop when the row was already deleted by a concurrent request while the cache still held the revoked token. The new test fails on exactly that mutation * test(mcp): cover invalidate on the redis-backed lazy store path Codecov flagged the redis fast path of LazyPerUserOAuthTokenStore.invalidate as unexercised; the existing invalidate tests only ran the no-redis chain. The new test builds the redis chain via a fetch and asserts a subsequent invalidate reaches the same store instance without a rebuild
This commit is contained in:
parent
bcd52754de
commit
1fb2b4aef4
9 changed files with 443 additions and 10 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -70,6 +70,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import
|
|||
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,
|
||||
)
|
||||
|
|
@ -689,9 +692,16 @@ class MCPServerManager:
|
|||
"""
|
||||
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] = {}
|
||||
|
|
@ -3922,6 +3932,19 @@ class MCPServerManager:
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -24,8 +24,8 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.dual_cache_toke
|
|||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
CachedOAuthTokenStore,
|
||||
InvalidatableOAuthTokenStore,
|
||||
OAuthToken,
|
||||
OAuthTokenStore,
|
||||
RefreshCoordinator,
|
||||
RefreshingTokenStore,
|
||||
TokenCacheBackend,
|
||||
|
|
@ -51,7 +51,7 @@ if TYPE_CHECKING:
|
|||
_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 @@ class LazyPerUserOAuthTokenStore:
|
|||
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 @@ class LazyPerUserOAuthTokenStore:
|
|||
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
|
||||
|
|
|
|||
|
|
@ -1913,6 +1913,11 @@ if MCP_AVAILABLE:
|
|||
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).
|
||||
|
|
@ -1953,6 +1958,11 @@ if MCP_AVAILABLE:
|
|||
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,
|
||||
|
|
|
|||
|
|
@ -3,8 +3,8 @@ import asyncio
|
|||
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 @@ class _BlockingStore:
|
|||
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 @@ class _BlockingStore:
|
|||
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 @@ async def test_lazy_store_waits_for_in_flight_local_fetch_before_redis_rebuild()
|
|||
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")]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -2940,6 +2940,39 @@ class TestMCPServerManager:
|
|||
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."""
|
||||
|
|
|
|||
|
|
@ -3566,6 +3566,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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue