diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py index 84f714db449..6f1e5baa109 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py @@ -121,6 +121,23 @@ class TokenEndpointClient: return Ok(ExchangedToken(access_token=parsed.access_token, expires_in=parsed.expires_in)) +class _KeyGuard: + """The per-key single-flight lock plus the invalidation generation that lock protects. + + Both live on one object so their lifetimes cannot diverge. `get_or_compute` binds the guard to + a local for its whole critical section, which keeps the weak map's entry alive for as long as + that compute could still write; an `invalidate` overlapping the compute therefore reaches the + very same object and its bump is guaranteed to be observed. Conversely a guard nobody holds is + collectible precisely because no write is outstanding for it to fence. + """ + + __slots__ = ("__weakref__", "generation", "lock") + + def __init__(self) -> None: + self.lock = asyncio.Lock() + self.generation = 0 + + class ExchangedTokenCache: """Memoizes the final token string per key, single-flighting concurrent misses on one lock.""" @@ -129,7 +146,7 @@ class ExchangedTokenCache: max_size_in_memory=MCP_TOKEN_EXCHANGE_CACHE_MAX_SIZE, default_ttl=MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL, ) - self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + self._guards: weakref.WeakValueDictionary[str, _KeyGuard] = weakref.WeakValueDictionary() async def get_or_compute( self, @@ -144,28 +161,50 @@ class ExchangedTokenCache: guaranteeing the token it gets back was minted for the *current* inputs: a stored entry whose fingerprint differs reads as a miss and is re-minted over. That keeps eviction addressable without the key having to encode the credential material it protects. + + An `invalidate` landing while `compute` is in flight wins over that compute's write. The + token is still returned to the caller it was minted for, but it is not stored, so the next + resolution re-mints rather than serving a bearer that predates the invalidation for the + rest of its TTL. """ cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) - async with self._lock(cache_key): + guard = self._guard(cache_key) + async with guard.lock: cached = self._get(cache_key, fingerprint) if cached is not None: return Ok(cached) + generation = guard.generation match await compute(): case Ok(token): - self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped - cache_key, - (fingerprint, token.access_token), - ttl=_cache_ttl_seconds(token.expires_in), - ) + if guard.generation == generation: + self._store(cache_key, fingerprint, token) return Ok(token.access_token) case Error(err): return Error(err) def invalidate(self, cache_key: str) -> None: - """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401).""" + """Evict one cached token so the next `get_or_compute` re-mints (e.g. after an upstream 401). + + Bumping the guard's generation is what makes the eviction stick against a compute already + awaiting the token endpoint: that compute snapshotted the old generation and so skips its + write. No guard means no compute is in flight, since an in-flight one pins its own. + + Stays synchronous: callers invalidate from plain `def`s. + """ self._cache.delete_cache(cache_key) # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + guard = self._guards.get(cache_key) + if guard is None: + return + guard.generation += 1 + + def _store(self, cache_key: str, fingerprint: str, token: ExchangedToken) -> None: + self._cache.set_cache( # pyright: ignore[reportUnknownMemberType] # InMemoryCache is untyped + cache_key, + (fingerprint, token.access_token), + ttl=_cache_ttl_seconds(token.expires_in), + ) def _get(self, cache_key: str, fingerprint: str) -> str | None: """The stored token, or None when absent or minted for different inputs. @@ -180,12 +219,12 @@ class ExchangedTokenCache: return None return token if stored_fingerprint == fingerprint else None - def _lock(self, cache_key: str) -> asyncio.Lock: - lock = self._locks.get(cache_key) - if lock is None: - lock = asyncio.Lock() - self._locks[cache_key] = lock - return lock + def _guard(self, cache_key: str) -> _KeyGuard: + guard = self._guards.get(cache_key) + if guard is None: + guard = _KeyGuard() + self._guards[cache_key] = guard + return guard def _cache_ttl_seconds(expires_in: int | None) -> int: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py index 5f277db2f72..db3a1a386a3 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_token_endpoint.py @@ -7,6 +7,7 @@ cache's hit/single-flight behavior. Each assertion fails under a real mutation o """ import asyncio +import gc import json from unittest.mock import AsyncMock, MagicMock, patch @@ -359,6 +360,110 @@ async def test_cache_invalidate_only_evicts_the_named_key(): assert calls == 2 +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_is_not_overwritten_by_that_compute(): + """A bearer minted before an invalidation must never be served after it. + + The compute is suspended at the token endpoint when the invalidation lands, so its write is + the one that would resurrect the evicted bearer for the rest of its TTL. The caller it was + minted for still gets it; the *cache* is what the invalidation is about. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + release_mint.set() + + raced = await in_flight + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_invalidate_mid_compute_survives_garbage_collection(): + """The record of an invalidation must outlive a collection cycle taken mid-compute. + + Per-key state is held weakly so idle keys do not accumulate. If the state a compute checks + before writing were collectible while that compute is suspended, the check would read as + "nothing was invalidated" and the stale write would land; the running compute has to pin it. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + + assert not in_flight.done() + cache.invalidate("slot") + gc.collect() + release_mint.set() + await in_flight + + after = await cache.get_or_compute("slot", re_mint, fingerprint="fp") + assert isinstance(after, Ok) and after.ok == "bearer-minted-after-invalidation" + + +@pytest.mark.asyncio +async def test_cache_stores_a_compute_that_started_after_the_invalidation(): + """Only the mint that predates the invalidation loses its write. + + A caller queued behind the single-flight lock computes after the eviction, so its token is + fresh and must be cached; otherwise the fix would trade one stale bearer for re-minting on + every subsequent resolution. + """ + cache = ExchangedTokenCache() + mint_started, release_mint = asyncio.Event(), asyncio.Event() + + async def slow_mint(): + mint_started.set() + await release_mint.wait() + return _ok_token("bearer-minted-before-invalidation") + + async def re_mint(): + return _ok_token("bearer-minted-after-invalidation") + + async def must_not_run(): + pytest.fail("the mint that followed the invalidation should have been cached") + + in_flight = asyncio.create_task(cache.get_or_compute("slot", slow_mint, fingerprint="fp")) + await mint_started.wait() + queued = asyncio.create_task(cache.get_or_compute("slot", re_mint, fingerprint="fp")) + await asyncio.sleep(0) + + assert not queued.done() + cache.invalidate("slot") + release_mint.set() + + raced, fresh = await asyncio.gather(in_flight, queued) + assert isinstance(raced, Ok) and raced.ok == "bearer-minted-before-invalidation" + assert isinstance(fresh, Ok) and fresh.ok == "bearer-minted-after-invalidation" + + served = await cache.get_or_compute("slot", must_not_run, fingerprint="fp") + assert isinstance(served, Ok) and served.ok == "bearer-minted-after-invalidation" + + @pytest.mark.asyncio async def test_cache_does_not_store_a_failed_compute(): cache = ExchangedTokenCache()