From 6d5c2d85ef31621e3fcee288e8fd18ddafa5d061 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:44:13 -0700 Subject: [PATCH 1/3] fix(mcp): use gateway authentication for root discovery --- .../mcp_server/byok_oauth_endpoints.py | 13 -- .../mcp_server/discoverable_endpoints.py | 27 ++-- litellm/proxy/_lazy_openapi_snapshot.json | 98 +++++--------- .../mcp_server/test_byok_oauth_endpoints.py | 9 -- .../mcp_server/test_discoverable_endpoints.py | 128 +++++++++++------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 32 +---- 6 files changed, 138 insertions(+), 169 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 28f3ec6521a..e8f0f9b8e8b 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -7,7 +7,6 @@ just a form that asks the user for their API key — not a full identity-provide Endpoints implemented here: GET /.well-known/oauth-authorization-server — OAuth authorization server metadata - GET /.well-known/oauth-protected-resource — OAuth protected resource metadata GET /v1/mcp/oauth/authorize — Shows HTML form to collect the API key POST /v1/mcp/oauth/authorize — Stores temp auth code and redirects POST /v1/mcp/oauth/token — Exchanges code for a bearer JWT token @@ -612,18 +611,6 @@ async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: ) -@router.get("/.well-known/oauth-protected-resource", include_in_schema=False) -async def oauth_protected_resource_metadata(request: Request) -> JSONResponse: - """RFC 9728 Protected Resource Metadata pointing back at this server.""" - base_url: Final = get_request_base_url(request) - return JSONResponse( - { - "resource": base_url, - "authorization_servers": [base_url], - } - ) - - # --------------------------------------------------------------------------- # Authorization endpoint — GET (show form) and POST (process form) # --------------------------------------------------------------------------- diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 42fbe82531c..eb24ed169de 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2393,8 +2393,7 @@ async def _build_oauth_protected_resource_response( per-server URL completes the same sign-in flow the aggregate ``/mcp`` endpoint supports and is admitted with a gateway session bearer. The per-server relay authorize/token endpoints stay registered for the keyed interactive flow (which - is challenged with an explicit ``authorization_uri``), and the root-resolved - (unnamed) legacy shape keeps the relay authorization server. + is challenged with an explicit ``authorization_uri``). Args: request: FastAPI Request object @@ -2405,15 +2404,11 @@ async def _build_oauth_protected_resource_response( Returns: OAuth protected resource metadata dict """ + if mcp_server_name is None: + return oauth_protected_resource_root(request) + request_base_url: Final = get_request_base_url(request) client_ip: Final = IPAddressUtils.get_mcp_client_ip(request) - explicitly_named: Final = mcp_server_name is not None - - # When no server name provided, try to resolve the single OAuth2 server - if mcp_server_name is None: - resolved: Final = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip) - if resolved: - mcp_server_name = resolved.server_name or resolved.name mcp_server: MCPServer | None = None if mcp_server_name: @@ -2478,7 +2473,7 @@ async def _build_oauth_protected_resource_response( if obo_response is not None: return obo_response - if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: + if mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, @@ -2542,6 +2537,17 @@ def _jwt_auth_issuers() -> list: return issuers +@router.get("/.well-known/oauth-protected-resource") +def oauth_protected_resource_root(request: Request) -> dict[str, str | list[str]]: + request_base_url: Final = get_request_base_url(request) + parsed: Final = urlparse(request_base_url) + return { + "resource": f"{parsed.scheme}://{parsed.netloc}", + "authorization_servers": [f"{request_base_url}/mcp"], + "scopes_supported": [], + } + + def _build_aggregate_protected_resource_response(request: Request) -> dict: """RFC 9728 metadata for the aggregate /mcp resource: the gateway itself is the authorization server. No per-server names or scopes leak here; access @@ -2645,7 +2651,6 @@ async def oauth_protected_resource_mcp_standard(request: Request, mcp_server_nam # LiteLLM legacy pattern: /.well-known/oauth-protected-resource/{server_name}/mcp # Kept for backward compatibility with existing deployments @router.get(f"/.well-known/oauth-protected-resource{well_known_root_suffix()}/{{mcp_server_name}}/mcp") -@router.get("/.well-known/oauth-protected-resource") async def oauth_protected_resource_mcp(request: Request, mcp_server_name: str | None = None): """ OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 53af85baac6..16681613c5d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -22415,47 +22415,34 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get", - "parameters": [ - { - "in": "query", - "name": "mcp_server_name", - "required": false, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Mcp Server Name" - } - } - ], + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get", "responses": { "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/HTTPValidationError" + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ] + }, + "title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get", + "type": "object" } } }, - "description": "Validation Error" + "description": "Successful Response" } }, - "summary": "Oauth Protected Resource Mcp", + "summary": "Oauth Protected Resource Root", "tags": [ "mcp_byok_oauth" ] @@ -24636,47 +24623,34 @@ }, "/.well-known/oauth-protected-resource": { "get": { - "description": "OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern.\n\nLegacy pattern: /{server_name}/mcp\nDiscovery path: /.well-known/oauth-protected-resource/{server_name}/mcp\n\nThis endpoint is kept for backward compatibility. New integrations should\nuse the standard MCP pattern (/mcp/{server_name}) instead.", - "operationId": "oauth_protected_resource_mcp__well_known_oauth_protected_resource_get_2", - "parameters": [ - { - "in": "query", - "name": "mcp_server_name", - "required": false, - "schema": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Mcp Server Name" - } - } - ], + "operationId": "oauth_protected_resource_root__well_known_oauth_protected_resource_get_2", "responses": { "200": { - "content": { - "application/json": { - "schema": {} - } - }, - "description": "Successful Response" - }, - "422": { "content": { "application/json": { "schema": { - "$ref": "#/components/schemas/HTTPValidationError" + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "items": { + "type": "string" + }, + "type": "array" + } + ] + }, + "title": "Response Oauth Protected Resource Root Well Known Oauth Protected Resource Get", + "type": "object" } } }, - "description": "Validation Error" + "description": "Successful Response" } }, - "summary": "Oauth Protected Resource Mcp", + "summary": "Oauth Protected Resource Root", "tags": [ "mcp_discoverable" ] diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 9f2feddb0e3..f5bc16375fd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -107,15 +107,6 @@ def test_oauth_authorization_server_metadata(client): assert "S256" in data["code_challenge_methods_supported"] -def test_oauth_protected_resource_metadata(client): - resp = client.get("/.well-known/oauth-protected-resource") - assert resp.status_code == 200 - data = resp.json() - assert "resource" in data - assert "authorization_servers" in data - assert len(data["authorization_servers"]) == 1 - - # --------------------------------------------------------------------------- # Authorization GET endpoint # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 1b7e3d6a1c3..7da8010869e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3433,51 +3433,84 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa global_mcp_server_manager.registry.clear() +@pytest.mark.parametrize("server_count", [0, 1, 2]) +@pytest.mark.parametrize("byok_first", [True, False]) +@pytest.mark.parametrize( + ("base_url", "origin"), + [ + ("https://gateway.example.com", "https://gateway.example.com"), + ("https://gateway.example.com/proxy", "https://gateway.example.com"), + ("http://[::1]:4000/proxy", "http://[::1]:4000"), + ], +) +def test_root_protected_resource_discovers_gateway(monkeypatch, server_count, byok_first, base_url, origin): + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy._types import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("PROXY_BASE_URL", base_url) + monkeypatch.setattr( + global_mcp_server_manager, + "registry", + { + f"oauth_{index}": MCPServer( + server_id=f"oauth_{index}", + name=f"oauth_{index}", + server_name=f"oauth_{index}", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + authorization_url="https://idp.example.com/authorize", + token_url="https://idp.example.com/token", + ) + for index in range(server_count) + }, + ) + app = FastAPI() + routers = (byok_oauth_endpoints.router, discoverable_endpoints.router) + for router in routers if byok_first else reversed(routers): + app.include_router(router) + with TestClient(app) as client: + response = client.get("/.well-known/oauth-protected-resource", params={"mcp_server_name": "oauth_0"}) + assert response.status_code == 200 + assert response.json() == { + "resource": origin, + "authorization_servers": [f"{base_url}/mcp"], + "scopes_supported": [], + } + authorization = client.get("/.well-known/oauth-authorization-server/mcp") + assert authorization.status_code == 200 + metadata = authorization.json() + assert metadata["issuer"] == response.json()["authorization_servers"][0] + assert metadata["authorization_endpoint"] == f"{base_url}/authorize" + assert metadata["token_endpoint"] == f"{base_url}/token" + assert metadata["registration_endpoint"] == f"{base_url}/register" + aggregate = client.get("/.well-known/oauth-protected-resource/mcp") + assert aggregate.status_code == 200 + assert aggregate.json()["resource"] == f"{base_url}/mcp" + + @pytest.mark.asyncio -async def test_oauth_protected_resource_root_resolved_single_server_keeps_relay_as(): - """The unnamed (bare-root) legacy shape resolves the single configured oauth2 server and - must keep advertising the per-server relay authorization server: only an EXPLICITLY - named request opts into the gateway-as-AS flow (LIT-4864), so pre-existing single-server - deployments discovering through the root document are byte-identical.""" - try: - from fastapi import Request +async def test_unnamed_protected_resource_builder_uses_gateway_origin(monkeypatch): + from fastapi import Request - from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - _build_oauth_protected_resource_response, - ) - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) - from litellm.proxy._types import MCPTransport - from litellm.types.mcp import MCPAuth - from litellm.types.mcp_server.mcp_server_manager import MCPServer - except ImportError: - pytest.skip("MCP discoverable endpoints not available") - - only_server = MCPServer( - server_id="solo_mcp", - name="solo_mcp", - server_name="solo_mcp", - alias="solo_mcp", - transport=MCPTransport.http, - auth_type=MCPAuth.oauth2, - authorization_url="https://idp.example.com/authorize", - token_url="https://idp.example.com/oauth/token", + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + _build_oauth_protected_resource_response, ) - mock_request = MagicMock(spec=Request) - mock_request.base_url = "https://litellm.example.com/" - mock_request.headers = {} - - global_mcp_server_manager.registry.clear() - try: - global_mcp_server_manager.registry[only_server.server_id] = only_server - response = await _build_oauth_protected_resource_response( - request=mock_request, mcp_server_name=None, use_standard_pattern=False - ) - assert response["authorization_servers"] == ["https://litellm.example.com/solo_mcp"] - finally: - global_mcp_server_manager.registry.clear() + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + request = Request( + {"type": "http", "scheme": "https", "server": ("gateway.example.com", 443), "path": "/", "headers": []} + ) + response = await _build_oauth_protected_resource_response(request, None, False) + assert response == { + "resource": "https://gateway.example.com", + "authorization_servers": ["https://gateway.example.com/mcp"], + "scopes_supported": [], + } @pytest.mark.asyncio @@ -4137,7 +4170,8 @@ async def test_discovery_root_does_not_expose_private_server_for_external_client assert "/test_oauth/" not in authorization_response["authorization_endpoint"] assert "/test_oauth/" not in authorization_response["token_endpoint"] assert authorization_response["scopes_supported"] == [] - assert resource_response["authorization_servers"] == ["https://llm.example.com"] + assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert resource_response["resource"] == "https://llm.example.com" assert resource_response["scopes_supported"] == [] finally: global_mcp_server_manager.registry.clear() @@ -9077,11 +9111,7 @@ def test_well_known_root_suffix_reflects_server_root_path(): @pytest.mark.asyncio -async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): - """The always-on aggregate front door must not change bare-origin discovery: with one - oauth2 server configured, the no-suffix /.well-known/oauth-{authorization-server, - protected-resource} still resolves THAT server, so an existing single-server deployment's - discovery is unchanged. The aggregate document lives only at the /mcp-suffixed routes.""" +async def test_root_resource_uses_gateway_without_changing_authorization_relay(): from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( @@ -9105,10 +9135,10 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate(): resource_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name=None, use_standard_pattern=True ) - # per-server, not aggregate: the single server's name is in the endpoints assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] assert authorization_response["issuer"] == "https://llm.example.com" - assert resource_response["authorization_servers"] == ["https://llm.example.com/test_oauth"] + assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert resource_response["resource"] == "https://llm.example.com" finally: global_mcp_server_manager.registry.clear() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index dd455cba44a..ea76aeb2a36 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -207,17 +207,8 @@ export interface paths { path?: never; cookie?: never; }; - /** - * Oauth Protected Resource Mcp - * @description OAuth protected resource discovery endpoint using LiteLLM legacy URL pattern. - * - * Legacy pattern: /{server_name}/mcp - * Discovery path: /.well-known/oauth-protected-resource/{server_name}/mcp - * - * This endpoint is kept for backward compatibility. New integrations should - * use the standard MCP pattern (/mcp/{server_name}) instead. - */ - get: operations["oauth_protected_resource_mcp__well_known_oauth_protected_resource_get"]; + /** Oauth Protected Resource Root */ + get: operations["oauth_protected_resource_root__well_known_oauth_protected_resource_get"]; put?: never; post?: never; delete?: never; @@ -40285,11 +40276,9 @@ export interface operations { }; }; }; - oauth_protected_resource_mcp__well_known_oauth_protected_resource_get: { + oauth_protected_resource_root__well_known_oauth_protected_resource_get: { parameters: { - query?: { - mcp_server_name?: string | null; - }; + query?: never; header?: never; path?: never; cookie?: never; @@ -40302,16 +40291,9 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; - }; - }; - /** @description Validation Error */ - 422: { - headers: { - [name: string]: unknown; - }; - content: { - "application/json": components["schemas"]["HTTPValidationError"]; + "application/json": { + [key: string]: string | string[]; + }; }; }; }; From 731f79fa31397ba5f1960fdea90f9fff5601db67 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:10:33 -0700 Subject: [PATCH 2/3] fix(mcp): preserve BYOK discovery and isolate session authorization --- .../mcp_server/byok_oauth_endpoints.py | 39 ++++-- .../mcp_server/discoverable_endpoints.py | 34 ++++- .../mcp_server/mcp_server_manager.py | 3 +- .../_experimental/mcp_server/oauth_utils.py | 8 ++ .../proxy/_experimental/mcp_server/server.py | 13 +- litellm/proxy/_lazy_openapi_snapshot.json | 123 ++++++++++++++++++ .../mcp_server/test_byok_oauth_endpoints.py | 89 ++++++++++++- .../mcp_server/test_discoverable_endpoints.py | 78 +++++++++-- .../mcp_server/test_mcp_hook_extra_headers.py | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 54 ++++++++ 10 files changed, 413 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index e8f0f9b8e8b..0721c2c5ace 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -18,7 +18,7 @@ import html as _html_module import time import uuid from typing import Final, cast -from urllib.parse import urlencode +from urllib.parse import urlencode, urlparse import jwt from fastapi import APIRouter, Depends, Form, HTTPException, Request @@ -26,12 +26,12 @@ from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import store_user_credential -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - get_request_base_url, -) from litellm.proxy._experimental.mcp_server.oauth_utils import ( + BYOK_RESOURCE_METADATA_PATH, TOKEN_NO_CACHE_HEADERS, + get_request_base_url, validate_loopback_redirect_uri, + well_known_root_suffix, ) from litellm.proxy._types import UserAPIKeyAuth @@ -595,13 +595,10 @@ def _build_authorize_html( # --------------------------------------------------------------------------- -@router.get("/.well-known/oauth-authorization-server", include_in_schema=False) -async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: - """RFC 8414 Authorization Server Metadata for the BYOK OAuth flow.""" - base_url: Final = get_request_base_url(request) +def _byok_authorization_server_response(base_url: str, issuer: str) -> JSONResponse: return JSONResponse( { - "issuer": base_url, + "issuer": issuer, "authorization_endpoint": f"{base_url}/v1/mcp/oauth/authorize", "token_endpoint": f"{base_url}/v1/mcp/oauth/token", "response_types_supported": ["code"], @@ -611,6 +608,30 @@ async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: ) +@router.get("/.well-known/oauth-authorization-server", include_in_schema=False) +async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: + base_url: Final = get_request_base_url(request) + return _byok_authorization_server_response(base_url, base_url) + + +@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/v1/mcp/oauth", include_in_schema=False) +async def byok_authorization_server_metadata(request: Request) -> JSONResponse: + base_url: Final = get_request_base_url(request) + return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth") + + +@router.get(BYOK_RESOURCE_METADATA_PATH, include_in_schema=False) +async def byok_protected_resource_metadata(request: Request) -> JSONResponse: + base_url: Final = get_request_base_url(request) + parsed: Final = urlparse(base_url) + return JSONResponse( + { + "resource": f"{parsed.scheme}://{parsed.netloc}", + "authorization_servers": (f"{base_url}/v1/mcp/oauth",), + } + ) + + # --------------------------------------------------------------------------- # Authorization endpoint — GET (show form) and POST (process form) # --------------------------------------------------------------------------- diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index eb24ed169de..09892eff1c6 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1842,6 +1842,30 @@ async def register_client_with_server( return JSONResponse(token_response) +@router.get("/authorize/mcp-session") +async def authorize_mcp_session( + request: Request, + redirect_uri: str, + client_id: str, + state: str = "", + code_challenge: str | None = None, + code_challenge_method: str | None = None, + response_type: str | None = None, + resource: str | None = None, +) -> Response: + return aggregate_authorize( + request=request, + client_id=client_id, + redirect_uri=redirect_uri, + state=state, + code_challenge=code_challenge, + code_challenge_method=code_challenge_method, + response_type=response_type, + session_user_id=_session_cookie_user_id(request), + resource=resource, + ) + + @router.get("/{mcp_server_name}/authorize") @router.get("/authorize") async def authorize( @@ -2538,13 +2562,13 @@ def _jwt_auth_issuers() -> list: @router.get("/.well-known/oauth-protected-resource") -def oauth_protected_resource_root(request: Request) -> dict[str, str | list[str]]: +def oauth_protected_resource_root(request: Request) -> dict[str, str | tuple[str, ...]]: request_base_url: Final = get_request_base_url(request) parsed: Final = urlparse(request_base_url) return { "resource": f"{parsed.scheme}://{parsed.netloc}", - "authorization_servers": [f"{request_base_url}/mcp"], - "scopes_supported": [], + "authorization_servers": (f"{request_base_url}/mcp",), + "scopes_supported": (), } @@ -2574,14 +2598,14 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: The issuer is ``{base}/mcp`` and must stay equal to the value the aggregate protected-resource document advertises: spec clients verify the issuer in the metadata matches the one that derived the well-known URL. - Advertises the root /authorize, /token, and /register endpoints and + Advertises the MCP session authorize endpoint, root /token and /register endpoints, and ``token_endpoint_auth_methods_supported: ["none", ...]`` because DCR clients (Claude Desktop, MCP Inspector) register as public clients; PKCE S256 is mandatory in the gateway's authorize flow.""" request_base_url: Final = get_request_base_url(request) return { "issuer": f"{request_base_url}/mcp", - "authorization_endpoint": f"{request_base_url}/authorize", + "authorization_endpoint": f"{request_base_url}/authorize/mcp-session", "token_endpoint": f"{request_base_url}/token", "introspection_endpoint": f"{request_base_url}/introspect", "registration_endpoint": f"{request_base_url}/register", diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index cc7ea0c1ea5..b46997c4d92 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -91,6 +91,7 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import ( from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, canonicalize_url_identity, + get_byok_www_authenticate, ) from litellm.proxy._experimental.mcp_server.outbound_credentials import ( Error, @@ -1232,7 +1233,7 @@ async def _resolve_byok_mcp_auth_header( "Complete the OAuth authorization flow to provide your API key." ), }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) return byok_cred diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index 1ca2ffc703d..fe35c87e07c 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -126,6 +126,14 @@ def _resolve_proxy_base_url_env() -> str | None: return None +BYOK_RESOURCE_METADATA_PATH: Final = "/v1/mcp/oauth/protected-resource" + + +def get_byok_www_authenticate() -> str: + base_url: Final = _resolve_proxy_base_url_env() or well_known_root_suffix() + return f'Bearer resource_metadata="{base_url}{BYOK_RESOURCE_METADATA_PATH}"' + + def get_request_base_url(request: Request) -> str: """ Get the base URL for the request, considering X-Forwarded-* headers. diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 43a7f894db8..fc87db69e16 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -56,6 +56,7 @@ from litellm.proxy._experimental.mcp_server.mcp_debug import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( _redact_mcp_resource_url, + get_byok_www_authenticate, get_passthrough_www_authenticate, get_route_relative_request_path, well_known_root_suffix, @@ -2852,7 +2853,7 @@ if MCP_AVAILABLE: "server_name": mcp_server.server_name or mcp_server.name, "message": "User identity is required for BYOK servers", }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) # Check shared credential cache before hitting the DB. @@ -2873,9 +2874,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) return @@ -2914,7 +2913,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'}, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) async def execute_mcp_tool( @@ -3068,9 +3067,7 @@ if MCP_AVAILABLE: "Complete the OAuth authorization flow to provide your API key." ), }, - headers={ - "WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"' - }, + headers={"WWW-Authenticate": get_byok_www_authenticate()}, ) mcp_auth_header = byok_cred elif mcp_server.is_byok: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 16681613c5d..cb7a18cd107 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -25026,6 +25026,129 @@ ] } }, + "/authorize/mcp-session": { + "get": { + "operationId": "authorize_mcp_session_authorize_mcp_session_get", + "parameters": [ + { + "in": "query", + "name": "redirect_uri", + "required": true, + "schema": { + "title": "Redirect Uri", + "type": "string" + } + }, + { + "in": "query", + "name": "client_id", + "required": true, + "schema": { + "title": "Client Id", + "type": "string" + } + }, + { + "in": "query", + "name": "state", + "required": false, + "schema": { + "default": "", + "title": "State", + "type": "string" + } + }, + { + "in": "query", + "name": "code_challenge", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Code Challenge" + } + }, + { + "in": "query", + "name": "code_challenge_method", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Code Challenge Method" + } + }, + { + "in": "query", + "name": "response_type", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Response Type" + } + }, + { + "in": "query", + "name": "resource", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Resource" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Mcp Session", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index f5bc16375fd..2b4fa572b76 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -97,6 +97,51 @@ def unauthenticated_client(): # --------------------------------------------------------------------------- +@pytest.mark.parametrize("byok_first", [True, False]) +def test_byok_challenge_discovers_api_key_flow(monkeypatch, byok_first): + from litellm.proxy._experimental.mcp_server import byok_oauth_endpoints, discoverable_endpoints + from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + app = FastAPI() + routers = (byok_oauth_endpoints.router, discoverable_endpoints.router) + for item in routers if byok_first else reversed(routers): + app.include_router(item) + with TestClient(app) as session: + challenge = get_byok_www_authenticate() + assert challenge == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"' + response = session.get(challenge.split('"')[1]) + assert response.status_code == 200 + assert response.json() == { + "resource": "http://testserver", + "authorization_servers": ["http://testserver/v1/mcp/oauth"], + } + authorization = session.get("/.well-known/oauth-authorization-server/v1/mcp/oauth") + assert authorization.status_code == 200 + metadata = authorization.json() + assert metadata["issuer"] == response.json()["authorization_servers"][0] + assert metadata["authorization_endpoint"] == "http://testserver/v1/mcp/oauth/authorize" + assert metadata["token_endpoint"] == "http://testserver/v1/mcp/oauth/token" + assert metadata["code_challenge_methods_supported"] == ["S256"] + + +@pytest.mark.parametrize( + ("base_url", "root_path", "expected"), + [ + ("", "", "/v1/mcp/oauth/protected-resource"), + ("", "/proxy", "/proxy/v1/mcp/oauth/protected-resource"), + ("https://gateway.example.com/proxy", "/proxy", "https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"), + ], +) +def test_byok_challenge_preserves_external_base(monkeypatch, base_url, root_path, expected): + from litellm.proxy._experimental.mcp_server.oauth_utils import get_byok_www_authenticate + + monkeypatch.setenv("PROXY_BASE_URL", base_url) + monkeypatch.setenv("SERVER_ROOT_PATH", root_path) + assert get_byok_www_authenticate() == f'Bearer resource_metadata="{expected}"' + + def test_oauth_authorization_server_metadata(client): resp = client.get("/.well-known/oauth-authorization-server") assert resp.status_code == 200 @@ -492,7 +537,7 @@ async def test_check_byok_credential_no_user_id(): @pytest.mark.asyncio -async def test_check_byok_credential_missing_credential(): +async def test_check_byok_credential_missing_credential(monkeypatch): """BYOK server with a known user but no stored credential → 401.""" from litellm.proxy._experimental.mcp_server.server import _check_byok_credential from litellm.proxy._types import UserAPIKeyAuth @@ -506,6 +551,11 @@ async def test_check_byok_credential_missing_credential(): ) user_auth = UserAPIKeyAuth(user_id="user-99", api_key="sk-test") + from litellm.proxy._experimental.mcp_server import server as server_module + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(server_module, "_byok_cred_cache", {}) mock_prisma = MagicMock() with ( @@ -517,6 +567,10 @@ async def test_check_byok_credential_missing_credential(): ): with pytest.raises(HTTPException) as exc_info: await _check_byok_credential(server, user_auth) + with pytest.raises(HTTPException) as cached_exc: + await _check_byok_credential(server, user_auth) + assert cached_exc.value.status_code == 401 + assert cached_exc.value.headers == exc_info.value.headers assert exc_info.value.status_code == 401 detail: Any = exc_info.value.detail @@ -524,7 +578,38 @@ async def test_check_byok_credential_missing_credential(): assert detail["server_id"] == "byok-2" headers = exc_info.value.headers or {} assert "WWW-Authenticate" in headers # type: ignore[operator] - assert "oauth-protected-resource" in headers["WWW-Authenticate"] # type: ignore[index] + assert headers["WWW-Authenticate"] == 'Bearer resource_metadata="/v1/mcp/oauth/protected-resource"' + + +@pytest.mark.asyncio +async def test_execute_byok_tool_missing_credential_advertises_api_key_flow(monkeypatch): + from datetime import datetime, timezone + + from litellm.proxy._experimental.mcp_server import server as mcp_module + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") + monkeypatch.setattr(mcp_module, "_byok_cred_cache", {}) + server = MCPServer(server_id="byok-discovery", name="byok-discovery", transport=MCPTransport.http, is_byok=True) + prisma = MagicMock() + prisma.db.litellm_mcpusercredentials.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) + with pytest.raises(HTTPException) as exc_info: + await mcp_module.execute_mcp_tool( + name="list_regions", + arguments={}, + allowed_mcp_servers=[server], + requested_server_id=server.server_id, + start_time=datetime.now(timezone.utc), + user_api_key_auth=UserAPIKeyAuth(user_id="byok-discovery-user"), + ) + assert exc_info.value.status_code == 401 + assert exc_info.value.detail["server_id"] == server.server_id + assert exc_info.value.headers == { + "WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"' + } @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 7da8010869e..c921106ed1c 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -3485,7 +3485,7 @@ def test_root_protected_resource_discovers_gateway(monkeypatch, server_count, by assert authorization.status_code == 200 metadata = authorization.json() assert metadata["issuer"] == response.json()["authorization_servers"][0] - assert metadata["authorization_endpoint"] == f"{base_url}/authorize" + assert metadata["authorization_endpoint"] == f"{base_url}/authorize/mcp-session" assert metadata["token_endpoint"] == f"{base_url}/token" assert metadata["registration_endpoint"] == f"{base_url}/register" aggregate = client.get("/.well-known/oauth-protected-resource/mcp") @@ -3508,8 +3508,8 @@ async def test_unnamed_protected_resource_builder_uses_gateway_origin(monkeypatc response = await _build_oauth_protected_resource_response(request, None, False) assert response == { "resource": "https://gateway.example.com", - "authorization_servers": ["https://gateway.example.com/mcp"], - "scopes_supported": [], + "authorization_servers": ("https://gateway.example.com/mcp",), + "scopes_supported": (), } @@ -4170,9 +4170,9 @@ async def test_discovery_root_does_not_expose_private_server_for_external_client assert "/test_oauth/" not in authorization_response["authorization_endpoint"] assert "/test_oauth/" not in authorization_response["token_endpoint"] assert authorization_response["scopes_supported"] == [] - assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",) assert resource_response["resource"] == "https://llm.example.com" - assert resource_response["scopes_supported"] == [] + assert not resource_response["scopes_supported"] finally: global_mcp_server_manager.registry.clear() @@ -9048,7 +9048,7 @@ def test_aggregate_wellknown_routes_serve_gateway_metadata(): assert asm.status_code == 200 assert asm.json()["issuer"] == "http://testserver/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/authorize" + assert asm.json()["authorization_endpoint"] == "http://testserver/authorize/mcp-session" assert "none" in asm.json()["token_endpoint_auth_methods_supported"] @@ -9137,7 +9137,7 @@ async def test_root_resource_uses_gateway_without_changing_authorization_relay() ) assert "/test_oauth/authorize" in authorization_response["authorization_endpoint"] assert authorization_response["issuer"] == "https://llm.example.com" - assert resource_response["authorization_servers"] == ["https://llm.example.com/mcp"] + assert tuple(resource_response["authorization_servers"]) == ("https://llm.example.com/mcp",) assert resource_response["resource"] == "https://llm.example.com" finally: global_mcp_server_manager.registry.clear() @@ -10670,7 +10670,7 @@ class TestPerRequestRootPathDiscovery: assert asm.status_code == 200 assert asm.json()["issuer"] == "http://testserver/tenant-a/mcp" - assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize" + assert asm.json()["authorization_endpoint"] == "http://testserver/tenant-a/authorize/mcp-session" # The prefixed authorize URL routes to the real handler (not 404): # under per-request root_path the whole app is reachable per-prefix, @@ -10829,6 +10829,68 @@ def _consent_flow_handle(page: str) -> str: return match.group(1) +@pytest.mark.parametrize("redirect_uri", ["http://127.0.0.1:51234/callback", "https://client.example.com/callback"]) +@pytest.mark.parametrize("signed_in", [True, False]) +def test_root_discovery_origin_authorizes_mcp_session(monkeypatch, redirect_uri, signed_in): + from urllib.parse import parse_qs, urlparse + + client, session_cookie, minted = _native_client_app(monkeypatch) + root = client.get("/.well-known/oauth-protected-resource") + assert root.status_code == 200 + assert root.json()["resource"] == "http://testserver" + authorization = client.get("/.well-known/oauth-authorization-server/mcp") + assert authorization.status_code == 200 + metadata = authorization.json() + registered = client.post(metadata["registration_endpoint"], json={"redirect_uris": [redirect_uri]}) + assert registered.status_code == 201 + if signed_in: + client.cookies.set("token", session_cookie) + response = client.get( + metadata["authorization_endpoint"], + params={ + "response_type": "code", + "client_id": registered.json()["client_id"], + "redirect_uri": redirect_uri, + "state": "mcp-state", + "code_challenge": _s256("v" * 43), + "code_challenge_method": "S256", + "resource": root.json()["resource"], + }, + follow_redirects=False, + ) + assert response.status_code == 303 + target = urlparse(response.headers["location"]) + assert target.path == ("/ui/connect" if signed_in else "/sso/key/generate") + if signed_in: + flow = parse_qs(target.query)["connect_flow"][0] + described = client.get("/authorize/flow", params={"flow": flow}) + assert described.status_code == 200 + assert described.json()["state"] == "unscoped" + assert minted == [] + + +@pytest.mark.parametrize("valid_client", [True, False]) +def test_mcp_session_authorize_rejects_invalid_registration_or_pkce(monkeypatch, valid_client): + client, session_cookie, minted = _native_client_app(monkeypatch) + redirect_uri = "https://client.example.com/callback" + registered = client.post("/register", json={"redirect_uris": [redirect_uri]}) + assert registered.status_code == 201 + client.cookies.set("token", session_cookie) + response = client.get( + "/authorize/mcp-session", + params={ + "client_id": registered.json()["client_id"] if valid_client else "unknown-client", + "redirect_uri": redirect_uri, + "response_type": "code", + }, + follow_redirects=False, + ) + assert response.status_code == 400 + assert response.json()["error"] == ("invalid_request" if valid_client else "invalid_client") + assert "location" not in response.headers + assert minted == [] + + def test_native_client_login_walks_discovery_consent_token_refresh_and_revoke(monkeypatch): """The whole ``lite login --pkce`` server side over the real router: a Go CLI reads the versioned discovery document, registers a loopback public client, the signed-in user consents to a team, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py index 56851d31241..67948375403 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py @@ -1233,13 +1233,14 @@ class TestResolveByokMcpAuthHeader: assert result == "stored-cred" @pytest.mark.asyncio - async def test_byok_server_raises_401_when_no_credential_stored(self): + async def test_byok_server_raises_401_when_no_credential_stored(self, monkeypatch): from fastapi import HTTPException from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _resolve_byok_mcp_auth_header, ) + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example.com/proxy") server = self._server(is_byok=True) user_auth = UserAPIKeyAuth(user_id="user-1", api_key="sk-dashboard") @@ -1252,6 +1253,9 @@ class TestResolveByokMcpAuthHeader: assert exc_info.value.status_code == 401 assert exc_info.value.detail["error"] == "byok_auth_required" + assert exc_info.value.headers == { + "WWW-Authenticate": 'Bearer resource_metadata="https://gateway.example.com/proxy/v1/mcp/oauth/protected-resource"' + } @pytest.mark.asyncio async def test_byok_server_checks_credential_and_keeps_caller_header_when_supplied(self): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ea76aeb2a36..fae668ad0f1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1177,6 +1177,23 @@ export interface paths { patch?: never; trace?: never; }; + "/authorize/mcp-session": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Authorize Mcp Session */ + get: operations["authorize_mcp_session_authorize_mcp_session_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/auto_router/benchmarks": { parameters: { query?: never; @@ -41648,6 +41665,43 @@ export interface operations { }; }; }; + authorize_mcp_session_authorize_mcp_session_get: { + parameters: { + query: { + redirect_uri: string; + client_id: string; + state?: string; + code_challenge?: string | null; + code_challenge_method?: string | null; + response_type?: string | null; + resource?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_auto_router_benchmarks_auto_router_benchmarks_get: { parameters: { query?: { From 87e2e7aface962b4ec8bb6a1b652d6f259a77eec Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:25:03 -0700 Subject: [PATCH 3/3] fix(mcp): preserve request prefixes in BYOK discovery --- .../mcp_server/byok_oauth_endpoints.py | 11 ++++++ .../_experimental/mcp_server/oauth_utils.py | 3 +- .../mcp_server/test_byok_oauth_endpoints.py | 37 +++++++++++++++++++ 3 files changed, 50 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index 0721c2c5ace..0ab76588b1f 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -34,6 +34,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( well_known_root_suffix, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.middleware.per_request_root_path_middleware import get_server_root_paths # --------------------------------------------------------------------------- # In-memory store for pending authorization codes. @@ -620,6 +621,16 @@ async def byok_authorization_server_metadata(request: Request) -> JSONResponse: return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth") +@router.get("/.well-known/oauth-authorization-server/{root_path:path}/v1/mcp/oauth", include_in_schema=False) +async def byok_prefixed_authorization_server_metadata(request: Request, root_path: str) -> JSONResponse: + prefix: Final = f"/{root_path}" + if prefix not in get_server_root_paths(): + raise HTTPException(status_code=404, detail="Unknown proxy root path") + parsed: Final = urlparse(get_request_base_url(request)) + base_url: Final = f"{parsed.scheme}://{parsed.netloc}{prefix}" + return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth") + + @router.get(BYOK_RESOURCE_METADATA_PATH, include_in_schema=False) async def byok_protected_resource_metadata(request: Request) -> JSONResponse: base_url: Final = get_request_base_url(request) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index fe35c87e07c..39865a35ec6 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -16,6 +16,7 @@ from litellm.proxy._experimental.mcp_server.auth.token_endpoint_auth import ( normalize_token_endpoint_auth_method, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.middleware.per_request_root_path_middleware import get_request_root_path if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer @@ -130,7 +131,7 @@ BYOK_RESOURCE_METADATA_PATH: Final = "/v1/mcp/oauth/protected-resource" def get_byok_www_authenticate() -> str: - base_url: Final = _resolve_proxy_base_url_env() or well_known_root_suffix() + base_url: Final = _resolve_proxy_base_url_env() or get_request_root_path().rstrip("/") return f'Bearer resource_metadata="{base_url}{BYOK_RESOURCE_METADATA_PATH}"' diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 2b4fa572b76..55accfb169d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -142,6 +142,43 @@ def test_byok_challenge_preserves_external_base(monkeypatch, base_url, root_path assert get_byok_www_authenticate() == f'Bearer resource_metadata="{expected}"' +def test_byok_discovery_preserves_per_request_prefixes(monkeypatch): + from fastapi import FastAPI + + from litellm.proxy._experimental.mcp_server.server import _check_byok_credential + from litellm.proxy.middleware.per_request_root_path_middleware import PerRequestRootPathMiddleware + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + monkeypatch.delenv("PROXY_BASE_URL", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setenv("SERVER_ROOT_PATHS", "/tenant-a,/tenant-b") + app = FastAPI() + app.include_router(router) + app.add_middleware(PerRequestRootPathMiddleware, root_paths=("/tenant-a", "/tenant-b")) + server = MCPServer(server_id="byok-prefix", name="byok-prefix", transport=MCPTransport.http, is_byok=True) + + @app.get("/challenge") + async def challenge(): + await _check_byok_credential(server, None) + + with TestClient(app) as client: + for prefix in ("/tenant-a", "/tenant-b", ""): + challenge_response = client.get(f"{prefix}/challenge") + assert challenge_response.status_code == 401 + metadata_path = f"{prefix}/v1/mcp/oauth/protected-resource" + assert challenge_response.headers["www-authenticate"] == f'Bearer resource_metadata="{metadata_path}"' + prm = client.get(metadata_path) + assert prm.status_code == 200 + issuer = f"http://testserver{prefix}/v1/mcp/oauth" + assert prm.json()["authorization_servers"] == [issuer] + asm = client.get(f"/.well-known/oauth-authorization-server{prefix}/v1/mcp/oauth") + assert asm.status_code == 200 + assert asm.json()["issuer"] == issuer + assert asm.json()["authorization_endpoint"] == f"{issuer}/authorize" + assert asm.json()["token_endpoint"] == f"{issuer}/token" + assert client.get("/.well-known/oauth-authorization-server/unknown/v1/mcp/oauth").status_code == 404 + + def test_oauth_authorization_server_metadata(client): resp = client.get("/.well-known/oauth-authorization-server") assert resp.status_code == 200