From 33afeedd47a109758d64495a882ad53f54399035 Mon Sep 17 00:00:00 2001 From: Varshith Puli Date: Thu, 15 Oct 2026 10:29:17 +0000 Subject: [PATCH 1/3] fix(proxy): serve stale JWKS keys while a background refresh runs (#46518) --- litellm/proxy/auth/handle_jwt.py | 136 ++++++++++++- tests/unit/proxy/auth/test_handle_jwt.py | 239 +++++++++++++++++++++-- 2 files changed, 358 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 9d800a1fd63..771ed59c1b6 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: @@ -826,11 +828,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 = 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 +860,91 @@ 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; failures keep the stale copy in use. + + 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. + """ + 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 = 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 # background refreshes never surface; staleness bounds model the failure + 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 +971,31 @@ 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 = 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 +1005,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 a0ea6ce0403..40f38775a06 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -4253,6 +4253,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.""" @@ -4371,8 +4407,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" @@ -4388,6 +4427,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") @@ -4594,6 +4636,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.""" @@ -4620,6 +4814,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] @@ -4635,8 +4830,14 @@ 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 request itself is still served from the trusted copy while the refresh runs off-path; what the + stale fallback never does is store the garbled body as the key set. + """ + import logging + from litellm.caching.dual_cache import DualCache jwks_url = "https://garbled-issuer.example.com/keys" @@ -4647,20 +4848,32 @@ 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,) - with pytest.raises(Exception, match="Error parsing response"): - await jwt_handler._get_public_key_from_jwks_url(jwks_url=jwks_url, kid="garbled-key") + 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}") == [jwk] + 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) @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 request itself is still served from the trusted copy while the refresh runs off-path. + """ + import logging + from litellm.caching.dual_cache import DualCache jwks_url = "https://erroring-issuer.example.com/keys" @@ -4679,11 +4892,14 @@ 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] + 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) @pytest.mark.asyncio @@ -4708,6 +4924,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 From 216072b73a72067c599ae936d095d26ea954cba0 Mon Sep 17 00:00:00 2001 From: Varshith Puli Date: Wed, 21 Oct 2026 16:45:37 +0000 Subject: [PATCH 2/3] style(proxy): apply ruff format to the JWKS stale-serve call The lint gate flagged handle_jwt.py: the _schedule_background_refresh call fits on one line at the configured width. --- litellm/proxy/auth/handle_jwt.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 771ed59c1b6..efb16b2b2e3 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -846,9 +846,7 @@ class JWTHandler: 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 - ) + 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()) From 700652d012e9b17ed3cc562342619afdeab1c41f Mon Sep 17 00:00:00 2001 From: Varshith Puli Date: Thu, 22 Oct 2026 00:12:32 +0000 Subject: [PATCH 3/3] fix(proxy): fail closed when a reachable JWKS endpoint answers badly The background refresh swallowed every non-transport failure, so a 503 or garbled JWKS body kept the stale key set trusted for ttl + stale_ttl instead of surfacing the error. Only an unreachable IdP should keep the stale copy in use. A non-transport refresh failure now invalidates the stale copy, so the next request takes the blocking path and raises as before. The garbled-body and 503 regression tests now pin this: the in-flight request is still served from the trusted copy, the stale copy is dropped, and the next request surfaces the IdP error. Also annotate the three servable bindings explicitly, clearing the 8 new reportUnknownVariableType diagnostics the branch added. Signed-off-by: Varshith Puli --- litellm/proxy/auth/handle_jwt.py | 29 +++++++++++++++++++----- tests/unit/proxy/auth/test_handle_jwt.py | 21 +++++++++++++---- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index efb16b2b2e3..833f015cb7f 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -821,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, @@ -840,7 +844,9 @@ class JWTHandler: if cached is not None: return cached - servable: Final = await self._get_servable_stale_copy(cache_key=cache_key, ttl=ttl) + 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( @@ -884,10 +890,13 @@ class JWTHandler: refresh: Callable[[], Awaitable[_CachedValueT]], log_context: str, ) -> None: - """Refresh ``cache_key`` off the request path; failures keep the stale copy in use. + """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. + 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()) @@ -900,7 +909,9 @@ class JWTHandler: cache_key=cache_key, ttl=ttl, stale_ttl=stale_ttl, refresh=refresh ) if isinstance(outcome, JWKSUnreachableError): - servable: Final = await self._get_servable_stale_copy(cache_key=cache_key, ttl=ttl) + 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", @@ -917,7 +928,11 @@ class JWTHandler: outcome=outcome, age=age, ) - except Exception as e: # noqa: BLE001 # background refreshes never surface; staleness bounds model the failure + 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 ) @@ -969,7 +984,9 @@ class JWTHandler: if stale_ttl <= 0: raise outcome - servable: Final = await self._get_servable_stale_copy(cache_key=cache_key, ttl=ttl) + 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 diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 40f38775a06..e1624283d40 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -4833,8 +4833,9 @@ async def test_stale_fallback_warns_with_the_kid_and_how_stale_the_jwks_copy_is( 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 request itself is still served from the trusted copy while the refresh runs off-path; what the - stale fallback never does is store the garbled body as the key set. + 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 @@ -4861,16 +4862,22 @@ async def test_unparseable_jwks_response_never_poisons_the_cached_keys(caplog): 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 "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(caplog): """An IdP error body must never be stored as the key set, least of all as the stale copy. - The request itself is still served from the trusted copy while the refresh runs off-path. + 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 @@ -4897,10 +4904,14 @@ async def test_jwks_error_response_is_not_cached_over_the_last_known_good_keys(c 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 async def test_sustained_jwks_outage_refetches_once_per_backoff_window_not_once_per_request():