mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(mcp): token refresh and M2M egress honor the admin-entered token URL
This commit is contained in:
parent
465ebb1bdd
commit
ce8d6f7d25
8 changed files with 104 additions and 12 deletions
|
|
@ -1600,7 +1600,7 @@ async def refresh_user_oauth_token(
|
|||
) -> OAuthCredentialPayload | None:
|
||||
"""Attempt to refresh a per-user OAuth2 token using its stored refresh_token.
|
||||
|
||||
POSTs to ``server.token_url`` with ``grant_type=refresh_token``.
|
||||
POSTs to ``server.effective_token_url`` with ``grant_type=refresh_token``.
|
||||
|
||||
On success: persists the new credential via ``store_user_oauth_credential``
|
||||
and returns the updated payload dict.
|
||||
|
|
@ -1609,7 +1609,7 @@ async def refresh_user_oauth_token(
|
|||
stale credential and triggering re-authentication.
|
||||
"""
|
||||
refresh_token: Final[str | None] = cred.get("refresh_token")
|
||||
token_url: Final[str | None] = getattr(server, "token_url", None)
|
||||
token_url: Final[str | None] = getattr(server, "effective_token_url", None) or getattr(server, "token_url", None)
|
||||
server_id: Final[str] = getattr(server, "server_id", "")
|
||||
client_id: Final[str | None] = getattr(server, "client_id", None)
|
||||
client_secret: Final[str | None] = getattr(server, "client_secret", None)
|
||||
|
|
|
|||
|
|
@ -67,7 +67,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
rest of the identity rather than stored in a key."""
|
||||
material: Final = "\x00".join(
|
||||
(
|
||||
server.token_url or "",
|
||||
server.effective_token_url or "",
|
||||
server.client_id or "",
|
||||
server.client_secret or "",
|
||||
" ".join(server.scopes or ()),
|
||||
|
|
@ -82,7 +82,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
|
||||
@staticmethod
|
||||
def _has_client_credentials_config(server: "MCPServer") -> bool:
|
||||
return bool(server.client_id and server.client_secret and server.token_url)
|
||||
return bool(server.client_id and server.client_secret and server.effective_token_url)
|
||||
|
||||
async def async_get_token(self, server: "MCPServer") -> str | None:
|
||||
"""Return a valid access token, fetching or refreshing as needed.
|
||||
|
|
@ -112,19 +112,20 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
return token
|
||||
|
||||
async def _fetch_token(self, server: "MCPServer") -> tuple[str, int]:
|
||||
"""POST to ``token_url`` with ``grant_type=client_credentials``.
|
||||
"""POST to ``effective_token_url`` with ``grant_type=client_credentials``.
|
||||
|
||||
Returns ``(access_token, ttl_seconds)`` where ttl accounts for the
|
||||
expiry buffer so the cache entry expires before the real token does.
|
||||
"""
|
||||
client: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
|
||||
if not server.client_id or not server.client_secret or not server.token_url:
|
||||
token_url: Final = server.effective_token_url
|
||||
if not server.client_id or not server.client_secret or not token_url:
|
||||
raise ValueError(
|
||||
f"MCP server '{server.server_id}' missing required OAuth2 fields: "
|
||||
f"client_id={bool(server.client_id)}, "
|
||||
f"client_secret={bool(server.client_secret)}, "
|
||||
f"token_url={bool(server.token_url)}"
|
||||
f"token_url={bool(token_url)}"
|
||||
)
|
||||
|
||||
token_request: Final = build_upstream_oauth2_token_request(
|
||||
|
|
@ -146,7 +147,7 @@ class MCPOAuth2TokenCache(InMemoryCache):
|
|||
)
|
||||
|
||||
try:
|
||||
response: Final = await client.post(server.token_url, data=data, headers=token_request.headers or None)
|
||||
response: Final = await client.post(token_url, data=data, headers=token_request.headers or None)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ValueError(
|
||||
|
|
|
|||
|
|
@ -142,7 +142,7 @@ def _client_credentials_spec(server: MCPServer, resource: str) -> ServerSpec:
|
|||
config=ClientCredentialsConfig(
|
||||
client_id=server.client_id,
|
||||
client_secret=SecretStr(server.client_secret) if server.client_secret else None,
|
||||
token_url=server.token_url,
|
||||
token_url=server.effective_token_url,
|
||||
scopes=tuple(server.scopes or ()),
|
||||
audience=server.audience,
|
||||
upstream_resource=resolve_upstream_resource(server),
|
||||
|
|
@ -163,7 +163,7 @@ def _token_exchange_spec(server: MCPServer, resource: str) -> ServerSpec | None:
|
|||
normalizes to ``rfc8693`` so a bad config value cannot crash spec-building. ``audience`` is
|
||||
forwarded only when the operator set it; a missing one is omitted, not derived.
|
||||
"""
|
||||
endpoint: Final = server.token_exchange_endpoint or server.token_url
|
||||
endpoint: Final = server.token_exchange_endpoint or server.effective_token_url
|
||||
if not server.client_id or not server.client_secret:
|
||||
return None
|
||||
profile: Final[Literal["rfc8693", "entra_obo"]] = (
|
||||
|
|
|
|||
|
|
@ -88,7 +88,10 @@ class AuthorizationCodeRefresher:
|
|||
if token.refresh_token is None:
|
||||
return None
|
||||
server: Final = self._server_lookup(server_id)
|
||||
if server is None or not server.token_url:
|
||||
if server is None:
|
||||
return None
|
||||
token_url: Final = server.effective_token_url
|
||||
if not token_url:
|
||||
return None
|
||||
|
||||
try:
|
||||
|
|
@ -106,7 +109,7 @@ class AuthorizationCodeRefresher:
|
|||
"refresh_token": token.refresh_token,
|
||||
**token_request.body,
|
||||
}
|
||||
body: Final = await self._token_endpoint(server.token_url, form, token_request.headers)
|
||||
body: Final = await self._token_endpoint(token_url, form, token_request.headers)
|
||||
if body is None:
|
||||
return None
|
||||
access_token: Final = body.get("access_token")
|
||||
|
|
|
|||
|
|
@ -579,3 +579,22 @@ def test_id_jag_honors_explicit_subject_token_type():
|
|||
def test_id_jag_half_configured_defers_to_v1(server):
|
||||
# A half-configured server must defer (None) rather than 500 at IdJagConfig construction.
|
||||
assert to_server_spec(server) is None
|
||||
|
||||
|
||||
def test_client_credentials_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
|
||||
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
|
||||
admin-entered value; the M2M spec must carry it so egress can mint."""
|
||||
spec = to_server_spec(
|
||||
_server(
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow="client_credentials",
|
||||
url="https://up.example.com/mcp",
|
||||
token_url=None,
|
||||
configured_token_url="https://idp.example.com/token",
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
)
|
||||
)
|
||||
assert spec is not None
|
||||
assert isinstance(spec.config, ClientCredentialsConfig)
|
||||
assert spec.config.token_url == "https://idp.example.com/token"
|
||||
|
|
|
|||
|
|
@ -20,8 +20,10 @@ class _Server:
|
|||
upstream_resource=None,
|
||||
url=None,
|
||||
server_id="srv",
|
||||
configured_token_url=None,
|
||||
):
|
||||
self.token_url = token_url
|
||||
self.configured_token_url = configured_token_url
|
||||
self.client_id = client_id
|
||||
self.client_secret = client_secret
|
||||
self.token_endpoint_auth_method = token_endpoint_auth_method
|
||||
|
|
@ -29,6 +31,10 @@ class _Server:
|
|||
self.url = url
|
||||
self.server_id = server_id
|
||||
|
||||
@property
|
||||
def effective_token_url(self):
|
||||
return self.token_url or self.configured_token_url
|
||||
|
||||
|
||||
def _lookup(server):
|
||||
return lambda server_id: server
|
||||
|
|
@ -262,3 +268,22 @@ async def test_returned_scope_overrides_prior_when_present():
|
|||
assert token is not None
|
||||
assert token.scopes == ("read",) # a present scope replaces the prior grant
|
||||
assert persisted[0][5] == ("read",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
|
||||
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
|
||||
admin-entered value; the refresh grant must POST there instead of silently failing."""
|
||||
posted = []
|
||||
refresher = _refresher(
|
||||
server=_Server(token_url=None, configured_token_url="https://idp.example.com/token"),
|
||||
body={"access_token": "new-at", "expires_in": 3600},
|
||||
post_sink=posted,
|
||||
)
|
||||
token = await refresher.refresh(
|
||||
"alice", "srv", OAuthToken(access_token="old", refresh_token="old-rt")
|
||||
)
|
||||
|
||||
assert token is not None
|
||||
assert token.access_token == "new-at"
|
||||
assert posted[0][0] == "https://idp.example.com/token"
|
||||
|
|
|
|||
|
|
@ -1337,3 +1337,28 @@ def test_mcp_oauth_token_identity_changes_when_only_upstream_resource_is_edited(
|
|||
assert mcp_oauth_token_identity(set_to_explicit) == mcp_oauth_token_identity(
|
||||
_identity_server(credentials={**creds, "upstream_resource": "api://audience-one"})
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer_yield_empties_resolved(monkeypatch):
|
||||
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
|
||||
admin-entered value; the silent per-user refresh must POST there instead of bailing."""
|
||||
from litellm.proxy._types import MCPTransport
|
||||
from litellm.types.mcp import MCPAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
server = MCPServer(
|
||||
server_id="srv-1",
|
||||
name="test",
|
||||
url="https://up.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
client_id="cid",
|
||||
client_secret="csec",
|
||||
token_url=None,
|
||||
configured_token_url="https://idp.example.com/token",
|
||||
)
|
||||
result, captured = await _run_refresh(monkeypatch, server)
|
||||
|
||||
assert result is not None
|
||||
assert captured["url"] == "https://idp.example.com/token"
|
||||
|
|
|
|||
|
|
@ -392,3 +392,22 @@ async def test_invalidate_clears_every_identity_for_a_server():
|
|||
|
||||
assert refetched == "tok-after-invalidate"
|
||||
assert mock_client.post.call_count == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_mint_uses_admin_entered_token_url_when_issuer_yield_empties_resolved():
|
||||
"""A pinned issuer empties the resolved token_url while configured_token_url keeps the
|
||||
admin-entered value; the client_credentials mint must POST there instead of raising."""
|
||||
server = _server(token_url=None, configured_token_url="https://auth.example.com/token")
|
||||
cache = MCPOAuth2TokenCache()
|
||||
mock_client = AsyncMock()
|
||||
mock_client.post.return_value = _token_response("m2m-token-configured")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.oauth2_token_cache.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await cache.async_get_token(server)
|
||||
|
||||
assert result == "m2m-token-configured"
|
||||
assert mock_client.post.call_args[0][0] == "https://auth.example.com/token"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue