fix(mcp): fence an outbound-token write against an overlapping invalidation (#35398)

get_or_compute single-flights concurrent misses under a per-key lock, but
invalidate() deletes outside it. A mint already awaiting the IdP when an
invalidation ran wrote its result into the slot afterwards, so a bearer minted
before the invalidation was served after it for its full TTL; the upstream-401
retry then re-presented the bearer the server had just rejected.

The cache now keeps a per-key generation beside the single-flight lock.
invalidate() bumps it and get_or_compute snapshots it inside the lock right
before minting, skipping the cache write when it no longer matches. The token
is still returned to the caller it was minted for, so only the caching is
fenced. Generation and lock share one object that a running mint binds to a
local, which pins the weak map's entry for exactly as long as a write can still
land; a generation held weakly on its own would be collected across the await
and the re-check would silently pass.

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:

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