From ce38b88b1eb503cf2a5c617d2101f301d0516648 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Wed, 24 Jun 2026 16:42:44 -0700 Subject: [PATCH] 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. --- .../proxy/_experimental/mcp_server/server.py | 37 +++++ .../mcp_server/test_mcp_server.py | 147 ++++++++++++++++++ 2 files changed, 184 insertions(+) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 3df7b01e57c..a5429cd6c76 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 1c00372fe5f..90cd5f86c9b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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"""