mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-09 12:24:21 +00:00
[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 <ishaan@berri.ai> * 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 <ishaan@berri.ai> * 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 <ishaan@berri.ai> * Add test: MCP route without model still returns team_id Co-authored-by: ishaan <ishaan@berri.ai> * 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 <ishaan@berri.ai> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com> Co-authored-by: ishaan <ishaan@berri.ai>
This commit is contained in:
Ishaan Jaff
co-authored by
Cursor Agent
ishaan
Cursor Agent
ishaan
parent
12b8cd5971
commit
66eadfabe4
@@ -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
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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!"
|
||||
)
|
||||
Reference in New Issue
Block a user