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:
tin-berri 2026-07-07 23:59:35 -07:00 committed by GitHub
parent bcd52754de
commit 1fb2b4aef4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 443 additions and 10 deletions

View file

@ -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(

View file

@ -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,

View file

@ -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.

View file

@ -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

View file

@ -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,

View file

@ -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")]

View file

@ -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()

View file

@ -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."""

View file

@ -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."""