From b8fe5f7b17ff5604f3e76ac6bcea391de7af0db7 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Thu, 14 Aug 2025 16:32:18 -0700 Subject: [PATCH] [MCP Gateway] LiteLLM Fix MCP gateway key auth (#13630) * Fix - add safe divide by 0 for most places to prevent crash * Enhance MCPRequestHandler to support permission inheritance and intersection logic for access groups. Added integration tests to verify behavior when keys have no permissions and when both keys and teams have overlapping permissions. * Remove redundant assertions for permission checks in test_user_api_key_auth_mcp.py to streamline test logic. * Refactor integration tests for MCPRequestHandler to simplify mocking. Replace complex database mocks with direct function mocks for permission inheritance and intersection scenarios, improving test clarity and maintainability. * Revert "Fix - add safe divide by 0 for most places to prevent crash" This reverts commit 265d40e39051e148996b9fb7f354730c57ff23ac. --- .../mcp_server/auth/user_api_key_auth_mcp.py | 13 ++- tests/mcp_tests/test_mcp_server.py | 106 ++++++++++++++++++ .../auth/test_user_api_key_auth_mcp.py | 80 +++++++++++++ 3 files changed, 195 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7469848e2f2..a075de13fb1 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -236,12 +236,17 @@ class MCPRequestHandler: ) ######################################################### - # If team has mcp_servers, then key must have a subset of the team's mcp_servers + # If team has mcp_servers, handle inheritance and intersection logic ######################################################### if len(allowed_mcp_servers_for_team) > 0: - for _mcp_server in allowed_mcp_servers_for_key: - if _mcp_server in allowed_mcp_servers_for_team: - allowed_mcp_servers.append(_mcp_server) + if len(allowed_mcp_servers_for_key) > 0: + # Key has its own MCP permissions - use intersection with team permissions + for _mcp_server in allowed_mcp_servers_for_key: + if _mcp_server in allowed_mcp_servers_for_team: + allowed_mcp_servers.append(_mcp_server) + else: + # Key has no MCP permissions - inherit from team + allowed_mcp_servers = allowed_mcp_servers_for_team else: allowed_mcp_servers = allowed_mcp_servers_for_key diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 43130e62d0d..8a390412495 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1773,3 +1773,109 @@ async def test_list_tool_rest_api_all_servers_with_auth(): assert calls[1][0][2] == "2025-06-18" # mcp_protocol_version +@pytest.mark.asyncio +async def test_mcp_access_group_permission_inheritance_integration(): + """Integration test for MCP access group permission inheritance""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._types import UserAPIKeyAuth + + # Test scenario: team has access groups, key has no permissions -> should inherit + # Use direct mocking of the helper functions instead of complex database mocking + with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key: + with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team: + # Key has no permissions, team has servers + mock_key.return_value = [] # Key inherits nothing directly + mock_team.return_value = ["staff-server-1", "staff-server-2", "ops-server-1"] # Team has servers + + # Create user auth object + user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="team-staff", + object_permission_id=None # Key has no explicit permissions + ) + + # Test the inheritance logic + allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_auth) + + # Should inherit all team servers since key has no permissions + expected_servers = ["staff-server-1", "staff-server-2", "ops-server-1"] + assert sorted(allowed_servers) == sorted(expected_servers) + + +@pytest.mark.asyncio +async def test_mcp_access_group_permission_intersection_integration(): + """Integration test for MCP access group permission intersection""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._types import UserAPIKeyAuth + + # Test scenario: both team and key have access groups -> should intersect + # Use direct mocking of the helper functions instead of complex database mocking + with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key: + with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team: + # Both key and team have permissions - should intersect + mock_key.return_value = ["ops-server", "external-server"] # Key has these servers + mock_team.return_value = ["staff-server", "ops-server", "admin-server"] # Team has these servers + + # Create user auth object + user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="team-staff", + object_permission_id="key-permission-id" # Key has explicit permissions + ) + + # Test the intersection logic + allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_auth) + + # Should only get intersection (ops-server is common) + expected_servers = ["ops-server"] + assert sorted(allowed_servers) == sorted(expected_servers) + + +@pytest.mark.asyncio +async def test_mcp_server_manager_with_access_groups_integration(): + """Integration test for MCPServerManager with access group filtering""" + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._types import UserAPIKeyAuth + + # Create a test manager + test_manager = MCPServerManager() + + # Load servers with access groups + test_manager.load_servers_from_config({ + "staff_server": { + "url": "https://staff-server.com/mcp", + "access_groups": ["staff"], + "transport": MCPTransport.http, + }, + "ops_server": { + "url": "https://ops-server.com/mcp", + "access_groups": ["ops"], + "transport": MCPTransport.http, + }, + "admin_server": { + "url": "https://admin-server.com/mcp", + "access_groups": ["admin"], + "transport": MCPTransport.http, + } + }) + + # Mock user with specific access groups + user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="team-staff" + ) + + # Mock the permission lookup to return staff access group + with patch.object(MCPRequestHandler, "get_allowed_mcp_servers") as mock_get_allowed: + mock_get_allowed.return_value = ["staff-server-id", "ops-server-id"] # User has access to staff and ops + + allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth) + + # Should only get servers user has access to + assert len(allowed_servers) >= 0 # At least verify no errors + mock_get_allowed.assert_called_once_with(user_auth) + + diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 2249b9a6fa9..a9f1f8b12d2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -118,6 +118,86 @@ class TestMCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.object_permission_id: mock_find_unique.assert_not_called() + @pytest.mark.parametrize( + "team_servers,key_servers,expected_servers,scenario", + [ + # Test case 1: Key has no permissions, should inherit from team + (["server1", "server2"], [], ["server1", "server2"], "inherit_from_team"), + # Test case 2: Key has permissions, should use intersection with team + (["server1", "server2", "server3"], ["server2", "server4"], ["server2"], "intersection_logic"), + # Test case 3: Key has permissions but no overlap with team + (["server1", "server2"], ["server3", "server4"], [], "no_overlap"), + # Test case 4: Team has no permissions, use key permissions + ([], ["server1", "server2"], ["server1", "server2"], "no_team_permissions"), + # Test case 5: Both team and key have no permissions + ([], [], [], "no_permissions"), + # Test case 6: Team has permissions, key has subset + (["server1", "server2", "server3"], ["server1", "server3"], ["server1", "server3"], "key_subset"), + # Test case 7: Team has permissions, key has superset (intersection should limit) + (["server1", "server2"], ["server1", "server2", "server3"], ["server1", "server2"], "key_superset"), + ], + ) + async def test_get_allowed_mcp_servers_inheritance_logic( + self, team_servers, key_servers, expected_servers, scenario + ): + """Test the inheritance and intersection logic in get_allowed_mcp_servers""" + + # Create mock user_api_key_auth + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team" if team_servers else None, + object_permission_id="test-permission" if key_servers else None + ) + + # Mock the helper functions + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key_servers: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_team" + ) as mock_team_servers: + + # Configure mocks to return the test data + mock_key_servers.return_value = key_servers + mock_team_servers.return_value = team_servers + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + + # Assert the result (order-independent comparison) + assert sorted(result) == sorted(expected_servers) + + # Verify the mock functions were called correctly + mock_key_servers.assert_called_once_with(user_api_key_auth) + mock_team_servers.assert_called_once_with(user_api_key_auth) + + async def test_permission_inheritance_edge_cases(self): + """Test edge cases in permission inheritance""" + + # Test case: None values in database + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.find_unique.return_value = None + mock_prisma_client.db.litellm_teamtable.find_unique.return_value = None + + user_api_key_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + object_permission_id="test-permission" + ) + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + assert result == [] + + # Test case: Exception handling + mock_prisma_client.db.litellm_objectpermissiontable.find_unique.side_effect = Exception("DB Error") + + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth) + assert result == [] # Should handle exception gracefully + @pytest.mark.parametrize( "headers,expected_api_key,expected_mcp_auth_header,expected_server_auth_headers", [