diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 8fa9964b1f..10a6ab8cb0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -102,7 +102,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails(): working_server if server_id == "working_server" else failing_server ) - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, mcp_auth_header=None, extra_headers=None + ): if server.name == "working_server": # Working server returns tools tool1 = MagicMock() @@ -184,7 +186,9 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing(): failing_server1 if server_id == "failing_server1" else failing_server2 ) - async def mock_get_tools_from_server(server, mcp_auth_header=None): + async def mock_get_tools_from_server( + server, mcp_auth_header=None, extra_headers=None + ): # All servers fail raise Exception(f"Server {server.name} connection failed") @@ -448,3 +452,115 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name(): assert ( called_servers[0].server_id == specific_server.server_id ), "Should have contacted the specific server alias, not the group." + + +@pytest.mark.asyncio +async def test_oauth2_headers_passed_to_mcp_client(): + """Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp""" + try: + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy._experimental.mcp_server.server import ( + _get_tools_from_mcp_servers, + set_auth_context, + ) + 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 server not available") + + # Clear the registry to avoid conflicts with other tests + global_mcp_server_manager.registry.clear() + + # Create an OAuth2 MCP server similar to github_mcp configuration + oauth2_server = MCPServer( + server_id="github_mcp_server_id", + name="github_mcp", + alias="github_mcp", + transport=MCPTransport.http, + url="https://api.githubcopilot.com/mcp", + auth_type=MCPAuth.oauth2, + client_id="test_github_client_id", + client_secret="test_github_client_secret", + scopes=["public_repo", "user:email"], + authorization_url="https://github.com/login/oauth/authorize", + token_url="https://github.com/login/oauth/access_token", + ) + global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server + + # Mock user auth + user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user") + + # Set up OAuth2 headers that would come from the client + oauth2_headers = {"Authorization": "Bearer github_oauth_token_12345"} + + # Set auth context with OAuth2 headers + set_auth_context(user_api_key_auth=user_api_key_auth, oauth2_headers=oauth2_headers) + + # This will capture the arguments passed to _create_mcp_client + captured_client_args = {} + + def mock_create_mcp_client(server, mcp_auth_header=None, extra_headers=None): + # Capture the arguments for verification + captured_client_args.update( + { + "server": server, + "mcp_auth_header": mcp_auth_header, + "extra_headers": extra_headers, + } + ) + # Return a mock client that doesn't actually connect + mock_client = MagicMock() + mock_client.disconnect = AsyncMock() + return mock_client + + # Mock _fetch_tools_with_timeout to avoid actual network calls + async def mock_fetch_tools_with_timeout(client, server_name): + return [] # Return empty list of tools + + with patch.object( + global_mcp_server_manager, + "_create_mcp_client", + side_effect=mock_create_mcp_client, + ) as mock_create_client, patch.object( + global_mcp_server_manager, + "_fetch_tools_with_timeout", + side_effect=mock_fetch_tools_with_timeout, + ), patch.object( + global_mcp_server_manager, + "get_allowed_mcp_servers", + AsyncMock(return_value=[oauth2_server.server_id]), + ): + # Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client + await _get_tools_from_mcp_servers( + user_api_key_auth=user_api_key_auth, + mcp_auth_header=None, + mcp_servers=None, # Will use all allowed servers + oauth2_headers=oauth2_headers, + ) + + # Verify that _create_mcp_client was called + assert ( + mock_create_client.call_count == 1 + ), "Expected _create_mcp_client to be called once" + + # Verify the server passed to _create_mcp_client is the OAuth2 server + assert captured_client_args["server"].server_id == oauth2_server.server_id + assert captured_client_args["server"].auth_type == MCPAuth.oauth2 + + # Most importantly: verify that OAuth2 headers were passed as extra_headers + assert ( + captured_client_args["extra_headers"] is not None + ), "Expected extra_headers to be passed for OAuth2 server" + assert ( + captured_client_args["extra_headers"] == oauth2_headers + ), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}" + + # Verify the Authorization header specifically + assert "Authorization" in captured_client_args["extra_headers"] + assert ( + captured_client_args["extra_headers"]["Authorization"] + == "Bearer github_oauth_token_12345" + )