mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
Merge pull request #33586 from BerriAI/litellm_mcp_oauth_challenge_mode_aware
fix(mcp): make the preemptive-401 OAuth challenge decision mode-aware
This commit is contained in:
tin-berri
GitHub
commit
ecef9e6c9b
3 changed files with 290 additions and 37 deletions
|
|
@ -3582,48 +3582,70 @@ if MCP_AVAILABLE:
|
|||
# preemptive challenge and let downstream authorization
|
||||
# return 403.
|
||||
continue
|
||||
if server and server.auth_type == MCPAuth.oauth2 and not oauth2_headers:
|
||||
# For per-user OAuth servers, only skip the pre-emptive 401 when
|
||||
# a stored token actually exists for this user+server pair.
|
||||
# If no stored token exists, fail fast with 401 so clients can
|
||||
# kick off PKCE/interactive OAuth flow immediately.
|
||||
if server.needs_user_oauth_token:
|
||||
if getattr(server, "delegate_auth_to_upstream", False) is True:
|
||||
# Delegate-auth servers run upstream PKCE: challenge with
|
||||
# the proxied resource_metadata (RFC 9728), not the
|
||||
# gateway authorization_uri below which would authorize
|
||||
# against the gateway instead of the upstream IdP.
|
||||
www_authenticate = _get_passthrough_www_authenticate(
|
||||
scope=scope,
|
||||
server_name=server_name,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Unauthorized",
|
||||
headers={"www-authenticate": www_authenticate},
|
||||
)
|
||||
# The v2 resolver owns the existence check, so every authorization_code
|
||||
# resolution (egress and this discovery challenge) runs through it.
|
||||
if server and server.auth_type == MCPAuth.oauth2:
|
||||
# The challenge decision is per oauth2 sub-mode, not per header:
|
||||
# gateway-managed modes (M2M and interactive authorization_code)
|
||||
# never receive a client-supplied upstream token, so a bearer in
|
||||
# Authorization is a LiteLLM key (surfaced here as oauth2_headers)
|
||||
# and must not suppress the challenge. Only the delegate mode
|
||||
# treats a present bearer as the upstream token. The sub-mode is
|
||||
# resolved the same way egress resolves it, via
|
||||
# effective_oauth2_flow: an unstamped (null oauth2_flow) row with
|
||||
# the M2M shape resolves to client_credentials, so the bare
|
||||
# has_client_credentials column is never trusted here.
|
||||
if MCPServerManager.effective_oauth2_flow(server) == "client_credentials":
|
||||
# M2M: the gateway mints its own token at egress from the
|
||||
# stored client credentials, so there is nothing to challenge.
|
||||
continue
|
||||
|
||||
if getattr(server, "delegate_auth_to_upstream", False) is not True:
|
||||
# Gateway-managed interactive (authorization_code): the only
|
||||
# thing that authorizes egress is a stored per-user token, so
|
||||
# challenge whenever one is absent, regardless of any bearer.
|
||||
# The v2 resolver owns the existence check, so every
|
||||
# authorization_code resolution (egress and this discovery
|
||||
# challenge) runs through it.
|
||||
if await global_mcp_server_manager.has_user_oauth_token(server, user_api_key_auth):
|
||||
continue
|
||||
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
_path = scope.get("_original_path") or scope.get("path", "") or ""
|
||||
request = StarletteRequest(scope)
|
||||
base_url = get_request_base_url(request)
|
||||
_path = scope.get("_original_path") or scope.get("path", "") or ""
|
||||
|
||||
# Pick the well-known AS-metadata form that matches the inbound route
|
||||
# so strict RFC 9728 §3.2 clients can resolve it correctly.
|
||||
if _path.startswith(f"/mcp/{server_name}"):
|
||||
_as_url = f"{base_url}/.well-known/oauth-authorization-server/mcp/{server_name}"
|
||||
else:
|
||||
_as_url = f"{base_url}/.well-known/oauth-authorization-server/{server_name}"
|
||||
authorization_uri = f'Bearer authorization_uri="{_as_url}"'
|
||||
# Pick the well-known AS-metadata form that matches the inbound route
|
||||
# so strict RFC 9728 §3.2 clients can resolve it correctly.
|
||||
if _path.startswith(f"/mcp/{server_name}"):
|
||||
_as_url = f"{base_url}/.well-known/oauth-authorization-server/mcp/{server_name}"
|
||||
else:
|
||||
_as_url = f"{base_url}/.well-known/oauth-authorization-server/{server_name}"
|
||||
authorization_uri = f'Bearer authorization_uri="{_as_url}"'
|
||||
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Unauthorized",
|
||||
headers={"www-authenticate": authorization_uri},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Unauthorized",
|
||||
headers={"www-authenticate": authorization_uri},
|
||||
)
|
||||
|
||||
if not oauth2_headers:
|
||||
# Delegate-auth servers run upstream PKCE: a present bearer is
|
||||
# the upstream token, so only challenge when it is absent, with
|
||||
# the proxied resource_metadata (RFC 9728), not the gateway
|
||||
# authorization_uri above which would authorize against the
|
||||
# gateway instead of the upstream IdP.
|
||||
www_authenticate = _get_passthrough_www_authenticate(
|
||||
scope=scope,
|
||||
server_name=server_name,
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Unauthorized",
|
||||
headers={"www-authenticate": www_authenticate},
|
||||
)
|
||||
# Delegate server with a bearer present: it is the upstream token,
|
||||
# so admit the session and move to the next target. Every oauth2
|
||||
# sub-mode is terminal here (continue or raise) so no oauth2 server
|
||||
# reaches the token_exchange / pass-through blocks below.
|
||||
continue
|
||||
|
||||
# token_exchange (OBO): the caller supplied no subject token. Challenge at connect
|
||||
# (transport level, where WWW-Authenticate survives) with the RFC 9728 resource_metadata
|
||||
|
|
|
|||
|
|
@ -7437,3 +7437,139 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error():
|
|||
)
|
||||
|
||||
proxy_logging_mock.post_call_failure_hook.assert_not_awaited()
|
||||
|
||||
|
||||
def _make_oauth2_server(
|
||||
alias: str,
|
||||
*,
|
||||
oauth2_flow=None,
|
||||
delegate_auth_to_upstream: bool = False,
|
||||
client_id=None,
|
||||
client_secret=None,
|
||||
token_url=None,
|
||||
) -> MCPServer:
|
||||
"""An auth_type=oauth2 MCP server in one of its sub-modes. oauth2_flow
|
||||
'client_credentials' is M2M; delegate_auth_to_upstream toggles the
|
||||
upstream-PKCE delegate mode; the default is gateway-managed interactive
|
||||
(authorization_code). client_id/client_secret/token_url set the M2M shape
|
||||
that effective_oauth2_flow infers as client_credentials when oauth2_flow is
|
||||
left unstamped (null)."""
|
||||
return MCPServer(
|
||||
server_id=f"id-{alias}",
|
||||
name=alias,
|
||||
alias=alias,
|
||||
server_name=alias,
|
||||
url=f"https://{alias}.test/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
oauth2_flow=oauth2_flow,
|
||||
delegate_auth_to_upstream=delegate_auth_to_upstream,
|
||||
client_id=client_id,
|
||||
client_secret=client_secret,
|
||||
token_url=token_url,
|
||||
mcp_info={"server_name": alias},
|
||||
)
|
||||
|
||||
|
||||
class TestPreemptive401ModeAware:
|
||||
"""The preemptive-401 challenge for auth_type=oauth2 servers is decided by
|
||||
the server's sub-mode, not by whether an Authorization header is present.
|
||||
|
||||
Regression guard for the bug where a LiteLLM virtual key presented as
|
||||
``Authorization: Bearer sk-...`` (indistinguishable at header-parse time
|
||||
from an upstream OAuth bearer, so it lands in oauth2_headers) suppressed
|
||||
the challenge on a gateway-managed authorization_code server, opening a
|
||||
session with no upstream token whose tools/list masks as 200 + empty.
|
||||
"""
|
||||
|
||||
LITELLM_KEY_HEADERS = {"Authorization": "Bearer sk-litellm-virtual-key"}
|
||||
|
||||
def _scope(self, alias: str):
|
||||
return {"type": "http", "method": "POST", "path": f"/mcp/{alias}", "headers": []}
|
||||
|
||||
async def _run(self, server, oauth2_headers, has_stored_token: bool):
|
||||
from litellm.proxy._experimental.mcp_server import server as server_module
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
server_module.global_mcp_server_manager,
|
||||
"get_mcp_server_by_name",
|
||||
return_value=server,
|
||||
),
|
||||
patch.object(
|
||||
server_module.global_mcp_server_manager,
|
||||
"has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
return_value=has_stored_token,
|
||||
),
|
||||
):
|
||||
await server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope=self._scope(server.alias),
|
||||
mcp_servers=[server.alias],
|
||||
oauth2_headers=oauth2_headers,
|
||||
mcp_server_auth_headers=None,
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"),
|
||||
client_ip=None,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_managed_interactive_no_token_challenges_with_x_litellm_api_key(self):
|
||||
"""No stored token, key in x-litellm-api-key (oauth2_headers empty): 401."""
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await self._run(_make_oauth2_server("interactive"), None, has_stored_token=False)
|
||||
assert exc.value.status_code == 401
|
||||
assert "www-authenticate" in {k.lower() for k in exc.value.headers}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_managed_interactive_no_token_challenges_with_authorization_bearer(self):
|
||||
"""The bug fix: no stored token, key in Authorization (oauth2_headers
|
||||
populated) must still get the 401 challenge, not a suppressed session."""
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await self._run(
|
||||
_make_oauth2_server("interactive"),
|
||||
self.LITELLM_KEY_HEADERS,
|
||||
has_stored_token=False,
|
||||
)
|
||||
assert exc.value.status_code == 401
|
||||
assert "www-authenticate" in {k.lower() for k in exc.value.headers}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gateway_managed_interactive_with_stored_token_does_not_challenge(self):
|
||||
"""A stored per-user token exists: no challenge, under either header."""
|
||||
await self._run(_make_oauth2_server("interactive"), None, has_stored_token=True)
|
||||
await self._run(_make_oauth2_server("interactive"), self.LITELLM_KEY_HEADERS, has_stored_token=True)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_m2m_never_challenges(self):
|
||||
"""client_credentials (M2M): the gateway mints its own token, so no
|
||||
challenge regardless of header or stored-token state."""
|
||||
m2m = _make_oauth2_server("m2m", oauth2_flow="client_credentials")
|
||||
await self._run(m2m, None, has_stored_token=False)
|
||||
await self._run(m2m, self.LITELLM_KEY_HEADERS, has_stored_token=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unstamped_m2m_shape_never_challenges(self):
|
||||
"""A legacy row with oauth2_flow left null but the M2M shape
|
||||
(client_id + client_secret + token_url) resolves to client_credentials
|
||||
via effective_oauth2_flow exactly as egress does, so it is treated as
|
||||
M2M and never challenged. The bare oauth2_flow column would misread it
|
||||
as interactive and raise a spurious 401."""
|
||||
unstamped = _make_oauth2_server(
|
||||
"unstampedm2m",
|
||||
oauth2_flow=None,
|
||||
client_id="cid",
|
||||
client_secret="csecret",
|
||||
token_url="https://idp.test/token",
|
||||
)
|
||||
await self._run(unstamped, None, has_stored_token=False)
|
||||
await self._run(unstamped, self.LITELLM_KEY_HEADERS, has_stored_token=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delegate_challenges_only_when_bearer_absent(self):
|
||||
"""delegate_auth_to_upstream: a present bearer IS the upstream token,
|
||||
so challenge only when it is absent."""
|
||||
delegate = _make_oauth2_server("delegate", delegate_auth_to_upstream=True)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await self._run(delegate, None, has_stored_token=False)
|
||||
assert exc.value.status_code == 401
|
||||
await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
|
||||
|
|
|
|||
|
|
@ -664,6 +664,101 @@ async def test_per_user_oauth_missing_stored_token_returns_preemptive_401():
|
|||
assert "Bearer authorization_uri=" in exc_info.value.headers["www-authenticate"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"m2m_fields",
|
||||
[
|
||||
{"oauth2_flow": "client_credentials"},
|
||||
{"client_id": "cid", "client_secret": "csec", "token_url": "https://idp.example.com/token"},
|
||||
],
|
||||
ids=["stamped", "unstamped_m2m_shape"],
|
||||
)
|
||||
async def test_client_credentials_server_is_not_preemptively_challenged(m2m_fields):
|
||||
"""
|
||||
An OAuth2 client_credentials (M2M) server mints its own upstream token;
|
||||
there is no user OAuth flow to bootstrap. The connect-time gate must let
|
||||
the request through to the session manager rather than pushing the client
|
||||
into an interactive OAuth flow it can never complete (the per-user token
|
||||
store is never even consulted for M2M). Covers both a stamped row and a
|
||||
legacy null-flow row with the M2M field shape: the gate must classify the
|
||||
flow through the same request-time chokepoint egress uses, or the two
|
||||
disagree and the unstamped server is challenged for a token egress would
|
||||
never look for.
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
session_manager_stateless,
|
||||
)
|
||||
from litellm.proxy._types import MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp",
|
||||
"scheme": "http",
|
||||
"query_string": b"",
|
||||
"root_path": "",
|
||||
"server": ("localhost", 8000),
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"host", b"localhost:8000"),
|
||||
],
|
||||
}
|
||||
receive = AsyncMock(
|
||||
return_value={
|
||||
"type": "http.request",
|
||||
"body": b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}',
|
||||
"more_body": False,
|
||||
}
|
||||
)
|
||||
send = AsyncMock()
|
||||
user_auth = MagicMock()
|
||||
user_auth.user_id = "test-user-id"
|
||||
m2m_server = MCPServer(
|
||||
server_id="m2m-server-id",
|
||||
name="m2m_server",
|
||||
server_name="m2m_server",
|
||||
alias="m2m_server",
|
||||
url="https://upstream.example.com/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
**m2m_fields,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(user_auth, None, ["m2m_server"], None, None, None),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"),
|
||||
patch("litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._handle_stale_mcp_session",
|
||||
new_callable=AsyncMock,
|
||||
return_value=False,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.has_user_oauth_token",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_has_token,
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_name",
|
||||
return_value=m2m_server,
|
||||
),
|
||||
patch.object(session_manager_stateless, "handle_request", new_callable=AsyncMock) as mock_handle_request,
|
||||
patch.object(session_manager_stateless, "_server_instances", {}),
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert mock_handle_request.await_count == 1
|
||||
assert mock_has_token.await_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_streamable_http_mcp_delegated_server_surfaces_upstream_challenge():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue