"""Tests for MCP OAuth discoverable endpoints""" from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @pytest.fixture(autouse=True) def mock_mcp_client_ip(): """Mock IPAddressUtils.get_mcp_client_ip to return None for all tests. This bypasses IP-based access control in tests, since the MCP server's available_on_public_internet defaults to False and mock requests don't have proper client IP context. """ with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip", return_value=None, ): yield @pytest.mark.asyncio async def test_authorize_endpoint_includes_response_type(): """Test that authorize endpoint includes response_type=code parameter (fixes #15684)""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} # Mock the encryption functions to avoid needing a signing key with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint response = await authorize( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", redirect_uri="https://client.example.com/callback", state="test_state", ) # Verify response is a redirect assert response.status_code == 307 # FastAPI RedirectResponse default # Verify response_type is in the redirect URL assert "response_type=code" in response.headers["location"] assert "https://provider.com/oauth/authorize" in response.headers["location"] assert "client_id=test_client_id" in response.headers["location"] assert "scope=read+write" in response.headers["location"] @pytest.mark.asyncio async def test_authorize_endpoint_preserves_existing_query_params(): """Test that authorize endpoint merges OAuth params with existing query params in authorization_url""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) 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") global_mcp_server_manager.registry.clear() # Authorization URL already has query params (e.g. multi-tenant OAuth) oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize?tenant=system", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" response = await authorize( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", redirect_uri="https://client.example.com/callback", state="test_state", ) location = response.headers["location"] # Must NOT have double '?' — existing params must be merged correctly assert location.count("?") == 1, ( f"Expected exactly one '?' in URL but got {location.count('?')}: {location}" ) assert "tenant=system" in location assert "client_id=test_client_id" in location assert "response_type=code" in location assert "scope=read+write" in location @pytest.mark.asyncio async def test_authorize_endpoint_forwards_pkce_parameters(): """Test that authorize endpoint forwards PKCE parameters (code_challenge and code_challenge_method)""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server (simulating Google OAuth) oauth2_server = MCPServer( server_id="google_mcp", name="google_mcp", server_name="google_mcp", alias="google_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="669428968603-test.apps.googleusercontent.com", client_secret="GOCSPX-test_secret", authorization_url="https://accounts.google.com/o/oauth2/v2/auth", token_url="https://oauth2.googleapis.com/token", scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm-proxy.example.com/" mock_request.headers = {} # Mock the encryption function with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state_with_pkce" # Call authorize endpoint with PKCE parameters response = await authorize( request=mock_request, client_id="669428968603-test.apps.googleusercontent.com", mcp_server_name="google_mcp", redirect_uri="http://localhost:60108/callback", state="test_client_state", code_challenge="x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk", code_challenge_method="S256", ) # Verify response is a redirect assert response.status_code == 307 # Verify PKCE parameters are included in the redirect URL location = response.headers["location"] assert "https://accounts.google.com/o/oauth2/v2/auth" in location assert "code_challenge=x6YH_qgwbvOzbsHDuL1sW9gYkR9-gObUiIB5RkPwxDk" in location assert "code_challenge_method=S256" in location assert "client_id=669428968603-test.apps.googleusercontent.com" in location assert "response_type=code" in location @pytest.mark.asyncio async def test_token_endpoint_forwards_code_verifier(): """Test that token endpoint forwards code_verifier for PKCE flow""" try: import httpx from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( token_endpoint, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="google_mcp", name="google_mcp", server_name="google_mcp", alias="google_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="669428968603-test.apps.googleusercontent.com", client_secret="GOCSPX-test_secret", authorization_url="https://accounts.google.com/o/oauth2/v2/auth", token_url="https://oauth2.googleapis.com/token", scopes=["https://www.googleapis.com/auth/drive", "openid", "email"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm-proxy.example.com/" mock_request.headers = {} # Mock httpx client response mock_response = MagicMock() mock_response.json.return_value = { "access_token": "ya29.test_access_token", "token_type": "Bearer", "expires_in": 3599, "scope": "openid email https://www.googleapis.com/auth/drive", } mock_response.raise_for_status = MagicMock() # Mock the async httpx client with AsyncMock for async methods from unittest.mock import AsyncMock with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" ) as mock_get_client: mock_async_client = MagicMock() # Use AsyncMock for the async post method mock_async_client.post = AsyncMock(return_value=mock_response) mock_get_client.return_value = mock_async_client # Call token endpoint with code_verifier response = await token_endpoint( request=mock_request, grant_type="authorization_code", code="4/test_authorization_code", redirect_uri="http://localhost:60108/callback", client_id="669428968603-test.apps.googleusercontent.com", mcp_server_name="google_mcp", client_secret="GOCSPX-test_secret", code_verifier="test_code_verifier_from_client", ) # Verify that the token endpoint was called with code_verifier mock_async_client.post.assert_called_once() call_args = mock_async_client.post.call_args # Check the data parameter includes code_verifier assert call_args[1]["data"]["code_verifier"] == "test_code_verifier_from_client" assert call_args[1]["data"]["code"] == "4/test_authorization_code" assert ( call_args[1]["data"]["client_id"] == "669428968603-test.apps.googleusercontent.com" ) assert call_args[1]["data"]["client_secret"] == "GOCSPX-test_secret" assert call_args[1]["data"]["grant_type"] == "authorization_code" # Verify response response_data = response.body import json token_data = json.loads(response_data) assert token_data["access_token"] == "ya29.test_access_token" assert token_data["token_type"] == "Bearer" @pytest.mark.asyncio async def test_register_client_without_mcp_server_name_returns_dummy(): try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") # Clear registry to ensure no OAuth2 servers exist (otherwise resolver would find one) global_mcp_server_manager.registry.clear() mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): result = await register_client(request=mock_request) assert result == { "client_id": "dummy_client", "client_secret": "dummy", "redirect_uris": ["https://proxy.litellm.example/callback"], } @pytest.mark.asyncio async def test_register_client_returns_existing_server_credentials(): try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) 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") global_mcp_server_manager.registry.clear() oauth2_server = MCPServer( server_id="stored_server", name="stored_server", server_name="stored_server", alias="stored_server", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="existing-client", client_secret="existing-secret", authorization_url="https://provider.example/oauth/authorize", token_url="https://provider.example/oauth/token", ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} try: with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): result = await register_client( request=mock_request, mcp_server_name=oauth2_server.server_name ) finally: global_mcp_server_manager.registry.clear() assert result == { "client_id": "stored_server", "client_secret": "dummy", "redirect_uris": ["https://proxy.litellm.example/callback"], } @pytest.mark.asyncio async def test_register_client_remote_registration_success(): try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) 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") global_mcp_server_manager.registry.clear() oauth2_server = MCPServer( server_id="remote_server", name="remote_server", server_name="remote_server", alias="remote_server", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id=None, client_secret=None, authorization_url="https://provider.example/oauth/authorize", token_url="https://provider.example/oauth/token", registration_url="https://provider.example/oauth/register", ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://proxy.litellm.example/" mock_request.headers = {} request_payload = { "client_name": "Litellm Proxy", "grant_types": ["authorization_code", "refresh_token"], "response_types": ["code"], "token_endpoint_auth_method": "client_secret_post", } mock_response = MagicMock() mock_response.json.return_value = { "client_id": "generated-client", "client_secret": "generated-secret", } mock_response.raise_for_status = MagicMock() mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) try: with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value=request_payload), ), patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client", return_value=mock_async_client, ): response = await register_client( request=mock_request, mcp_server_name=oauth2_server.server_name ) finally: global_mcp_server_manager.registry.clear() import json assert response.status_code == 200 payload = json.loads(response.body.decode("utf-8")) assert payload == mock_response.json.return_value mock_async_client.post.assert_called_once() call_args = mock_async_client.post.call_args assert call_args.args[0] == oauth2_server.registration_url assert call_args.kwargs["headers"] == { "Content-Type": "application/json", "Accept": "application/json", } assert call_args.kwargs["json"]["redirect_uris"] == [ "https://proxy.litellm.example/callback" ] assert call_args.kwargs["json"]["grant_types"] == request_payload["grant_types"] assert ( call_args.kwargs["json"]["token_endpoint_auth_method"] == request_payload["token_endpoint_auth_method"] ) @pytest.mark.asyncio async def test_authorize_endpoint_respects_x_forwarded_proto(): """Test that authorize endpoint uses X-Forwarded-Proto header to construct correct redirect_uri""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy # Mock the encryption functions with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint response = await authorize( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", redirect_uri="https://client.example.com/callback", state="test_state", ) # Verify redirect URL uses HTTPS in the redirect_uri parameter location = response.headers["location"] # The redirect_uri parameter sent to the OAuth provider should use HTTPS assert ( "redirect_uri=https%3A%2F%2Flitellm.example.com%2Fcallback" in location or "redirect_uri=https://litellm.example.com/callback" in location ) @pytest.mark.asyncio async def test_token_endpoint_respects_x_forwarded_proto(): """Test that token endpoint uses X-Forwarded-Proto header for redirect_uri""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( token_endpoint, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="google_mcp", name="google_mcp", server_name="google_mcp", alias="google_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_secret", authorization_url="https://accounts.google.com/o/oauth2/v2/auth", token_url="https://oauth2.googleapis.com/token", scopes=["openid", "email"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm-proxy.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy # Mock httpx client response mock_response = MagicMock() mock_response.json.return_value = { "access_token": "test_token", "token_type": "Bearer", "expires_in": 3599, } mock_response.raise_for_status = MagicMock() # Mock the async httpx client mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" ) as mock_get_client: mock_get_client.return_value = mock_async_client # Call token endpoint response = await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", redirect_uri="http://localhost:60108/callback", client_id="test_client_id", mcp_server_name="google_mcp", client_secret="test_secret", ) # Verify that the redirect_uri sent to the provider uses HTTPS call_args = mock_async_client.post.call_args assert ( call_args[1]["data"]["redirect_uri"] == "https://litellm-proxy.example.com/callback" ) @pytest.mark.asyncio async def test_oauth_protected_resource_respects_x_forwarded_proto(): """Test that oauth_protected_resource_mcp uses X-Forwarded-Proto for URLs""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( oauth_protected_resource_mcp, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy # Call the endpoint response = await oauth_protected_resource_mcp( request=mock_request, mcp_server_name="test_oauth", ) # Verify response uses HTTPS URLs assert response["authorization_servers"][0].startswith( "https://litellm.example.com/" ) assert response["scopes_supported"] == oauth2_server.scopes @pytest.mark.asyncio async def test_oauth_authorization_server_respects_x_forwarded_proto(): """Test that oauth_authorization_server_mcp uses X-Forwarded-Proto for URLs""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( oauth_authorization_server_mcp, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://litellm.example.com/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy # Call the endpoint response = await oauth_authorization_server_mcp( request=mock_request, mcp_server_name="test_oauth", ) # Verify response uses HTTPS URLs assert response["authorization_endpoint"].startswith("https://litellm.example.com/") assert response["token_endpoint"].startswith("https://litellm.example.com/") assert response["registration_endpoint"].startswith("https://litellm.example.com/") assert response["grant_types_supported"] == ["authorization_code", "refresh_token"] assert response["scopes_supported"] == oauth2_server.scopes @pytest.mark.asyncio async def test_register_client_respects_x_forwarded_proto(): """Test that register_client uses X-Forwarded-Proto for redirect_uris""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") # Clear registry to ensure no OAuth2 servers exist (otherwise resolver would find one) global_mcp_server_manager.registry.clear() # Mock request with http base_url but X-Forwarded-Proto: https mock_request = MagicMock(spec=Request) mock_request.base_url = "http://proxy.litellm.example/" # HTTP mock_request.headers = {"X-Forwarded-Proto": "https"} # Behind HTTPS proxy with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): result = await register_client(request=mock_request) # Verify the redirect_uris use HTTPS assert result == { "client_id": "dummy_client", "client_secret": "dummy", "redirect_uris": ["https://proxy.litellm.example/callback"], } @pytest.mark.asyncio async def test_authorize_endpoint_respects_x_forwarded_host(): """Test that authorize endpoint uses X-Forwarded-Host and X-Forwarded-Proto to construct correct redirect_uri""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_client_secret", authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request simulating nginx proxy: # Internal: http://localhost:8888/github/mcp # External: https://proxy.example.com/github/mcp mock_request = MagicMock(spec=Request) mock_request.base_url = "http://localhost:8888/github/mcp" mock_request.headers = { "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } # Mock the encryption functions with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call authorize endpoint response = await authorize( request=mock_request, client_id="test_client_id", mcp_server_name="test_oauth", redirect_uri="https://client.example.com/callback", state="test_state", ) # Verify redirect URL uses the forwarded host and scheme location = response.headers["location"] # The redirect_uri parameter should use the external URL assert ( "redirect_uri=https%3A%2F%2Fproxy.example.com%2Fgithub%2Fmcp%2Fcallback" in location or "redirect_uri=https://proxy.example.com/github/mcp/callback" in location ) @pytest.mark.asyncio async def test_token_endpoint_respects_x_forwarded_host(): """Test that token endpoint uses X-Forwarded-Host and X-Forwarded-Proto for redirect_uri""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( token_endpoint, ) 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") # Clear registry global_mcp_server_manager.registry.clear() # Create mock OAuth2 server oauth2_server = MCPServer( server_id="google_mcp", name="google_mcp", server_name="google_mcp", alias="google_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="test_client_id", client_secret="test_secret", authorization_url="https://accounts.google.com/o/oauth2/v2/auth", token_url="https://oauth2.googleapis.com/token", scopes=["openid", "email"], ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server # Mock request simulating nginx proxy without port in host mock_request = MagicMock(spec=Request) mock_request.base_url = "http://localhost:8888/github/mcp" mock_request.headers = { "X-Forwarded-Proto": "https", "X-Forwarded-Host": "proxy.example.com", } # Mock httpx client response mock_response = MagicMock() mock_response.json.return_value = { "access_token": "test_token", "token_type": "Bearer", "expires_in": 3599, } mock_response.raise_for_status = MagicMock() # Mock the async httpx client mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" ) as mock_get_client: mock_get_client.return_value = mock_async_client # Call token endpoint response = await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_code", redirect_uri="http://localhost:60108/callback", client_id="test_client_id", mcp_server_name="google_mcp", client_secret="test_secret", ) # Verify that the redirect_uri sent to the provider uses the external URL call_args = mock_async_client.post.call_args assert ( call_args[1]["data"]["redirect_uri"] == "https://proxy.example.com/github/mcp/callback" ) @pytest.mark.parametrize( "base_url,x_forwarded_proto,x_forwarded_host,x_forwarded_port,expected_url", [ # Case 1: No forwarded headers - use original URL as-is (no trailing slash) ( "http://localhost:4000/", None, None, None, "http://localhost:4000", ), # Case 2: Only X-Forwarded-Proto - change scheme only ( "http://localhost:4000/", "https", None, None, "https://localhost:4000", ), # Case 3: X-Forwarded-Proto + X-Forwarded-Host - change scheme and host ( "http://localhost:4000/", "https", "proxy.example.com", None, "https://proxy.example.com", ), # Case 4: X-Forwarded-Host with port included in host header ( "http://localhost:4000/", "https", "proxy.example.com:8080", None, "https://proxy.example.com:8080", ), # Case 5: X-Forwarded-Host + X-Forwarded-Port as separate headers ( "http://localhost:4000/", "https", "proxy.example.com", "8443", "https://proxy.example.com:8443", ), # Case 6: Only X-Forwarded-Host without proto - use original scheme ( "http://localhost:4000/", None, "proxy.example.com", None, "http://proxy.example.com", ), # Case 7: Only X-Forwarded-Port without host - preserves original port if present # (This is safer behavior - X-Forwarded-Port alone is unusual) ( "http://localhost:4000/", None, None, "8443", "http://localhost:4000", # Original port preserved when already present ), # Case 8: Complex internal URL with path (path is preserved) ( "http://localhost:8888/github/mcp", "https", "proxy.example.com", None, "https://proxy.example.com/github/mcp", ), # Case 9: IPv6 address in X-Forwarded-Host (should not treat :: as port separator) ( "http://localhost:4000/", "https", "[2001:db8::1]", None, "https://[2001:db8::1]", ), # Case 10: IPv6 address with port ( "http://localhost:4000/", "https", "[2001:db8::1]:8080", None, "https://[2001:db8::1]:8080", ), # Case 11: X-Forwarded-Host already has port, X-Forwarded-Port also provided (host wins) ( "http://localhost:4000/", "https", "proxy.example.com:9000", "8443", "https://proxy.example.com:9000", ), # Case 12: Standard proxy setup (most common case) ( "http://127.0.0.1:8888/", "https", "chatproxy.company.com", None, "https://chatproxy.company.com", ), # Case 13: Internal URL already has port, X-Forwarded-Port does NOT override # (safer behavior - preserves original port when X-Forwarded-Host not provided) ( "http://localhost:4000/", None, None, "443", "http://localhost:4000", # Original port preserved ), # Case 14: Original URL with existing port in netloc, X-Forwarded-Host replaces it ( "http://internal.local:8888/", "https", "external.com", None, "https://external.com", ), ], ) def test_get_request_base_url_comprehensive( base_url, x_forwarded_proto, x_forwarded_host, x_forwarded_port, expected_url ): """Comprehensive test for get_request_base_url with various header combinations""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") # Create mock request mock_request = MagicMock(spec=Request) mock_request.base_url = base_url # Build headers dict headers = {} if x_forwarded_proto: headers["X-Forwarded-Proto"] = x_forwarded_proto if x_forwarded_host: headers["X-Forwarded-Host"] = x_forwarded_host if x_forwarded_port: headers["X-Forwarded-Port"] = x_forwarded_port # Mock headers.get() to return our test values def mock_get(header_name, default=None): return headers.get(header_name, default) mock_request.headers.get = mock_get # Test the function result = get_request_base_url(mock_request) # Verify result assert result == expected_url, ( f"Expected '{expected_url}' but got '{result}'\n" f"Input: base_url={base_url}, " f"X-Forwarded-Proto={x_forwarded_proto}, " f"X-Forwarded-Host={x_forwarded_host}, " f"X-Forwarded-Port={x_forwarded_port}" ) # ------------------------------------------------------------------- # Tests for scopes_supported when mcp_server.scopes is None # ------------------------------------------------------------------- @pytest.mark.asyncio async def test_oauth_protected_resource_returns_empty_scopes_when_none(): """ When an MCP server exists but has scopes=None (e.g. Atlassian OAuth), scopes_supported should be [] not None. """ try: 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") global_mcp_server_manager.registry.clear() # Create an OAuth2 server with scopes=None (like Atlassian) oauth2_server = MCPServer( server_id="atlassian_mcp", name="atlassian_mcp", server_name="atlassian_mcp", alias="atlassian_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="atlassian_client_id", client_secret="atlassian_secret", authorization_url="https://auth.atlassian.com/authorize", token_url="https://auth.atlassian.com/oauth/token", scopes=None, # Atlassian doesn't set scopes ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} try: response = _build_oauth_protected_resource_response( request=mock_request, mcp_server_name="atlassian_mcp", use_standard_pattern=False, ) assert response["scopes_supported"] == [] finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_oauth_authorization_server_returns_empty_scopes_when_none(): """ When an MCP server exists but has scopes=None (e.g. Atlassian OAuth), scopes_supported should be [] not None. """ try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_authorization_server_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") global_mcp_server_manager.registry.clear() # Create an OAuth2 server with scopes=None oauth2_server = MCPServer( server_id="atlassian_mcp", name="atlassian_mcp", server_name="atlassian_mcp", alias="atlassian_mcp", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id="atlassian_client_id", client_secret="atlassian_secret", authorization_url="https://auth.atlassian.com/authorize", token_url="https://auth.atlassian.com/oauth/token", scopes=None, ) global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} try: response = _build_oauth_authorization_server_response( request=mock_request, mcp_server_name="atlassian_mcp", ) assert response["scopes_supported"] == [] finally: global_mcp_server_manager.registry.clear() # ------------------------------------------------------------------- # Tests for root-level OAuth endpoint resolution (no server name) # ------------------------------------------------------------------- def _create_oauth2_server( server_id="test_oauth_server", name="test_oauth", server_name="test_oauth", alias="test_oauth", client_id="test_client_id", client_secret="test_client_secret", ): """Helper to create a mock OAuth2 MCPServer.""" from litellm.proxy._types import MCPTransport from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPServer return MCPServer( server_id=server_id, name=name, server_name=server_name, alias=alias, transport=MCPTransport.http, auth_type=MCPAuth.oauth2, client_id=client_id, client_secret=client_secret, authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", scopes=["read", "write"], ) @pytest.mark.asyncio async def test_authorize_root_resolves_single_oauth2_server(): """When /authorize is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://llm.example.com/" mock_request.headers = {} try: with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "mocked_encrypted_state" # Call /authorize WITHOUT mcp_server_name, with dummy_client as client_id response = await authorize( request=mock_request, client_id="dummy_client", mcp_server_name=None, redirect_uri="http://localhost:62646/callback", state="test_state", ) # Should resolve to the single OAuth2 server and redirect assert response.status_code == 307 location = response.headers["location"] assert "https://provider.com/oauth/authorize" in location assert "client_id=test_client_id" in location finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_authorize_root_fails_with_multiple_oauth2_servers(): """When /authorize is hit without server name and multiple OAuth2 servers exist, return 404.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() server1 = _create_oauth2_server( server_id="server1", name="server1", server_name="server1", alias="server1" ) server2 = _create_oauth2_server( server_id="server2", name="server2", server_name="server2", alias="server2" ) global_mcp_server_manager.registry[server1.server_id] = server1 global_mcp_server_manager.registry[server2.server_id] = server2 mock_request = MagicMock(spec=Request) mock_request.base_url = "https://llm.example.com/" mock_request.headers = {} try: with pytest.raises(HTTPException) as exc_info: await authorize( request=mock_request, client_id="dummy_client", mcp_server_name=None, redirect_uri="http://localhost:62646/callback", state="test_state", ) assert exc_info.value.status_code == 404 assert "MCP server not found" in str(exc_info.value.detail) finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_token_root_resolves_single_oauth2_server(): """When /token is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( token_endpoint, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://llm.example.com/" mock_request.headers = {} mock_response = MagicMock() mock_response.json.return_value = { "access_token": "ya29.test_token", "token_type": "Bearer", "expires_in": 3599, } mock_response.raise_for_status = MagicMock() mock_async_client = MagicMock() mock_async_client.post = AsyncMock(return_value=mock_response) try: with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.get_async_httpx_client" ) as mock_get_client: mock_get_client.return_value = mock_async_client # Call /token WITHOUT mcp_server_name response = await token_endpoint( request=mock_request, grant_type="authorization_code", code="test_auth_code", redirect_uri="http://localhost:62646/callback", client_id="dummy_client", mcp_server_name=None, client_secret=None, code_verifier="test_verifier", ) # Should resolve and exchange token with the upstream server import json token_data = json.loads(response.body) assert token_data["access_token"] == "ya29.test_token" # Verify it called the correct upstream token URL call_args = mock_async_client.post.call_args assert call_args.args[0] == "https://provider.com/oauth/token" finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_register_root_resolves_single_oauth2_server(): """When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( register_client, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://llm.example.com/" mock_request.headers = {} try: with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body", new=AsyncMock(return_value={}), ): result = await register_client(request=mock_request, mcp_server_name=None) # Should resolve to the single server and return its name as client_id assert result["client_id"] == "test_oauth" assert "redirect_uris" in result finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_discovery_root_includes_server_name_prefix(): """When root discovery is hit and exactly 1 OAuth2 server exists, include server name in URLs.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( _build_oauth_authorization_server_response, ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") global_mcp_server_manager.registry.clear() oauth2_server = _create_oauth2_server() global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server mock_request = MagicMock(spec=Request) mock_request.base_url = "https://llm.example.com/" mock_request.headers = {} try: # Call with mcp_server_name=None (root discovery) response = _build_oauth_authorization_server_response( request=mock_request, mcp_server_name=None, ) # Should resolve to the single server and include its name in endpoint URLs assert "/test_oauth/authorize" in response["authorization_endpoint"] assert "/test_oauth/token" in response["token_endpoint"] assert "/test_oauth/register" in response["registration_endpoint"] assert response["scopes_supported"] == ["read", "write"] finally: global_mcp_server_manager.registry.clear() @pytest.mark.asyncio async def test_oauth_callback_redirects_with_state(): """Test OAuth callback endpoint properly decodes state and redirects to client callback URL.""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( callback, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") # Mock the state decoding mock_state_data = { "base_url": "http://localhost:3000/ui/mcp/oauth/callback", "original_state": "test-uuid-state-123", "code_challenge": "test_challenge", "code_challenge_method": "S256", "client_redirect_uri": "http://localhost:3000/ui/mcp/oauth/callback", } with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: mock_decode.return_value = mock_state_data # Call callback endpoint with code and state response = await callback( code="test_authorization_code_12345", state="encrypted_state_value", ) # Should redirect to the client callback URL with code and original state assert response.status_code == 302 assert "http://localhost:3000/ui/mcp/oauth/callback" in response.headers["location"] assert "code=test_authorization_code_12345" in response.headers["location"] assert "state=test-uuid-state-123" in response.headers["location"] # Verify state was decoded mock_decode.assert_called_once_with("encrypted_state_value") @pytest.mark.asyncio async def test_oauth_callback_handles_invalid_state(): """Test OAuth callback returns error page when state decryption fails.""" try: from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( callback, ) except ImportError: pytest.skip("MCP discoverable endpoints not available") # Mock state decoding to raise an exception with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" ) as mock_decode: mock_decode.side_effect = Exception("Failed to decrypt state") # Call callback endpoint with invalid state response = await callback( code="test_code", state="invalid_encrypted_state", ) # Should return HTML error page assert response.status_code == 200 assert "Authentication incomplete" in response.body.decode() @pytest.mark.asyncio async def test_oauth_authorize_includes_scopes_from_server_config(): """Test that authorize endpoint includes scopes from server configuration.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize_with_server, ) 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") # Create server with specific scopes (e.g., GitLab requires 'ai_workflows') oauth_server = MCPServer( server_id="gitlab_server", name="gitlab", server_name="gitlab", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, authorization_url="https://gitlab.com/oauth/authorize", token_url="https://gitlab.com/oauth/token", client_id="test_client", scopes=["api", "read_user", "ai_workflows"], # GitLab-specific scopes ) mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "encrypted_state" # Call authorize without explicit scope parameter response = await authorize_with_server( request=mock_request, mcp_server=oauth_server, client_id="test_client", redirect_uri="http://localhost:3000/callback", state="test_state", code_challenge="test_challenge", code_challenge_method="S256", response_type="code", scope=None, # No scope in request, should use server's scopes ) # Should redirect with scopes from server config assert response.status_code in (307, 302) redirect_url = response.headers["location"] assert "scope=api+read_user+ai_workflows" in redirect_url or "scope=api%20read_user%20ai_workflows" in redirect_url @pytest.mark.asyncio async def test_oauth_authorize_prefers_request_scope_over_server_config(): """Test that explicit scope parameter takes precedence over server configuration.""" try: from fastapi import Request from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( authorize_with_server, ) 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") oauth_server = MCPServer( server_id="test_server", name="test", server_name="test", transport=MCPTransport.http, auth_type=MCPAuth.oauth2, authorization_url="https://provider.com/oauth/authorize", token_url="https://provider.com/oauth/token", client_id="test_client", scopes=["default_scope1", "default_scope2"], ) mock_request = MagicMock(spec=Request) mock_request.base_url = "https://litellm.example.com/" mock_request.headers = {} with patch( "litellm.proxy._experimental.mcp_server.discoverable_endpoints.encrypt_value_helper" ) as mock_encrypt: mock_encrypt.return_value = "encrypted_state" # Call authorize WITH explicit scope parameter response = await authorize_with_server( request=mock_request, mcp_server=oauth_server, client_id="test_client", redirect_uri="http://localhost:3000/callback", state="test_state", code_challenge="test_challenge", code_challenge_method="S256", response_type="code", scope="custom_scope1 custom_scope2", # Explicit scope should take precedence ) # Should use the explicit scope, not server config assert response.status_code in (307, 302) redirect_url = response.headers["location"] assert "scope=custom_scope1+custom_scope2" in redirect_url or "scope=custom_scope1%20custom_scope2" in redirect_url assert "default_scope" not in redirect_url