fix(mcp): bypass stale assertion cache during renewal

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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

@@ -387,12 +387,14 @@ class RefreshingSSOAssertionStore:
         inner: SSOAssertionStore,
         refresher: SSOAssertionRefresher,
         *,
+        fresh_read: AssertionRead,
         coordinator_factory: CoordinatorFactory = runtime_refresh_coordinator,
         expiry_skew_seconds: float = _DEFAULT_EXPIRY_SKEW_SECONDS,
         clock: Callable[[], datetime] = lambda: datetime.now(timezone.utc),
     ) -> None:
         self._inner = inner
         self._refresher = refresher
+        self._fresh_read = fresh_read
         self._coordinator_factory = coordinator_factory
         self._in_process_coordinator = InProcessRefreshCoordinator()
         self._distributed_coordinator: RefreshCoordinator | None = None
@@ -409,7 +411,7 @@ class RefreshingSSOAssertionStore:
             refresh=lambda: self._renew(user_id),
             reread=lambda: self._reread_renewed(user_id),
         )
-        return await self._inner.fetch(user_id)
+        return await self._fresh_read(user_id)
 
     def _expiring(self, assertion: SSOIdentityAssertion | None) -> bool:
         return assertion is not None and assertion_expired(assertion, self._clock() + self._skew)
@@ -428,7 +430,7 @@ class RefreshingSSOAssertionStore:
         """The elected renewal, judged from a fresh read so a rotation another replica just landed is
         never redeemed again. Returns nothing: the inner store, not this return value, is what every
         caller reads afterwards, so the winner and the losers cannot disagree."""
-        latest: Final = await self._inner.fetch(user_id)
+        latest: Final = await self._fresh_read(user_id)
         if latest is None or not self._expiring(latest):
             return
         match await self._refresher.refresh(user_id, latest):
@@ -448,7 +450,7 @@ class RefreshingSSOAssertionStore:
         It answers retryable 503 instead of guessing a sign-in challenge; the retry runs uncontended
         and settles the outcome itself.
         """
-        latest: Final = await self._inner.fetch(user_id)
+        latest: Final = await self._fresh_read(user_id)
         if self._expiring(latest):
             raise AssertionStoreUnavailable(
                 f"the IdP identity assertion for user_id={user_id} was being renewed by another replica "
@@ -458,4 +460,10 @@ class RefreshingSSOAssertionStore:
 
 def default_sso_assertion_store() -> SSOAssertionStore:
     """The live read seam for the ``id_jag`` arm: the stored assertion, renewed when it is stale."""
-    return RefreshingSSOAssertionStore(DbSSOAssertionStore(), SSOAssertionRefresher(HttpxTokenEndpointTransport()))
+    db_store: Final = DbSSOAssertionStore()
+    fresh_read: Final = db_store.fetch_uncached
+    return RefreshingSSOAssertionStore(
+        db_store,
+        SSOAssertionRefresher(HttpxTokenEndpointTransport(), read=fresh_read),
+        fresh_read=fresh_read,
+    )

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

@@ -278,6 +278,12 @@ class DbSSOAssertionStore:
         except Exception as exc:  # noqa: BLE001  # any driver/storage failure is an outage, not an absence
             raise AssertionStoreUnavailable(str(exc)) from exc
 
+    async def fetch_uncached(self, user_id: str) -> SSOIdentityAssertion | None:
+        try:
+            return await _read_assertion_from_db(user_id)
+        except Exception as exc:  # noqa: BLE001  # any driver/storage failure is an outage, not an absence
+            raise AssertionStoreUnavailable(str(exc)) from exc
+
 
 async def rotate_sso_identity_assertions_master_key(prisma_client: PrismaClient, new_master_key: str) -> None:
     """Re-encrypt every stored assertion under ``new_master_key`` during a salt-key rotation,

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py

@@ -613,7 +613,7 @@ async def test_id_jag_renews_an_expired_stored_assertion_instead_of_challenging(
     provider = UpstreamCredentialProvider(
         token_endpoint=endpoint,
         sso_assertion_store=RefreshingSSOAssertionStore(
-            _Inner(), refresher, coordinator_factory=lambda: None
+            _Inner(), refresher, fresh_read=_read, coordinator_factory=lambda: None
         ),
     )
 

tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py

@@ -86,6 +86,7 @@ class _FakeRows:
 
     def __init__(self, rows: dict[str, SSOIdentityAssertion] | None = None) -> None:
         self.rows: dict[str, SSOIdentityAssertion] = dict(rows or {})
+        self.cached_rows: dict[str, SSOIdentityAssertion] = {}
         self.reads: list[str] = []
         self.writes: list[tuple[str, SSOIdentityAssertion]] = []
 
@@ -94,6 +95,11 @@ class _FakeRows:
         # A real suspension point, so concurrent callers interleave here instead of running to
         # completion one at a time and never actually racing.
         await asyncio.sleep(0)
+        return self.cached_rows.get(user_id, self.rows.get(user_id))
+
+    async def fetch_fresh(self, user_id: str) -> SSOIdentityAssertion | None:
+        self.reads.append(user_id)
+        await asyncio.sleep(0)
         return self.rows.get(user_id)
 
     async def write(self, user_id: str, assertion: SSOIdentityAssertion) -> None:
@@ -145,6 +151,7 @@ def _store(
     return RefreshingSSOAssertionStore(
         rows,
         refresher,
+        fresh_read=rows.fetch_fresh,
         coordinator_factory=coordinator_factory,  # pyright: ignore[reportArgumentType]  # test doubles stand in for the runtime factory
     )
 
@@ -408,6 +415,7 @@ async def test_a_failed_write_does_not_tell_the_user_to_sign_in_again():
     store = RefreshingSSOAssertionStore(
         rows,
         refresher,
+        fresh_read=rows.fetch_fresh,
         coordinator_factory=lambda: None,  # pyright: ignore[reportArgumentType]  # test double stands in for the runtime factory
     )
 
@@ -586,6 +594,22 @@ async def test_a_cross_replica_loser_whose_winner_renewed_reads_the_renewal_with
     assert len(coordinator.runs) == 1
 
 
+@pytest.mark.asyncio
+async def test_a_cross_replica_loser_rereads_past_a_stale_process_local_cache():
+    stale = _stored(_id_token(exp_offset=-1), expires_in=-1)
+    fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2")
+    rows = _FakeRows({"alice": fresh})
+    rows.cached_rows["alice"] = stale
+    transport = _FakeTransport(_renewal(_id_token()))
+    coordinator = _HeldCoordinator()
+
+    served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
+
+    assert served is fresh
+    assert transport.calls == []
+    assert len(coordinator.runs) == 1
+
+
 @pytest.mark.asyncio
 async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_in_challenge():
     rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
@@ -595,6 +619,7 @@ async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_i
     store = RefreshingSSOAssertionStore(
         rows,
         refresher,
+        fresh_read=rows.fetch_fresh,
         coordinator_factory=lambda: coordinator,  # pyright: ignore[reportArgumentType]  # test double stands in for the runtime factory
     )