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 d8cf4ce83b3..7600fd7ab8a 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 @@ -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, + ) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py index 42f64a9f90e..f7b92df5ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py +++ b/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, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py index 2d38ede8b27..7bdd02f30b7 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py +++ b/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 ), ) 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 08f19362503..d80913f8d33 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 @@ -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 )