fix(mcp): token refresh and M2M egress honor the admin-entered token URL

This commit is contained in:
mateo-berri 2026-08-26 15:31:45 -07:00
parent 465ebb1bdd
commit ce8d6f7d25
8 changed files with 104 additions and 12 deletions

View file

@ -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)

View file

@ -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(

View file

@ -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"]] = (

View file

@ -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")

View file

@ -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"

View file

@ -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"

View file

@ -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"

View file

@ -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"