fix(anthropic): refresh the WIF bearer token on the connection error retry (#50570)

* Refresh the WIF bearer token on the connection error retry

* Only refresh the bearer when the retry goes to the host the token was minted for

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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
 

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)

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: