mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(mcp): challenge keyless BYOK servers with 401 on single-server routes
A single-server-scoped MCP route (/mcp/{server} or /{server}/mcp) to a BYOK
server with no stored per-user credential listed empty instead of prompting
the caller to authenticate, because the tools-list paths degrade on a missing
key and the byok_auth_required 401 only fired on the tool-call path. Add a BYOK
branch to the connect-time preemptive 401 so SSE and StreamableHTTP clients get
a proper WWW-Authenticate challenge and can run the key-entry flow.
Gated to single-server-scoped routes so a credential-less BYOK server in a
multi-server aggregated listing still degrades per-server rather than aborting
the whole listing. The challenge carries an absolute RFC 9728 resource_metadata
URI pointing at the gateway's BYOK protected-resource document.
This commit is contained in:
Tin Chi Lo
parent
9f859053a6
commit
ce38b88b1e
2 changed files with 184 additions and 0 deletions
|
|
@ -3493,6 +3493,8 @@ if MCP_AVAILABLE:
|
|||
excludes a passthrough server is not pushed into an OAuth flow for
|
||||
a server it will be 403'd on immediately after authentication.
|
||||
"""
|
||||
_req_path = scope.get("_original_path") or scope.get("path", "") or ""
|
||||
is_single_server_scoped = len(_get_mcp_servers_in_path(_req_path) or []) == 1
|
||||
for server_name in mcp_servers or []:
|
||||
server = global_mcp_server_manager.get_mcp_server_by_name(
|
||||
server_name, client_ip=client_ip
|
||||
|
|
@ -3562,6 +3564,41 @@ if MCP_AVAILABLE:
|
|||
headers={"www-authenticate": www_authenticate},
|
||||
)
|
||||
|
||||
# BYOK: on a single-server-scoped route, when the caller supplied no
|
||||
# key inline and has none stored, fail fast with 401 + challenge so
|
||||
# OAuth-capable MCP clients run the key-entry flow instead of seeing
|
||||
# a silently empty tool list. Gated to single-server routes so one
|
||||
# credential-less BYOK server cannot abort a multi-server aggregated
|
||||
# listing, which keeps degrading per-server.
|
||||
if (
|
||||
server
|
||||
and server.is_byok
|
||||
and is_single_server_scoped
|
||||
and not _client_has_passthrough_authorization(
|
||||
server, oauth2_headers, mcp_server_auth_headers
|
||||
)
|
||||
and await _get_byok_credential(server, user_api_key_auth) is None
|
||||
):
|
||||
base_url = get_request_base_url(StarletteRequest(scope))
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": "byok_auth_required",
|
||||
"server_id": server.server_id,
|
||||
"server_name": server.server_name or server.name,
|
||||
"message": (
|
||||
"No stored credential found for this BYOK server. "
|
||||
"Complete the OAuth authorization flow to provide your API key."
|
||||
),
|
||||
},
|
||||
headers={
|
||||
"www-authenticate": (
|
||||
"Bearer resource_metadata="
|
||||
f'"{base_url}/.well-known/oauth-protected-resource"'
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
def _get_forwarded_auth_from_scope(scope: Scope) -> Optional[str]:
|
||||
"""Return the upstream-bound ``Authorization`` header value, or None.
|
||||
|
||||
|
|
|
|||
|
|
@ -1101,6 +1101,153 @@ async def test_get_tools_from_mcp_servers_byok_lookup_error_does_not_abort_listi
|
|||
assert captured["mcp_auth_header"] is None
|
||||
|
||||
|
||||
def _byok_preemptive_scope(path: str) -> dict:
|
||||
return {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"scheme": "http",
|
||||
"server": ("testserver", 80),
|
||||
"path": path,
|
||||
"headers": [(b"host", b"testserver")],
|
||||
"client": ("1.2.3.4", 12345),
|
||||
}
|
||||
|
||||
|
||||
def _byok_preemptive_server() -> MagicMock:
|
||||
server = MagicMock()
|
||||
server.name = "byok_server"
|
||||
server.alias = "byok"
|
||||
server.server_id = "byok_server_id"
|
||||
server.server_name = "byok_server"
|
||||
server.auth_type = MCPAuth.none
|
||||
server.is_byok = True
|
||||
server.is_oauth_passthrough = False
|
||||
server.needs_user_oauth_token = False
|
||||
return server
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preemptive_401_byok_single_server_no_credential_raises():
|
||||
"""On a single-server-scoped route, a BYOK server with no stored credential
|
||||
and no inline key must fail fast with 401 + WWW-Authenticate so OAuth-capable
|
||||
MCP clients run the key-entry flow instead of seeing an empty tool list."""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
|
||||
server = _byok_preemptive_server()
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_name = MagicMock(return_value=server)
|
||||
|
||||
with (
|
||||
patch.object(mcp_server_module, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mcp_server_module, "_get_byok_credential", AsyncMock(return_value=None)
|
||||
) as mock_cred,
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await mcp_server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope=_byok_preemptive_scope("/mcp/byok_server"),
|
||||
mcp_servers=["byok_server"],
|
||||
oauth2_headers=None,
|
||||
mcp_server_auth_headers=None,
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="k", user_id="u"),
|
||||
client_ip=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
challenge = exc_info.value.headers["www-authenticate"]
|
||||
assert challenge.startswith("Bearer ")
|
||||
# Absolute RFC 9728 resource_metadata URI (not the relative form strict
|
||||
# clients reject), pointing at the gateway's BYOK protected-resource doc.
|
||||
assert 'resource_metadata="http' in challenge
|
||||
assert "/.well-known/oauth-protected-resource" in challenge
|
||||
mock_cred.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preemptive_401_byok_single_server_with_credential_no_raise():
|
||||
"""A stored credential satisfies BYOK auth, so the preemptive challenge must
|
||||
not fire (otherwise an already-authenticated user is blocked from listing)."""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
|
||||
server = _byok_preemptive_server()
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_name = MagicMock(return_value=server)
|
||||
|
||||
with (
|
||||
patch.object(mcp_server_module, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mcp_server_module,
|
||||
"_get_byok_credential",
|
||||
AsyncMock(return_value="stored-key"),
|
||||
),
|
||||
):
|
||||
await mcp_server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope=_byok_preemptive_scope("/mcp/byok_server"),
|
||||
mcp_servers=["byok_server"],
|
||||
oauth2_headers=None,
|
||||
mcp_server_auth_headers=None,
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="k", user_id="u"),
|
||||
client_ip=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preemptive_401_byok_multi_server_route_does_not_raise():
|
||||
"""A credential-less BYOK server reached through a multi-server (aggregated)
|
||||
route must not abort the listing; the challenge is single-server only, so the
|
||||
aggregator keeps degrading per-server."""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
|
||||
server = _byok_preemptive_server()
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_name = MagicMock(return_value=server)
|
||||
|
||||
with (
|
||||
patch.object(mcp_server_module, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mcp_server_module, "_get_byok_credential", AsyncMock(return_value=None)
|
||||
) as mock_cred,
|
||||
):
|
||||
await mcp_server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope=_byok_preemptive_scope("/mcp/byok_server,other_server"),
|
||||
mcp_servers=["byok_server"],
|
||||
oauth2_headers=None,
|
||||
mcp_server_auth_headers=None,
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="k", user_id="u"),
|
||||
client_ip=None,
|
||||
)
|
||||
|
||||
mock_cred.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_preemptive_401_byok_inline_credential_does_not_raise():
|
||||
"""When the caller supplies the BYOK key inline via a per-server auth header,
|
||||
the preemptive challenge must be skipped."""
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server_module
|
||||
|
||||
server = _byok_preemptive_server()
|
||||
manager = MagicMock()
|
||||
manager.get_mcp_server_by_name = MagicMock(return_value=server)
|
||||
|
||||
with (
|
||||
patch.object(mcp_server_module, "global_mcp_server_manager", manager),
|
||||
patch.object(
|
||||
mcp_server_module, "_get_byok_credential", AsyncMock(return_value=None)
|
||||
) as mock_cred,
|
||||
):
|
||||
await mcp_server_module._raise_preemptive_401_for_unauthenticated_servers(
|
||||
scope=_byok_preemptive_scope("/mcp/byok_server"),
|
||||
mcp_servers=["byok_server"],
|
||||
oauth2_headers=None,
|
||||
mcp_server_auth_headers={"byok": "user-supplied-key"},
|
||||
user_api_key_auth=UserAPIKeyAuth(api_key="k", user_id="u"),
|
||||
client_ip=None,
|
||||
)
|
||||
|
||||
mock_cred.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
|
||||
"""Test that _get_tools_from_mcp_servers handles all servers failing gracefully"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue