diff --git a/litellm/llms/base_llm/auth/token_exchange.py b/litellm/llms/base_llm/auth/token_exchange.py index 5bf1f4dd0db..ce0e3c45639 100644 --- a/litellm/llms/base_llm/auth/token_exchange.py +++ b/litellm/llms/base_llm/auth/token_exchange.py @@ -9,6 +9,7 @@ public exception contract. import asyncio import hashlib +import hmac import json import re import threading @@ -538,10 +539,22 @@ class _Entry: """Single-flight state for one cache key; mutable by design, confined to the engine, and only ever mutated under the engine lock.""" - __slots__ = ("backoff_until", "done", "force_refresh", "in_flight", "last_error", "lifetime_seconds", "token") + __slots__ = ( + "backoff_until", + "done", + "force_refresh", + "in_flight", + "last_error", + "lifetime_seconds", + "previous_token", + "spec", + "token", + ) - def __init__(self, force_refresh: bool = False) -> None: + def __init__(self, spec: TokenExchangeSpec, force_refresh: bool = False) -> None: + self.spec: Final = spec self.token: MintedToken | None = None + self.previous_token: MintedToken | None = None self.lifetime_seconds: float | None = None self.in_flight: bool = False self.done: Final = threading.Event() @@ -562,7 +575,14 @@ class _Entry: self.in_flight = False self.done.set() + def issued(self, access_token: str) -> bool: + return any( + minted is not None and hmac.compare_digest(minted.access_token.get_secret_value(), access_token) + for minted in (self.token, self.previous_token) + ) + def _store(self, token: MintedToken, now: float) -> None: + self.previous_token = self.token self.token = token self.lifetime_seconds = None if token.expires_at is None else max(token.expires_at - now, 0.0) self.last_error = None @@ -684,11 +704,17 @@ class JwtBearerTokenExchangeEngine: async def aget_token(self, spec: TokenExchangeSpec) -> ExchangeResult: return await asyncio.to_thread(self.get_token, spec) + def spec_for_access_token(self, access_token: str) -> TokenExchangeSpec | None: + """Finds the identity a bearer was minted for, so resending an already built request can swap in + the current token instead of one that expired while the first attempt was in flight.""" + with self._lock: + return next((entry.spec for entry in self._entries.values() if entry.issued(access_token)), None) + def invalidate(self, spec: TokenExchangeSpec) -> None: key: Final = _cache_key(spec) with self._lock: if key in self._entries: - self._entries[key] = _Entry(force_refresh=True) + self._entries[key] = _Entry(spec, force_refresh=True) if self._shared_store is not None: self._shared_store.delete(key) @@ -699,7 +725,7 @@ class JwtBearerTokenExchangeEngine: return existing if len(self._entries) >= self._max_entries: self._evict_locked() - created: Final = _Entry() + created: Final = _Entry(spec) self._entries[key] = created return created diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3a6fa5f51da..0cfdbc74227 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -44,6 +44,8 @@ from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine +from litellm.llms.base_llm.auth.types import MintedToken from litellm.llms.custom_httpx.aiohttp_pool_diagnostics import build_pool_wait_trace_config from litellm.llms.custom_httpx.http2_connection_spread import AsyncHTTP2ConnectionSpread from litellm.types.llms.custom_http import * @@ -586,6 +588,28 @@ def _raise_masked_sync_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn raise MaskedHTTPStatusError(e, message=_text, text=_text) from None +def _same_origin(first: httpx.URL, second: httpx.URL) -> bool: + return (first.scheme, first.host, first.port) == (second.scheme, second.host, second.port) + + +async def _with_current_bearer( + headers: Mapping[str, str] | None, url: str, engine: JwtBearerTokenExchangeEngine +) -> Mapping[str, str] | None: + auth_key: Final = next((key for key in headers or () if key.lower() == "authorization"), None) + if headers is None or auth_key is None: + return headers + auth_value: Final = headers[auth_key] + if not auth_value.startswith("Bearer "): + return headers + spec: Final = engine.spec_for_access_token(auth_value.removeprefix("Bearer ")) + if spec is None or not _same_origin(httpx.URL(spec.token_url), httpx.URL(url)): + return headers + current: Final = await engine.aget_token(spec) + if not isinstance(current, MintedToken): + return headers + return MappingProxyType({**headers, auth_key: f"Bearer {current.access_token.get_secret_value()}"}) + + async def _raise_masked_async_error(e: httpx.HTTPStatusError, stream: bool) -> NoReturn: """Raise a MaskedHTTPStatusError for async HTTP handlers.""" if stream: @@ -639,7 +663,7 @@ class HTTPResponseLimitError(ValueError): class MaskedHTTPStatusError(httpx.HTTPStatusError): - def __init__(self, original_error, message: str | None = None, text: str | None = None): + def __init__(self, original_error: httpx.HTTPStatusError, message: str | None = None, text: str | None = None): # Create a new error with the masked URL masked_url: Final = mask_sensitive_info(str(original_error.request.url)) # Mask the original exception message too (it contains the full URL) @@ -702,12 +726,14 @@ class AsyncHTTPHandler: transport: httpx.AsyncBaseTransport | None = None, follow_redirects: bool = True, http2: bool = False, + token_exchange_engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, ): self.timeout = timeout self.event_hooks = event_hooks self.ssl_verify = ssl_verify self.shared_session = shared_session self.transport = transport + self.token_exchange_engine = token_exchange_engine self.follow_redirects = follow_redirects self.http2 = http2 self._owns_client = True @@ -1037,7 +1063,7 @@ class AsyncHTTPHandler: data=retry_data, json=json, params=params, - headers=headers, + headers=await _with_current_bearer(headers, url, self.token_exchange_engine), timeout=timeout, content=retry_content, ) @@ -1193,7 +1219,7 @@ class AsyncHTTPHandler: data=request_data, json=json, params=params, - headers=headers, + headers=await _with_current_bearer(headers, url, self.token_exchange_engine), content=request_content, ) response: Final = await client.send(req, stream=stream) diff --git a/tests/unit/llms/custom_httpx/test_http_handler.py b/tests/unit/llms/custom_httpx/test_http_handler.py index 3e583bf8dff..0de8efa2662 100644 --- a/tests/unit/llms/custom_httpx/test_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_http_handler.py @@ -26,6 +26,8 @@ from yarl import URL import litellm from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.custom_httpx.aiohttp_transport import LiteLLMAiohttpTransport +from litellm.llms.base_llm.auth.token_exchange import JwtBearerTokenExchangeEngine, default_token_exchange_engine +from litellm.llms.base_llm.auth.types import MintedToken, TokenExchangeSpec from litellm.llms.custom_httpx.http_handler import ( _CLIENT_REFCOUNT_WHEN_HANDLER_IS_SOLE_REFERRER, AsyncHTTPHandler, @@ -1513,9 +1515,14 @@ async def test_finalizer_on_live_loop_disposes_foreign_loop_session_without_sche class _RetryClientHandler(AsyncHTTPHandler): - def __init__(self, first: httpx.AsyncClient, retry: httpx.AsyncClient) -> None: + def __init__( + self, + first: httpx.AsyncClient, + retry: httpx.AsyncClient, + token_exchange_engine: JwtBearerTokenExchangeEngine = default_token_exchange_engine, + ) -> None: self._retry_client: Final = retry - super().__init__() + super().__init__(token_exchange_engine=token_exchange_engine) self.client = first def create_client( @@ -1552,6 +1559,135 @@ async def test_connection_error_retry_forwards_content(method: str): await handler.close() +class _WifClock: + def __init__(self) -> None: + self.now = 1_000.0 # mutable-ok: fake clock advanced by the test + + def __call__(self) -> float: + return self.now + + +class _CountingIssuer: + def __init__(self) -> None: + self.mints = 0 # mutable-ok: counts tokens minted so each one is distinct + + def post(self, url: str, *, content: bytes, headers: Mapping[str, str], timeout: float) -> httpx.Response: + self.mints += 1 + return httpx.Response( + 200, json={"access_token": f"wif-token-{self.mints}", "expires_in": 600, "token_type": "Bearer"} + ) + + +_WIF_SPEC: Final = TokenExchangeSpec( + token_url="https://api.anthropic.com/v1/oauth/token", + assertion_ref="oidc/env/ANTHROPIC_IDENTITY_TOKEN", + assertion_field="assertion", + static_body={"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", "federation_rule_id": "fdrl_1"}, + body_encoding="json", + request_headers={}, + cache_key_identity=("fdrl_1", "org-1", "", ""), + timeout_seconds=2.0, +) + + +def _minted_bearer(engine: JwtBearerTokenExchangeEngine) -> str: + minted: Final = engine.get_token(_WIF_SPEC) + assert isinstance(minted, MintedToken) + return minted.access_token.get_secret_value() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["post", "put", "patch", "delete"]) +@pytest.mark.parametrize("refreshed_by_another_request", [False, True]) +async def test_connection_error_retry_sends_current_wif_token(method: str, refreshed_by_another_request: bool) -> None: + clock: Final = _WifClock() + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), + assertion_reader=lambda ref: "identity-jwt", + clock=clock, + refresh_executor=ThreadPoolExecutor(max_workers=1), + ) + first_token: Final = _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def expire_then_drop(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + clock.now += 601 + if refreshed_by_another_request: + _minted_bearer(engine) + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(expire_then_drop)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await getattr(handler, method)( + "https://api.anthropic.com/v1/messages", + headers={"authorization": f"Bearer {first_token}", "x-trace": "kept"}, + content=b"{}", + ) + await handler.close() + + assert sent == [f"Bearer {first_token}", f"Bearer {_minted_bearer(engine)}"] + assert first_token != _minted_bearer(engine) + + +@pytest.mark.asyncio +async def test_connection_error_retry_keeps_bearer_the_engine_did_not_mint() -> None: + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), assertion_reader=lambda ref: "identity-jwt", clock=_WifClock() + ) + _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def raise_connection_error(request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(raise_connection_error)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await handler.post("https://api.anthropic.com/v1/messages", headers={"Authorization": "Bearer sk-static"}) + await handler.close() + + assert sent == ["Bearer sk-static"] + + +@pytest.mark.asyncio +async def test_connection_error_retry_keeps_wif_token_sent_to_another_host() -> None: + clock: Final = _WifClock() + engine: Final = JwtBearerTokenExchangeEngine( + poster=_CountingIssuer(), assertion_reader=lambda ref: "identity-jwt", clock=clock + ) + first_token: Final = _minted_bearer(engine) + sent: Final[list[str]] = [] # mutable-ok: async closure capture buffer + + async def expire_then_drop(request: httpx.Request) -> httpx.Response: + clock.now += 601 + raise httpx.ConnectError("connection dropped", request=request) + + async def capture_and_succeed(request: httpx.Request) -> httpx.Response: + sent.append(request.headers["authorization"]) + return httpx.Response(200, request=request) + + first: Final = httpx.AsyncClient(transport=httpx.MockTransport(expire_then_drop)) + retry: Final = httpx.AsyncClient(transport=httpx.MockTransport(capture_and_succeed)) + async with first, retry: + handler: Final = _RetryClientHandler(first=first, retry=retry, token_exchange_engine=engine) + await handler.post("https://elsewhere.example/v1/messages", headers={"authorization": f"Bearer {first_token}"}) + await handler.close() + + assert sent == [f"Bearer {first_token}"] + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["post", "patch"]) async def test_disconnect_after_send_is_not_replayed(method: str) -> None: