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