fix(mcp): answer a cross-replica loser retryable instead of re-electing it

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

@@ -29,12 +29,12 @@ so the reader's own guard challenges the user to sign in again, while a transien
 
 One ambiguity remains under Redis-coordinated renewal across replicas. A cross-replica loser that
 finds the row still expiring after the holder finished cannot tell a refused refresh from a renewal
-that could not be recorded, so it runs its own election instead of guessing a sign-in challenge. Its
-own redemption answers definitively: a dead token is refused again and challenges, a recovered store
-records the renewal, and a still-failing store answers 503. A read that loses twice answers 503 so it
-is retried rather than told to sign in. On the refusal path, each cross-replica loser costs one extra
-token-endpoint call for an already-dead token, the same call the next uncontended request would make
-anyway.
+that could not be recorded. Redeeming itself could consume a refresh token the holder may already
+have rotated, so it answers retryable 503 rather than guessing a sign-in challenge. The next
+uncontended read settles the outcome itself: a refusal challenges, and a successful refresh persists.
+If the holder rotated the token but its write failed, that rotation is lost and the next uncontended
+read's refusal challenges, which is the only honest answer because the rotated token was never
+recorded. On the refusal path, the loser pays for one retry before that challenge.
 """
 
 from __future__ import annotations
@@ -92,12 +92,6 @@ _SINGLE_FLIGHT_KEY: Final = "sso_identity_assertion"
 # Renew this far ahead of ``exp`` so a token that would die between resolution and the second leg of
 # the exchange is replaced first. Matches the sibling per-user token store's skew.
 _DEFAULT_EXPIRY_SKEW_SECONDS: Final = 60.0
-# Elections one read sits through before it stops waiting on other replicas and answers 503.
-_MAX_ELECTIONS: Final = 2
-
-
-class _ElectionLost(Exception):
-    """A cross-replica loser re-read the row after the holder finished and found it still expiring."""
 
 
 class AssertionRead(Protocol):
@@ -381,8 +375,11 @@ class RefreshingSSOAssertionStore:
     A refusal leaves the expired assertion in place for the reader's own guard to reject, so the user
     sees the same sign-in-again challenge as before this store existed. A transient IdP failure
     raises ``AssertionStoreUnavailable``, the protocol's existing signal for "this is not the user's
-    fault"; concurrent in-process callers share that outcome, while a cross-replica loser re-elects
-    when its re-read still finds the row expiring.
+    fault"; concurrent in-process callers share that outcome, while a cross-replica loser answers 503
+    when its re-read still finds the row expiring. On the refusal path that costs the loser one retry,
+    which then challenges. If the holder rotated the token but its write failed, the rotation is lost
+    and the next uncontended read's refusal challenges, the only honest answer because that token was
+    never recorded.
     """
 
     def __init__(
@@ -406,20 +403,13 @@ class RefreshingSSOAssertionStore:
         assertion: Final = await self._inner.fetch(user_id)
         if not self._expiring(assertion):
             return assertion
-        for _ in range(_MAX_ELECTIONS):
-            try:
-                await self._coordinator().run(
-                    user_id,
-                    _SINGLE_FLIGHT_KEY,
-                    refresh=lambda: self._renew(user_id),
-                    reread=lambda: self._settled_or_lost(user_id),
-                )
-            except _ElectionLost:
-                continue
-            return await self._inner.fetch(user_id)
-        raise AssertionStoreUnavailable(
-            f"the IdP identity assertion for user_id={user_id} is being renewed by other replicas; retry shortly"
+        await self._coordinator().run(
+            user_id,
+            _SINGLE_FLIGHT_KEY,
+            refresh=lambda: self._renew(user_id),
+            reread=lambda: self._reread_renewed(user_id),
         )
+        return await self._inner.fetch(user_id)
 
     def _expiring(self, assertion: SSOIdentityAssertion | None) -> bool:
         return assertion is not None and assertion_expired(assertion, self._clock() + self._skew)
@@ -452,11 +442,18 @@ class RefreshingSSOAssertionStore:
                         raise AssertionStoreUnavailable(failure.detail)
                 assert_never(failure.kind)
 
-    async def _settled_or_lost(self, user_id: str) -> None:
-        """The cross-replica loser's re-read. A still-expiring row means the holder's renewal was refused
-        or could not be recorded, and this replica cannot tell which, so it runs its own election."""
-        if self._expiring(await self._inner.fetch(user_id)):
-            raise _ElectionLost
+    async def _reread_renewed(self, user_id: str) -> None:
+        """A loser cannot distinguish refusal from an unrecorded renewal without risking token replay.
+
+        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)
+        if self._expiring(latest):
+            raise AssertionStoreUnavailable(
+                f"the IdP identity assertion for user_id={user_id} was being renewed by another replica "
+                "and is not yet current; retry shortly"
+            )
 
 
 def default_sso_assertion_store() -> SSOAssertionStore:

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

@@ -537,11 +537,10 @@ class _ReplaceThenRefreshCoordinator:
         return await refresh()
 
 
-class _HeldThenWonCoordinator:
-    """Emulates a cross-replica holder finishing before a later election wins."""
+class _HeldCoordinator:
+    """Emulates a cross-replica holder finishing before the loser re-reads."""
 
-    def __init__(self, lost_elections: int, before_reread: Callable[[], None] | None = None) -> None:
-        self._lost_elections = lost_elections
+    def __init__(self, before_reread: Callable[[], None] | None = None) -> None:
         self._before_reread = before_reread
         self.runs: list[tuple[str, str]] = []
 
@@ -553,11 +552,9 @@ class _HeldThenWonCoordinator:
         reread: Callable[[], Awaitable[None]],
     ) -> None:
         self.runs.append((user_id, server_id))
-        if len(self.runs) <= self._lost_elections:
-            if self._before_reread is not None:
-                self._before_reread()
-            return await reread()
-        return await refresh()
+        if self._before_reread is not None:
+            self._before_reread()
+        return await reread()
 
 
 @pytest.mark.asyncio
@@ -580,7 +577,7 @@ async def test_a_cross_replica_loser_whose_winner_renewed_reads_the_renewal_with
     fresh = _stored(_id_token(), expires_in=3600, refresh_token="rt_2")
     rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
     transport = _FakeTransport(_renewal(_id_token()))
-    coordinator = _HeldThenWonCoordinator(lost_elections=1, before_reread=lambda: rows.rows.__setitem__("alice", fresh))
+    coordinator = _HeldCoordinator(before_reread=lambda: rows.rows.__setitem__("alice", fresh))
 
     served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
 
@@ -594,7 +591,7 @@ async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_i
     rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
     transport = _FakeTransport(_renewal(_id_token()))
     refresher = SSOAssertionRefresher(transport, client_config=lambda: _CLIENT, read=rows.fetch, write=_explode)
-    coordinator = _HeldThenWonCoordinator(lost_elections=1)
+    coordinator = _HeldCoordinator()
     store = RefreshingSSOAssertionStore(
         rows,
         refresher,
@@ -604,35 +601,21 @@ async def test_a_cross_replica_loser_does_not_turn_a_write_failure_into_a_sign_i
     with pytest.raises(AssertionStoreUnavailable):
         await store.fetch("alice")
 
-    assert len(transport.calls) == 1
-    assert len(coordinator.runs) == 2
+    assert transport.calls == []
+    assert len(coordinator.runs) == 1
 
 
 @pytest.mark.asyncio
-async def test_a_cross_replica_loser_whose_winner_was_refused_is_refused_too_and_challenges():
-    stale = _id_token(exp_offset=-1)
-    rows = _FakeRows({"alice": _stored(stale, expires_in=-1)})
-    transport = _FakeTransport(Error(RefreshFailure.of_rejected("dead")))
-    coordinator = _HeldThenWonCoordinator(lost_elections=1)
-
-    served = await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
-
-    assert served is not None
-    assert served.id_token.get_secret_value() == stale
-    assert len(transport.calls) == 1
-
-
-@pytest.mark.asyncio
-async def test_a_read_that_loses_every_election_answers_unavailable_not_sign_in_again():
+async def test_a_cross_replica_loser_never_redeems_the_token_the_holder_may_have_rotated():
     rows = _FakeRows({"alice": _stored(_id_token(exp_offset=-1), expires_in=-1)})
-    transport = _FakeTransport(_renewal(_id_token()))
-    coordinator = _HeldThenWonCoordinator(lost_elections=10)
+    transport = _FakeTransport(Error(RefreshFailure.of_rejected("dead")))
+    coordinator = _HeldCoordinator()
 
     with pytest.raises(AssertionStoreUnavailable):
         await _store(rows, transport, coordinator_factory=lambda: coordinator).fetch("alice")
 
     assert transport.calls == []
-    assert len(coordinator.runs) == 2
+    assert len(coordinator.runs) == 1
 
 
 @pytest.mark.asyncio