diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py new file mode 100644 index 00000000000..5e2f223282b --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py @@ -0,0 +1,207 @@ +""" +Simple test to validate MCP permissions are enforced when calling MCP routes with JWT. +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_TeamTable, + LiteLLM_ObjectPermissionTable, + UserAPIKeyAuth, +) + + +@pytest.mark.asyncio +async def test_simple_jwt_mcp_permissions_enforced(): + """ + Simple test: Call MCP route with JWT, verify team's MCP servers are returned. + + Setup: + - Team "my-team" has MCP servers: ["github-mcp", "slack-mcp"] + - JWT user belongs to "my-team" + + Expected: Only ["github-mcp", "slack-mcp"] should be allowed + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + # 1. Create a user authenticated via JWT with team_id set + user_auth = UserAPIKeyAuth( + api_key=None, # JWT auth doesn't have api_key + user_id="jwt-user-123", + team_id="my-team", # This is set by JWT auth when team is in groups + ) + + # 2. Team's MCP permissions + team_mcp_servers = ["github-mcp", "slack-mcp"] + team_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-123", + mcp_servers=team_mcp_servers, + ) + + # 3. Mock the team permission lookup + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock + ) as mock_team_perm: + mock_team_perm.return_value = team_object_permission + + # Mock key permissions (empty - user has no key-level MCP permissions) + with patch.object( + MCPRequestHandler, "_get_key_object_permission", new_callable=AsyncMock + ) as mock_key_perm: + mock_key_perm.return_value = None + + # Mock access groups (empty) + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock + ) as mock_access_groups: + mock_access_groups.return_value = [] + + # 4. Call get_allowed_mcp_servers - this is what MCP routes use + allowed = await MCPRequestHandler.get_allowed_mcp_servers(user_auth) + + # 5. Verify only team's MCP servers are returned + assert sorted(allowed) == sorted(team_mcp_servers), ( + f"Expected {team_mcp_servers}, got {allowed}" + ) + + # Verify team permission was looked up + mock_team_perm.assert_called_once_with(user_auth) + + +@pytest.mark.asyncio +async def test_simple_jwt_no_team_no_mcp_servers(): + """ + Simple test: JWT user with no team should get no MCP servers. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + # User with no team_id (JWT didn't have teams in groups) + user_auth = UserAPIKeyAuth( + api_key=None, + user_id="jwt-user-no-team", + team_id=None, # No team + ) + + # _get_allowed_mcp_servers_for_team returns [] when team_id is None + allowed = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_auth) + + assert allowed == [], f"Expected [], got {allowed}" + + +@pytest.mark.asyncio +async def test_simple_jwt_team_id_required_for_mcp_permissions(): + """ + Simple test: Verify that team_id must be set for team MCP permissions to work. + + This is the key insight - if JWT auth doesn't set team_id, + team MCP permissions won't be enforced. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + + # Case 1: team_id is set -> team permissions should be checked + user_with_team = UserAPIKeyAuth( + api_key=None, + user_id="user-1", + team_id="team-abc", + ) + + team_mcp_servers = ["server-1", "server-2"] + team_perm = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-1", + mcp_servers=team_mcp_servers, + ) + + with patch.object( + MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock + ) as mock_perm: + mock_perm.return_value = team_perm + + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.return_value = [] + + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_with_team) + + assert sorted(result) == sorted(team_mcp_servers) + mock_perm.assert_called_once() # Permission WAS checked + + # Case 2: team_id is None -> team permissions NOT checked + user_without_team = UserAPIKeyAuth( + api_key=None, + user_id="user-2", + team_id=None, + ) + + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_without_team) + assert result == [] # No permissions returned + + +@pytest.mark.asyncio +async def test_jwt_auth_sets_team_id_for_mcp_route(): + """ + Test that JWT auth properly sets team_id when accessing MCP routes. + + This is the critical test - when user calls /mcp/tools/list with JWT, + the team_id from JWT groups must be set on UserAPIKeyAuth. + """ + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + + # Setup + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_ids_jwt_field="groups", # Teams come from "groups" field in JWT + ) + + # Team exists with models + team = LiteLLM_TeamTable( + team_id="team-from-jwt", + models=["gpt-4"], + ) + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + # Mock JWT token with team in groups + jwt_payload = { + "sub": "user-123", + "groups": ["team-from-jwt"], + "scope": "", + } + + with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth: + mock_auth.return_value = jwt_payload + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", new_callable=AsyncMock + ) as mock_get_team: + mock_get_team.return_value = team + + # Simulate calling MCP route + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/mcp/tools/list", # MCP route + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # THE KEY ASSERTION: team_id must be set + assert result["team_id"] == "team-from-jwt", ( + f"team_id should be 'team-from-jwt' but got '{result['team_id']}'. " + "This means JWT auth is not properly setting team_id for MCP routes!" + )