diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py index 1c82fd27973..d8cf4ce83b3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_refresher.py +++ b/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: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py index f48a6fd8a35..08f19362503 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_refresher.py +++ b/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