From 66eadfabe466f1f59b62c8afe2da9529ec248508 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 3 Feb 2026 19:13:13 -0800 Subject: [PATCH] [Bug] Ensure MCP permissions are enforced when using JWT Auth (#20383) * fix: enforce team MCP permissions when using JWT authentication Root cause: When JWT auth was used with teams in groups (via team_ids_jwt_field), the team's MCP permissions were not being enforced because: 1. The default team_allowed_routes did not include mcp_routes 2. allowed_routes_check() failed for MCP endpoints like /mcp/tools/list 3. find_team_with_model_access() skipped the team due to failed route check 4. team_id was None in UserAPIKeyAuth 5. MCPRequestHandler._get_allowed_mcp_servers_for_team() returned empty list Fix: Add 'mcp_routes' to the default team_allowed_routes in LiteLLM_JWTAuth. This ensures that teams can access MCP endpoints by default, allowing the team's MCP server permissions to be properly enforced. Added tests: - test_reproduce_jwt_mcp_enforcement_issue: Reproduces the exact bug scenario - test_verify_mcp_routes_in_default_team_allowed_routes: Verifies fix - test_mcp_route_check_passes_for_team: Verifies route check works Co-authored-by: ishaan * test: add comprehensive E2E tests for JWT + team MCP permission enforcement Added tests: - test_e2e_jwt_team_mcp_permissions_enforced: Full E2E test verifying JWT auth with teams in groups properly sets team_id and MCPRequestHandler returns the team's MCP servers - test_e2e_jwt_without_team_no_mcp_servers: Verifies no MCP servers returned when JWT has no teams - test_e2e_jwt_team_mcp_key_intersection: Verifies intersection logic when both key and team have MCP permissions (result = intersection) These tests verify the complete flow: 1. JWT token with team in groups field 2. JWT auth properly sets team_id on UserAPIKeyAuth 3. MCPRequestHandler.get_allowed_mcp_servers() returns team's MCP servers 4. Key/team permission intersection works correctly Co-authored-by: ishaan * test: add simple tests for JWT + MCP permission enforcement Simple, focused tests that validate: 1. test_simple_jwt_mcp_permissions_enforced: JWT user with team gets team's MCP servers 2. test_simple_jwt_no_team_no_mcp_servers: JWT user without team gets no MCP servers 3. test_simple_jwt_team_id_required_for_mcp_permissions: Verifies team_id is required 4. test_jwt_auth_sets_team_id_for_mcp_route: JWT auth sets team_id for MCP routes These tests directly verify the core MCP permission enforcement logic works when using JWT authentication with teams. Co-authored-by: ishaan * Add test: MCP route without model still returns team_id Co-authored-by: ishaan * Add 2 debug logs for JWT+MCP troubleshooting - handle_jwt.py: Log team route check result (team_id, route, is_allowed) - user_api_key_auth_mcp.py: Log team_id when looking up MCP permissions Co-authored-by: ishaan --------- Co-authored-by: Cursor Agent Co-authored-by: ishaan --- .../mcp_server/auth/user_api_key_auth_mcp.py | 3 + litellm/proxy/_types.py | 2 +- litellm/proxy/auth/handle_jwt.py | 3 + .../mcp_server/test_jwt_mcp_enforcement.py | 480 ++++++++++++++++++ .../mcp_server/test_jwt_mcp_simple.py | 277 ++++++++++ 5 files changed, 764 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py 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 49d6ac7d89..7e70b5baae 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 @@ -387,6 +387,9 @@ class MCPRequestHandler: user_api_key_cache, ) + verbose_logger.debug( + f"MCP team permission lookup: team_id={user_api_key_auth.team_id if user_api_key_auth else None}" + ) if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: return None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9ae95085f5..131a6caab0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3673,7 +3673,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): team_id_upsert: bool = False team_ids_jwt_field: Optional[str] = None upsert_sso_user_to_team: bool = False - team_allowed_routes: List[str] = ["openai_routes", "info_routes"] + team_allowed_routes: List[str] = ["openai_routes", "info_routes", "mcp_routes"] team_id_default: Optional[str] = Field( default=None, description="If no team_id given, default permissions/spend-tracking to this team.s", diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 33667b5d8d..584be0a949 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -976,6 +976,9 @@ class JWTAuthManager: user_route=route, litellm_proxy_roles=jwt_handler.litellm_jwtauth, ) + verbose_proxy_logger.debug( + f"JWT team route check: team_id={team_id}, route={route}, is_allowed={is_allowed}" + ) if is_allowed: return team_id, team_object except Exception: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py new file mode 100644 index 0000000000..b4a5a8ca19 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_enforcement.py @@ -0,0 +1,480 @@ +""" +Test to verify Team MCP permissions are enforced when using JWT authentication. + +Scenario: +1. Team "ABC" exists with models configured and MCPs assigned +2. User JWT has team "ABC" in groups (via team_ids_jwt_field) +3. Call MCP list endpoint +4. EXPECTED: Team MCP permissions should be enforced +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_TeamTable, + LiteLLM_ObjectPermissionTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler +from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, +) +from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager + + +@pytest.mark.asyncio +async def test_reproduce_jwt_mcp_enforcement_issue(monkeypatch): + """ + Reproduce the bug where Team MCP permissions are NOT enforced when using JWT. + + Setup: + - Team "ABC" has models ["gpt-4"] and MCPs ["mcp-server-1"] assigned + - JWT has team "ABC" in groups field + - User calls MCP list endpoint (no model requested) + + Expected: team_id should be set to "ABC" so MCP permissions are enforced + Actual (BUG): team_id is None because route check fails for MCP routes + """ + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + # Setup mock router + router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}]) + import sys + import types + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + # Team "ABC" has models configured AND MCPs assigned + team_with_mcp = LiteLLM_TeamTable( + team_id="ABC", + models=["gpt-4"], # Team HAS models + object_permission=LiteLLM_ObjectPermissionTable( + object_permission_id="perm-123", + mcp_servers=["mcp-server-1"], # Team has MCPs assigned + ), + ) + + async def mock_get_team_object(*args, **kwargs): + team_id = kwargs.get("team_id") or args[0] + if team_id == "ABC": + return team_with_mcp + return None + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + # Setup JWT handler with team_ids_jwt_field (groups) + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_ids_jwt_field="groups", # Use groups field for teams + # NOTE: team_allowed_routes defaults to ["openai_routes", "info_routes"] + # which does NOT include "mcp_routes" + ) + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + # Simulate JWT payload with team in groups + jwt_token = { + "sub": "user-123", + "groups": ["ABC"], # Team "ABC" is in groups + "scope": "", + } + + # Mock auth_jwt to return our token + with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt: + mock_auth_jwt.return_value = jwt_token + + # Call auth_builder for MCP route (like /mcp/tools/list) + result = await JWTAuthManager.auth_builder( + api_key="test-jwt-token", + jwt_handler=jwt_handler, + request_data={}, # No model in request (MCP endpoint) + 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, + ) + + # THIS IS THE BUG: team_id should be "ABC" but it's None! + print(f"Result team_id: {result['team_id']}") + print(f"Result team_object: {result['team_object']}") + + # The test should FAIL if the bug exists (team_id is None) + # If the fix is applied, team_id should be "ABC" + assert result["team_id"] == "ABC", ( + f"BUG: team_id should be 'ABC' but got '{result['team_id']}'. " + f"This happens because default team_allowed_routes does not include 'mcp_routes', " + f"so allowed_routes_check() fails and the team is skipped in find_team_with_model_access()." + ) + + +@pytest.mark.asyncio +async def test_verify_mcp_routes_in_default_team_allowed_routes(): + """ + Verify that mcp_routes IS in the default team_allowed_routes. + This is required for team MCP permissions to work with JWT auth. + """ + default_jwt_auth = LiteLLM_JWTAuth() + + print(f"Default team_allowed_routes: {default_jwt_auth.team_allowed_routes}") + + # mcp_routes must be in defaults for team MCP permissions to work + assert "mcp_routes" in default_jwt_auth.team_allowed_routes, ( + "mcp_routes must be in default team_allowed_routes for JWT MCP enforcement to work" + ) + + +@pytest.mark.asyncio +async def test_mcp_route_check_passes_for_team(): + """ + Verify that allowed_routes_check returns True for MCP routes with default settings. + This is required for teams to access MCP endpoints with JWT auth. + """ + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.auth.auth_checks import allowed_routes_check + + jwt_auth = LiteLLM_JWTAuth() # Use defaults + + # Check if MCP route is allowed for TEAM role + is_allowed = allowed_routes_check( + user_role=LitellmUserRoles.TEAM, + user_route="/mcp/tools/list", + litellm_proxy_roles=jwt_auth, + ) + + print(f"Is /mcp/tools/list allowed for TEAM with defaults? {is_allowed}") + + # MCP routes should be allowed by default for teams + assert is_allowed is True, ( + "MCP routes must be allowed by default for teams for JWT MCP enforcement to work" + ) + + +@pytest.mark.asyncio +async def test_e2e_jwt_team_mcp_permissions_enforced(monkeypatch): + """ + End-to-end test verifying that team MCP permissions are properly enforced + when using JWT authentication with teams in groups. + + This test verifies the complete flow: + 1. JWT token contains team "ABC" in groups field + 2. Team "ABC" exists with MCP servers ["mcp-server-1", "mcp-server-2"] assigned + 3. JWT auth properly sets team_id on UserAPIKeyAuth + 4. MCPRequestHandler.get_allowed_mcp_servers() returns team's MCP servers + """ + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + # Setup mock router + router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}]) + import sys + import types + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + proxy_server_module.prisma_client = MagicMock() # Mock prisma client + proxy_server_module.user_api_key_cache = DualCache() + proxy_server_module.proxy_logging_obj = MagicMock() + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + # Team "ABC" has MCP servers assigned via object_permission + team_mcp_servers = ["mcp-server-1", "mcp-server-2"] + team_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-abc-123", + mcp_servers=team_mcp_servers, + mcp_access_groups=[], + vector_stores=[], + ) + + team_with_mcp = LiteLLM_TeamTable( + team_id="ABC", + models=["gpt-4"], + object_permission=team_object_permission, + object_permission_id="perm-abc-123", + ) + + async def mock_get_team_object(*args, **kwargs): + team_id = kwargs.get("team_id") or (args[0] if args else None) + if team_id == "ABC": + return team_with_mcp + return None + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.get_team_object", mock_get_team_object + ) + + # Setup JWT handler with team_ids_jwt_field (groups) + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + ) + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + # Simulate JWT payload with team in groups + jwt_token = { + "sub": "user-123", + "groups": ["ABC"], + "scope": "", + } + + # Step 1: Verify JWT auth returns correct team_id + with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt: + mock_auth_jwt.return_value = jwt_token + + result = await JWTAuthManager.auth_builder( + api_key="test-jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/mcp/tools/list", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # Verify team_id is set correctly + assert result["team_id"] == "ABC", f"Expected team_id='ABC', got '{result['team_id']}'" + assert result["team_object"] is not None, "team_object should not be None" + + # Step 2: Create UserAPIKeyAuth with the team_id from JWT auth + user_api_key_auth = UserAPIKeyAuth( + api_key=None, + team_id=result["team_id"], + user_id=result["user_id"], + ) + + # Step 3: Verify MCPRequestHandler returns team's MCP servers + # Mock _get_team_object_permission to return our team's object_permission + with patch.object( + MCPRequestHandler, "_get_team_object_permission" + ) as mock_get_team_perm: + mock_get_team_perm.return_value = team_object_permission + + # Mock _get_allowed_mcp_servers_for_key to return empty (no key-level permissions) + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_key" + ) as mock_key_servers: + mock_key_servers.return_value = [] + + # Mock _get_mcp_servers_from_access_groups to return empty + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups" + ) as mock_access_groups: + mock_access_groups.return_value = [] + + allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth + ) + + print(f"Allowed MCP servers: {allowed_servers}") + + # Verify team's MCP servers are returned + assert set(allowed_servers) == set(team_mcp_servers), ( + f"Expected team MCP servers {team_mcp_servers}, got {allowed_servers}" + ) + + +@pytest.mark.asyncio +async def test_e2e_jwt_without_team_no_mcp_servers(monkeypatch): + """ + End-to-end test verifying that when JWT has no teams, no MCP servers are returned. + + This ensures: + 1. JWT token with no groups returns no team_id + 2. MCPRequestHandler.get_allowed_mcp_servers() returns empty list + """ + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + # Setup mock router + router = Router(model_list=[]) + import sys + import types + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + async def mock_get_team_object(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + # Setup JWT handler + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + ) + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + # JWT payload with empty groups + jwt_token = { + "sub": "user-123", + "groups": [], # No teams + "scope": "", + } + + with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt: + mock_auth_jwt.return_value = jwt_token + + result = await JWTAuthManager.auth_builder( + api_key="test-jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/mcp/tools/list", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # Verify no team_id is set + assert result["team_id"] is None, f"Expected team_id=None, got '{result['team_id']}'" + + # Create UserAPIKeyAuth without team_id + user_api_key_auth = UserAPIKeyAuth( + api_key=None, + team_id=None, + user_id=result["user_id"], + ) + + # Verify no MCP servers are returned when there's no team + allowed_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_team( + user_api_key_auth + ) + + assert allowed_servers == [], f"Expected empty list, got {allowed_servers}" + + +@pytest.mark.asyncio +async def test_e2e_jwt_team_mcp_key_intersection(monkeypatch): + """ + End-to-end test verifying MCP permission intersection between key and team. + + Scenario: + - Team has MCP servers: ["server-1", "server-2", "server-3"] + - Key has MCP servers: ["server-2", "server-4"] + - Result should be intersection: ["server-2"] + """ + from litellm.caching import DualCache + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + # Setup mock router + router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}]) + import sys + import types + proxy_server_module = types.ModuleType("proxy_server") + proxy_server_module.llm_router = router + proxy_server_module.prisma_client = MagicMock() + proxy_server_module.user_api_key_cache = DualCache() + proxy_server_module.proxy_logging_obj = MagicMock() + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_module) + + # Team MCP servers + team_mcp_servers = ["server-1", "server-2", "server-3"] + team_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="team-perm", + mcp_servers=team_mcp_servers, + ) + + team_with_mcp = LiteLLM_TeamTable( + team_id="TEAM-X", + models=["gpt-4"], + object_permission=team_object_permission, + ) + + # Key MCP servers + key_mcp_servers = ["server-2", "server-4"] + key_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="key-perm", + mcp_servers=key_mcp_servers, + ) + + async def mock_get_team_object(*args, **kwargs): + team_id = kwargs.get("team_id") or (args[0] if args else None) + if team_id == "TEAM-X": + return team_with_mcp + return None + + monkeypatch.setattr( + "litellm.proxy.auth.handle_jwt.get_team_object", mock_get_team_object + ) + + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_ids_jwt_field="groups") + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_token = {"sub": "user-123", "groups": ["TEAM-X"], "scope": ""} + + with patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt: + mock_auth_jwt.return_value = jwt_token + + result = await JWTAuthManager.auth_builder( + api_key="test-jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/mcp/tools/list", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + assert result["team_id"] == "TEAM-X" + + user_api_key_auth = UserAPIKeyAuth( + api_key=None, + team_id=result["team_id"], + user_id=result["user_id"], + object_permission=key_object_permission, # Key has its own permissions + ) + + # Mock the helper methods to return our test data + with patch.object( + MCPRequestHandler, "_get_team_object_permission" + ) as mock_team_perm: + mock_team_perm.return_value = team_object_permission + + with patch.object( + MCPRequestHandler, "_get_key_object_permission" + ) as mock_key_perm: + mock_key_perm.return_value = key_object_permission + + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups" + ) as mock_access_groups: + mock_access_groups.return_value = [] + + allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth + ) + + # Should be intersection: only server-2 is in both + expected = ["server-2"] + assert sorted(allowed_servers) == sorted(expected), ( + f"Expected intersection {expected}, got {allowed_servers}" + ) 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 0000000000..9ad7736d01 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_jwt_mcp_simple.py @@ -0,0 +1,277 @@ +""" +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!" + ) + + +@pytest.mark.asyncio +async def test_mcp_route_without_model_still_returns_team_id(): + """ + Test that MCP routes (which don't specify a model) still get team_id assigned. + + Key insight: MCP routes don't require a model in the request, but the JWT auth + flow must still assign a team_id so that team MCP permissions are enforced. + + The flow is: + 1. JWT token contains team in "groups" field + 2. find_team_with_model_access() is called with requested_model=None + 3. Since `not requested_model` is True, model check passes + 4. Route check passes because "mcp_routes" is in team_allowed_routes + 5. team_id is returned and 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", + ) + + # Team exists - note: models is a list (can be empty or have values) + # The key is that when no model is requested, model check is skipped + team = LiteLLM_TeamTable( + team_id="my-team", + models=["gpt-4", "gpt-3.5-turbo"], # Team has models, but MCP request won't specify one + ) + + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + # JWT with team in groups + jwt_payload = { + "sub": "user-abc", + "groups": ["my-team"], + "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 + + # Call MCP route with NO MODEL in request_data + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={}, # <-- NO MODEL SPECIFIED + 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, + ) + + # Team ID must still be set even though no model was requested + assert result["team_id"] == "my-team", ( + f"Expected team_id='my-team' but got '{result['team_id']}'. " + "MCP routes without model should still get team_id from JWT!" + )