mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(mcp): remember the rotated M2M bearer for later requests in the session
The auth object is the httpx client's auth for the whole MCP session; after a 401 recovery it kept sending the rejected token first, burning a 401 round trip and the single retry on every subsequent call
This commit is contained in:
parent
5b64239afc
commit
8bfd8baab8
2 changed files with 22 additions and 0 deletions
|
|
@ -340,6 +340,7 @@ class ClientCredentialsBearerAuth(httpx.Auth):
|
|||
fresh = await self._refetch(token)
|
||||
if fresh is None:
|
||||
return
|
||||
self._access_token = SecretStr(fresh)
|
||||
request.headers[self.header_name] = f"Bearer {fresh}"
|
||||
yield request
|
||||
|
||||
|
|
|
|||
|
|
@ -340,6 +340,27 @@ async def test_bearer_auth_retries_a_401_once_with_a_fresh_token():
|
|||
assert seen == ["Bearer stale-token", "Bearer fresh-token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bearer_auth_remembers_the_rotated_token_for_later_requests():
|
||||
# The auth object lives for the whole MCP session (it is the httpx client's auth), so after a
|
||||
# 401 recovery it must send the fresh token first on subsequent requests; re-sending the
|
||||
# rejected one would burn a 401 round trip and the single retry on every call.
|
||||
transport, seen = _upstream([httpx.Response(401), httpx.Response(200), httpx.Response(200)])
|
||||
refetched: "list[str]" = []
|
||||
|
||||
async def refetch(failed: str) -> "str | None":
|
||||
refetched.append(failed)
|
||||
return "fresh-token"
|
||||
|
||||
auth = ClientCredentialsBearerAuth("stale-token", refetch)
|
||||
async with httpx.AsyncClient(transport=transport, auth=auth) as client:
|
||||
first = await client.get("https://upstream.example.com/mcp")
|
||||
second = await client.get("https://upstream.example.com/mcp")
|
||||
assert first.status_code == 200 and second.status_code == 200
|
||||
assert refetched == ["stale-token"]
|
||||
assert seen == ["Bearer stale-token", "Bearer fresh-token", "Bearer fresh-token"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bearer_auth_surfaces_the_401_when_the_refetch_fails():
|
||||
transport, seen = _upstream([httpx.Response(401)])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue