Merge 5ec24bd85d into ec94a1f82a

litellm/proxy/_experimental/mcp_server/outbound_credentials/token_endpoint.py

@@ -124,6 +124,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."""
 
@@ -132,7 +149,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,
@@ -147,28 +164,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.
@@ -183,12 +222,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
 
@@ -388,6 +389,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()