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] 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[]; + }; }; }; };