diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9e11179c686..cd50c7ee91f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -279,6 +279,8 @@ class JWTHandler: self.leeway = 0 # Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request. self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url + # Per-cache-key background refresh tasks so concurrent requests share one off-path refresh. + self._refresh_tasks: dict[str, asyncio.Task[None]] = {} # mutable-ok: task registry, keyed by JWKS url self.agent_lookup: AgentLookup = _NoRegisteredAgents() def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None: @@ -819,6 +821,10 @@ class JWTHandler: async def _put_cached_value(self, cache_key: str, value: JWKKeyValue | str | float, ttl: float) -> None: await self.user_api_key_cache.async_set_cache(key=cache_key, value=value, ttl=ttl) + @with_service_target(AUTH_OBJECTS_TARGET) + async def _delete_cached_value(self, cache_key: str) -> None: + await self.user_api_key_cache.async_delete_cache(key=cache_key) + async def _cached_with_stale_fallback( self, cache_key: str, @@ -826,11 +832,29 @@ class JWTHandler: refresh: Callable[[], Awaitable[_CachedValueT]], log_context: str, ) -> _CachedValueT: - """Read ``cache_key``, refreshing it through a single-flight lock on a miss.""" + """Read ``cache_key``, refreshing in the background while a trusted copy exists. + + On a fresh-cache miss the last-known-good copy is served immediately and one background refresh + is scheduled for the expiry window, so a slow identity provider cannot stall every concurrent + request behind the single-flight lock. Callers block on the lock only when no trusted copy exists + at all (cold start) or the last-known-good copy has aged past ``public_key_ttl`` + + ``public_key_stale_ttl``. + """ cached: Final[_CachedValueT | None] = await self._get_cached_value(cache_key) if cached is not None: return cached + servable: Final[tuple[_CachedValueT, float] | None] = await self._get_servable_stale_copy( + cache_key=cache_key, ttl=ttl + ) + if servable is not None: + stale, _ = servable + verbose_proxy_logger.debug( + "JWT Auth: serving a stale JWKS copy for %s while a background refresh runs", cache_key + ) + self._schedule_background_refresh(cache_key=cache_key, ttl=ttl, refresh=refresh, log_context=log_context) + return stale + lock: Final = self._refresh_locks.setdefault(cache_key, asyncio.Lock()) async with lock: cached_after_lock: Final[_CachedValueT | None] = await self._get_cached_value(cache_key) @@ -840,6 +864,100 @@ class JWTHandler: cache_key=cache_key, ttl=ttl, refresh=refresh, log_context=log_context ) + def _schedule_background_refresh( + self, + cache_key: str, + ttl: float, + refresh: Callable[[], Awaitable[_CachedValueT]], + log_context: str, + ) -> None: + """Schedule one background refresh for ``cache_key``, skipping when one is already in flight. + + The check and the schedule run with no await between them, so concurrent requests on one event + loop cannot schedule the same window twice. + """ + in_flight: Final = self._refresh_tasks.get(cache_key) + if in_flight is not None and not in_flight.done(): + return + self._refresh_tasks[cache_key] = asyncio.get_running_loop().create_task( + self._background_refresh(cache_key=cache_key, ttl=ttl, refresh=refresh, log_context=log_context) + ) + + async def _background_refresh( + self, + cache_key: str, + ttl: float, + refresh: Callable[[], Awaitable[_CachedValueT]], + log_context: str, + ) -> None: + """Refresh ``cache_key`` off the request path. + + Takes the single-flight lock so a concurrent cold-start caller cannot start a second fetch, and + reuses the unreachable backoff so a sustained outage still costs one fetch per window. An + unreachable IdP keeps the bounded stale copy in use; a reachable IdP that answers badly + invalidates it, so the next request takes the blocking path and surfaces the error instead of + trusting keys the IdP may have rotated away. + """ + try: + lock: Final = self._refresh_locks.setdefault(cache_key, asyncio.Lock()) + async with lock: + cached: Final[_CachedValueT | None] = await self._get_cached_value(cache_key) + if cached is not None: + return + stale_ttl: Final = self._get_public_key_stale_ttl() + outcome: Final = await self._refresh_or_record_outage( + cache_key=cache_key, ttl=ttl, stale_ttl=stale_ttl, refresh=refresh + ) + if isinstance(outcome, JWKSUnreachableError): + servable: Final[tuple[_CachedValueT, float] | None] = await self._get_servable_stale_copy( + cache_key=cache_key, ttl=ttl + ) + if servable is None: + verbose_proxy_logger.warning( + "JWT Auth: background JWKS refresh for %s failed and no trusted copy remains: %s", + log_context, + outcome, + ) + else: + _, age = servable + self._warn_stale_copy_in_use( + cache_key=cache_key, + ttl=ttl, + stale_ttl=stale_ttl, + log_context=log_context, + outcome=outcome, + age=age, + ) + except Exception as e: # noqa: BLE001 # only non-transport failures reach here; unreachable is an outcome value + # A reachable IdP that answers badly must surface its error, not keep the stale copy + # trusted: drop it so the next request takes the blocking refresh path and raises. + await self._delete_cached_value(f"{STALE_CACHE_KEY_PREFIX}{cache_key}") + await self._delete_cached_value(f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{cache_key}") + verbose_proxy_logger.warning( + "JWT Auth: background JWKS refresh for %s raised unexpectedly: %s", log_context, e + ) + finally: + self._refresh_tasks.pop(cache_key, None) + + async def _get_servable_stale_copy(self, cache_key: str, ttl: float) -> tuple[_CachedValueT, float] | None: + """Return the last-known-good copy and its age when it is still trusted, else None. + + Trust is bounded by ``public_key_ttl`` + ``public_key_stale_ttl`` measured from when the copy was + taken, enforced on every read rather than baked into the cache entry's own expiry. An operator who + lowers ``public_key_stale_ttl``, or sets it to 0 to fail closed, is usually doing it mid-incident, + and a copy written under the old longer setting would otherwise stay servable until it aged out on + its own. A copy whose write time cannot be established is not servable, so the bound cannot be + dodged by losing the timestamp. + """ + stale_ttl: Final = self._get_public_key_stale_ttl() + if stale_ttl <= 0: + return None + stale: Final[_CachedValueT | None] = await self._get_cached_value(f"{STALE_CACHE_KEY_PREFIX}{cache_key}") + age: Final = await self._stale_copy_age(cache_key) + if stale is None or age is None or age > ttl + stale_ttl: + return None + return (stale, age) + async def _refresh_or_serve_stale( self, cache_key: str, @@ -866,11 +984,33 @@ class JWTHandler: if stale_ttl <= 0: raise outcome - stale: Final[_CachedValueT | None] = await self._get_cached_value(f"{STALE_CACHE_KEY_PREFIX}{cache_key}") - age: Final = await self._stale_copy_age(cache_key) - lifetime: Final = ttl + stale_ttl - if stale is None or age is None or age > lifetime: + servable: Final[tuple[_CachedValueT, float] | None] = await self._get_servable_stale_copy( + cache_key=cache_key, ttl=ttl + ) + if servable is None: raise outcome + stale, age = servable + self._warn_stale_copy_in_use( + cache_key=cache_key, + ttl=ttl, + stale_ttl=stale_ttl, + log_context=log_context, + outcome=outcome, + age=age, + ) + return stale + + def _warn_stale_copy_in_use( + self, + cache_key: str, + ttl: float, + stale_ttl: float, + log_context: str, + outcome: JWKSUnreachableError, + age: float, + ) -> None: + """Log that requests are being authenticated against the stale copy because the refresh failed.""" + lifetime: Final = ttl + stale_ttl verbose_proxy_logger.warning( "JWT Auth: identity provider unreachable, authenticating %s against a stale JWKS copy of %s " "(last refreshed %.0fs ago, stops being trusted in %.0fs). Refresh failed: %s", @@ -880,7 +1020,6 @@ class JWTHandler: max(lifetime - age, 0), outcome, ) - return stale async def _stale_copy_age(self, cache_key: str) -> float | None: written_at: Final = await self._get_cached_timestamp(f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{cache_key}") diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 4c223edc20d..5f0bb268137 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -4278,6 +4278,42 @@ def _get_jwt_handler_with_scripted_endpoint( return jwt_handler +class _GatedJWKSEndpoint: + """Injected ``http_handler`` whose JWKS response waits on a test-controlled event. + + Lets a test prove requests finish while the refresh is still in flight, with no timing involved. + """ + + def __init__(self, body: Mapping[str, object], gate: asyncio.Event) -> None: + self.body = body + self.gate = gate + self.call_count = 0 + + async def get( + self, + url: str, + params: Mapping[str, str] | None = None, + headers: Mapping[str, str] | None = None, + ) -> MagicMock: + self.call_count += 1 + await self.gate.wait() + response = MagicMock() + response.status_code = 200 + response.json.return_value = self.body + return response + + +async def _await_background_jwks_refresh(jwt_handler: JWTHandler, cache_key: str) -> None: + """Wait out the background refresh scheduled for ``cache_key``, if one is still running. + + The schedule happens synchronously inside the request that triggered it, so once that request has + returned the task is either registered or already finished, never still to be created. + """ + pending = jwt_handler._refresh_tasks.get(cache_key) + if pending is not None: + await pending + + @pytest.mark.asyncio async def test_get_public_key_retries_transient_jwks_fetch_failure(): """A single connect timeout to the IdP must be retried, not surfaced to the caller.""" @@ -4397,8 +4433,11 @@ async def test_removed_signing_key_stops_being_trusted_once_the_stale_window_exp @pytest.mark.asyncio -async def test_key_removed_from_a_reachable_jwks_is_rejected_without_consulting_the_stale_copy(): - """A reachable IdP always wins: dropping a key revokes it immediately, stale copy included.""" +async def test_key_removed_from_a_reachable_jwks_is_rejected_once_the_triggered_refresh_lands(): + """A reachable IdP always wins, just not mid-flight: the expiry triggers a refresh, and once that refresh + lands the dropped key is rejected and the stale copy is replaced. Requests served while the refresh runs + still see the last trusted copy; the revocation delay is bounded by one fetch, never by the stale window. + """ from litellm.caching.dual_cache import DualCache jwks_url = "https://rotating-issuer.example.com/keys" @@ -4414,6 +4453,9 @@ async def test_key_removed_from_a_reachable_jwks_is_rejected_without_consulting_ await cache.async_delete_cache(key=active_cache_key) endpoint.outcomes = ({"keys": [current_jwk]},) + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key") == retired_jwk + await _await_background_jwks_refresh(jwt_handler, active_cache_key) + with pytest.raises(NoMatchingJWTPublicKeyError): await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="retired-key") @@ -4620,6 +4662,158 @@ async def test_public_key_stale_ttl_defaults_to_one_hour(): assert LiteLLM_JWTAuth().public_key_stale_ttl == 3600 +@pytest.mark.asyncio +async def test_expiry_serves_stale_keys_without_waiting_for_the_slow_refresh(): + """A slow IdP must not stall JWT-authenticated traffic when `public_key_ttl` lapses. + + Regression test: on a fresh-cache miss every caller used to queue behind the single-flight lock while the + first caller ran the live fetch. The fetch below never completes during the burst, so any caller that + waited for it would hang the gather. + """ + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://slow-idp.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="slow-idp-key") + cache = DualCache() + gate = asyncio.Event() + gate.set() + endpoint = _GatedJWKSEndpoint({"keys": [jwk]}, gate) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + active_key = f"litellm_jwt_auth_keys_{jwks_url}" + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="slow-idp-key") == jwk + + await cache.async_delete_cache(key=active_key) + gate.clear() + results = await asyncio.wait_for( + asyncio.gather( + *( + jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="slow-idp-key") + for _ in range(10) + ) + ), + timeout=5, + ) + + assert results == [jwk] * 10 + gate.set() + await _await_background_jwks_refresh(jwt_handler, active_key) + assert endpoint.call_count == 2 + assert await cache.async_get_cache(key=active_key) == [jwk] + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="slow-idp-key") == jwk + assert endpoint.call_count == 2 + + +@pytest.mark.asyncio +async def test_each_expiry_window_triggers_exactly_one_background_refresh(): + """Concurrent requests in one TTL window must share a single background refresh, not one per request.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://single-flight.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="single-flight-key") + cache = DualCache() + gate = asyncio.Event() + gate.set() + endpoint = _GatedJWKSEndpoint({"keys": [jwk]}, gate) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + active_key = f"litellm_jwt_auth_keys_{jwks_url}" + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="single-flight-key") == jwk + + for _ in range(3): + await cache.async_delete_cache(key=active_key) + results = await asyncio.wait_for( + asyncio.gather( + *( + jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="single-flight-key") + for _ in range(10) + ) + ), + timeout=5, + ) + assert results == [jwk] * 10 + await _await_background_jwks_refresh(jwt_handler, active_key) + + assert endpoint.call_count == 4 + + +@pytest.mark.asyncio +async def test_cold_start_blocks_concurrent_requests_until_keys_arrive(): + """With no cached copy at all there is nothing to serve, so callers wait on the single-flight refresh.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://cold-start.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="cold-start-key") + cache = DualCache() + gate = asyncio.Event() + endpoint = _GatedJWKSEndpoint({"keys": [jwk]}, gate) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + + pending = [ + asyncio.create_task(jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="cold-start-key")) + for _ in range(5) + ] + for _ in range(10): + await asyncio.sleep(0) + assert all(not task.done() for task in pending) + + gate.set() + assert await asyncio.gather(*pending) == [jwk] * 5 + assert endpoint.call_count == 1 + + +@pytest.mark.asyncio +async def test_failed_background_refresh_keeps_serving_stale_keys(): + """A refresh that fails off-path must not poison later requests; the stale copy keeps serving.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://flapping-idp.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="flapping-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint(({"keys": [jwk]},)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint(cache, endpoint) + active_key = f"litellm_jwt_auth_keys_{jwks_url}" + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flapping-key") == jwk + + await cache.async_delete_cache(key=active_key) + endpoint.outcomes = (httpx.ConnectTimeout("connect timed out"),) + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flapping-key") == jwk + await _await_background_jwks_refresh(jwt_handler, active_key) + assert endpoint.call_count == 1 + JWKS_FETCH_ATTEMPTS + + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="flapping-key") == jwk + await _await_background_jwks_refresh(jwt_handler, active_key) + assert endpoint.call_count == 1 + JWKS_FETCH_ATTEMPTS + + +@pytest.mark.asyncio +async def test_stale_copy_past_the_staleness_cap_is_not_served_off_path(): + """The staleness bound is enforced before anything is served: an aged-out copy still fails closed.""" + from litellm.caching.dual_cache import DualCache + + jwks_url = "https://aged-out.example.com/keys" + _, jwk = _get_rsa_key_and_jwk(kid="aged-out-key") + cache = DualCache() + endpoint = _ScriptedJWKSEndpoint((httpx.ConnectTimeout("connect timed out"),)) + jwt_handler = _get_jwt_handler_with_scripted_endpoint( + cache, endpoint, public_key_ttl=600, public_key_stale_ttl=3600 + ) + active_key = f"litellm_jwt_auth_keys_{jwks_url}" + + await cache.async_set_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_key}", value=[jwk], ttl=4200) + await cache.async_set_cache( + key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_key}", + value=time.time() - 4201, + ttl=4200, + ) + + with pytest.raises(JWKSUnreachableError): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="aged-out-key") + assert endpoint.call_count == JWKS_FETCH_ATTEMPTS + + @pytest.mark.asyncio async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is(caplog): """Serving an unrefreshed signing key is a security-relevant event, so it must be legible in the logs.""" @@ -4646,6 +4840,7 @@ async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is( caplog.set_level(logging.WARNING) await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="warned-key") + await _await_background_jwks_refresh(jwt_handler, active_cache_key) warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] stale_warnings = [m for m in warnings if "stale JWKS copy" in m] @@ -4661,8 +4856,15 @@ async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is( @pytest.mark.asyncio -async def test_unparseable_jwks_response_does_not_fall_back_to_the_stale_copy(): - """Only an unreachable IdP unlocks the stale copy. A reachable one that answers badly must surface the error.""" +async def test_unparseable_jwks_response_never_poisons_the_cached_keys(caplog): + """A reachable IdP that answers badly must not poison the cache, and the failure must be loud. + + The in-flight request is still served from the trusted copy while the refresh runs off-path, but the + garbled body invalidates the stale copy, so the next request takes the blocking path and the parse + error surfaces instead of silently trusting keys the IdP may have rotated away. + """ + import logging + from litellm.caching.dual_cache import DualCache jwks_url = "https://garbled-issuer.example.com/keys" @@ -4673,20 +4875,38 @@ async def test_unparseable_jwks_response_does_not_fall_back_to_the_stale_copy(): assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") == jwk - await cache.async_delete_cache(key=f"litellm_jwt_auth_keys_{jwks_url}") + active_cache_key = f"litellm_jwt_auth_keys_{jwks_url}" + await cache.async_delete_cache(key=active_cache_key) garbled = MagicMock() garbled.status_code = 200 garbled.text = "not json" garbled.json.side_effect = ValueError("Expecting value: line 1 column 1") endpoint.outcomes = (garbled,) + caplog.set_level(logging.WARNING) + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") == jwk + await _await_background_jwks_refresh(jwt_handler, active_cache_key) + + assert await cache.async_get_cache(key=active_cache_key) is None + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") is None + assert await cache.async_get_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}") is None + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("raised unexpectedly" in m and "garbled-key" in m for m in warnings) + with pytest.raises(Exception, match="Error parsing response"): await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") @pytest.mark.asyncio -async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys(): - """An IdP error body must never be stored as the key set, least of all as the stale copy.""" +async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys(caplog): + """An IdP error body must never be stored as the key set, least of all as the stale copy. + + The in-flight request is still served from the trusted copy while the refresh runs off-path, but the + error invalidates the stale copy, so the next request takes the blocking path and the 503 surfaces + instead of silently trusting keys the IdP may have rotated away. + """ + import logging + from litellm.caching.dual_cache import DualCache jwks_url = "https://erroring-issuer.example.com/keys" @@ -4705,11 +4925,18 @@ async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys() server_error.json.return_value = {"error": "upstream unavailable"} endpoint.outcomes = (server_error,) - with pytest.raises(Exception, match="returned status 503"): - await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") + caplog.set_level(logging.WARNING) + assert await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") == jwk + await _await_background_jwks_refresh(jwt_handler, active_cache_key) assert await cache.async_get_cache(key=active_cache_key) is None - assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") == [jwk] + assert await cache.async_get_cache(key=f"{STALE_CACHE_KEY_PREFIX}{active_cache_key}") is None + assert await cache.async_get_cache(key=f"{STALE_WRITTEN_AT_CACHE_KEY_PREFIX}{active_cache_key}") is None + warnings = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert any("raised unexpectedly" in m and "erroring-key" in m for m in warnings) + + with pytest.raises(Exception, match="returned status 503"): + await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="erroring-key") @pytest.mark.asyncio @@ -4734,6 +4961,7 @@ async def test_sustained_jwks_outage_refetches_once_per_backoff_window_not_once_ ) assert public_keys == [jwk] * 6 + await _await_background_jwks_refresh(jwt_handler, f"litellm_jwt_auth_keys_{jwks_url}") assert endpoint.call_count - calls_before_outage == JWKS_FETCH_ATTEMPTS