From cf6b0e45a7b5bd587b4414368c827f0959d6dc8f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 17 Feb 2026 15:47:37 -0800 Subject: [PATCH] feat: address greptile comments --- .../mcp_server/auth/user_api_key_auth_mcp.py | 50 +- .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{budgets.html => budgets/index.html} | 0 .../{caching.html => caching/index.html} | 0 .../index.html} | 0 .../{old-usage.html => old-usage/index.html} | 0 .../{prompts.html => prompts/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{login.html => login/index.html} | 0 .../out/{logs.html => logs/index.html} | 0 .../{callback.html => callback/index.html} | 0 .../{model-hub.html => model-hub/index.html} | 0 .../{model_hub.html => model_hub/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{policies.html => policies/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../{ui-theme.html => ui-theme/index.html} | 0 .../out/{teams.html => teams/index.html} | 0 .../{test-key.html => test-key/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{usage.html => usage/index.html} | 0 .../out/{users.html => users/index.html} | 0 .../index.html} | 0 .../customer_endpoints.py | 15 +- .../auth/test_user_api_key_auth_mcp.py | 670 +++++++++++++++++- 36 files changed, 698 insertions(+), 37 deletions(-) rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{_not-found.html => _not-found/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) 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 67cb783bbba..0d1a11e760c 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 @@ -334,11 +334,16 @@ class MCPRequestHandler: """ Get list of allowed MCP servers for the given user/key based on permissions. - Permission hierarchy (all rules are intersections): - 1. Get allowed servers from key permissions - 2. Get allowed servers from team permissions - 3. Get allowed servers from end_user permissions - 4. Final result = intersection of key/team AND end_user (if end_user has permissions set) + Permission hierarchy: + 1. Calculate key/team allowed servers: + - If team has permissions: key inherits from team (or intersects if key has own permissions) + - If team has no permissions: use key permissions + 2. Apply end_user permissions if end_user_id is set: + - If require_end_user_mcp_access_defined=True and end_user has no permissions: block all access + - If end_user has permissions: + * If key/team has no restrictions (empty): use end_user permissions + * If key/team has restrictions: intersect key/team AND end_user + - If end_user has no permissions and flag is False: use key/team permissions Returns: List[str]: List of allowed MCP servers by server id @@ -386,40 +391,41 @@ class MCPRequestHandler: ) ) + # Check if require_end_user_mcp_access flag is enabled + require_end_user_mcp_access = general_settings.get( + "require_end_user_mcp_access_defined", False + ) - # If end_user has explicit MCP server permissions, apply intersection + # If the flag is enabled and end_user has no permissions, block all access + if require_end_user_mcp_access and len(allowed_mcp_servers_for_end_user) == 0: + verbose_logger.debug( + f"require_end_user_mcp_access_defined=True and end_user {user_api_key_auth.end_user_id} has no MCP permissions - blocking MCP access" + ) + return [] + + # If end_user has explicit MCP server permissions, apply intersection logic if len(allowed_mcp_servers_for_end_user) > 0: verbose_logger.debug( f"End user {user_api_key_auth.end_user_id} has explicit MCP permissions: {allowed_mcp_servers_for_end_user}" ) - # Check if require_end_user_mcp_access flag is enabled - require_end_user_mcp_access = general_settings.get( - "require_end_user_mcp_access_defined", False - ) - - # If the flag is enabled and end_user has permissions, use ONLY end_user permissions - if require_end_user_mcp_access: + # If key/team has no restrictions (empty list), use end_user permissions + if len(allowed_mcp_servers) == 0: verbose_logger.debug( - "require_end_user_mcp_access_defined=True - using only end_user permissions" + "No key/team MCP restrictions - using end_user permissions" ) allowed_mcp_servers = allowed_mcp_servers_for_end_user else: - # Otherwise, apply intersection: key/team AND end_user + # Apply intersection: key/team AND end_user + # This ensures key/team restrictions are ALWAYS respected filtered_servers = [] for _mcp_server in allowed_mcp_servers: if _mcp_server in allowed_mcp_servers_for_end_user: filtered_servers.append(_mcp_server) allowed_mcp_servers = filtered_servers verbose_logger.debug( - f"Applied end_user intersection filter. Final allowed servers: {allowed_mcp_servers}" + f"Applied end_user intersection with key/team. Final allowed servers: {allowed_mcp_servers}" ) - # If flag is enabled but end_user has no permissions, block all access - elif general_settings.get("require_end_user_mcp_access_defined", False): - verbose_logger.debug( - f"require_end_user_mcp_access_defined=True and end_user {user_api_key_auth.end_user_id} has no MCP permissions - blocking MCP access" - ) - return [] return list(set(allowed_mcp_servers)) except Exception as e: diff --git a/litellm/proxy/_experimental/out/404.html b/litellm/proxy/_experimental/out/404/index.html similarity index 100% rename from litellm/proxy/_experimental/out/404.html rename to litellm/proxy/_experimental/out/404/index.html diff --git a/litellm/proxy/_experimental/out/_not-found.html b/litellm/proxy/_experimental/out/_not-found/index.html similarity index 100% rename from litellm/proxy/_experimental/out/_not-found.html rename to litellm/proxy/_experimental/out/_not-found/index.html diff --git a/litellm/proxy/_experimental/out/api-reference.html b/litellm/proxy/_experimental/out/api-reference/index.html similarity index 100% rename from litellm/proxy/_experimental/out/api-reference.html rename to litellm/proxy/_experimental/out/api-reference/index.html diff --git a/litellm/proxy/_experimental/out/experimental/api-playground.html b/litellm/proxy/_experimental/out/experimental/api-playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/api-playground.html rename to litellm/proxy/_experimental/out/experimental/api-playground/index.html diff --git a/litellm/proxy/_experimental/out/experimental/budgets.html b/litellm/proxy/_experimental/out/experimental/budgets/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/budgets.html rename to litellm/proxy/_experimental/out/experimental/budgets/index.html diff --git a/litellm/proxy/_experimental/out/experimental/caching.html b/litellm/proxy/_experimental/out/experimental/caching/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/caching.html rename to litellm/proxy/_experimental/out/experimental/caching/index.html diff --git a/litellm/proxy/_experimental/out/experimental/claude-code-plugins.html b/litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/claude-code-plugins.html rename to litellm/proxy/_experimental/out/experimental/claude-code-plugins/index.html diff --git a/litellm/proxy/_experimental/out/experimental/old-usage.html b/litellm/proxy/_experimental/out/experimental/old-usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/old-usage.html rename to litellm/proxy/_experimental/out/experimental/old-usage/index.html diff --git a/litellm/proxy/_experimental/out/experimental/prompts.html b/litellm/proxy/_experimental/out/experimental/prompts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/prompts.html rename to litellm/proxy/_experimental/out/experimental/prompts/index.html diff --git a/litellm/proxy/_experimental/out/experimental/tag-management.html b/litellm/proxy/_experimental/out/experimental/tag-management/index.html similarity index 100% rename from litellm/proxy/_experimental/out/experimental/tag-management.html rename to litellm/proxy/_experimental/out/experimental/tag-management/index.html diff --git a/litellm/proxy/_experimental/out/guardrails.html b/litellm/proxy/_experimental/out/guardrails/index.html similarity index 100% rename from litellm/proxy/_experimental/out/guardrails.html rename to litellm/proxy/_experimental/out/guardrails/index.html diff --git a/litellm/proxy/_experimental/out/login.html b/litellm/proxy/_experimental/out/login/index.html similarity index 100% rename from litellm/proxy/_experimental/out/login.html rename to litellm/proxy/_experimental/out/login/index.html diff --git a/litellm/proxy/_experimental/out/logs.html b/litellm/proxy/_experimental/out/logs/index.html similarity index 100% rename from litellm/proxy/_experimental/out/logs.html rename to litellm/proxy/_experimental/out/logs/index.html diff --git a/litellm/proxy/_experimental/out/mcp/oauth/callback.html b/litellm/proxy/_experimental/out/mcp/oauth/callback/index.html similarity index 100% rename from litellm/proxy/_experimental/out/mcp/oauth/callback.html rename to litellm/proxy/_experimental/out/mcp/oauth/callback/index.html diff --git a/litellm/proxy/_experimental/out/model-hub.html b/litellm/proxy/_experimental/out/model-hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model-hub.html rename to litellm/proxy/_experimental/out/model-hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub.html b/litellm/proxy/_experimental/out/model_hub/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub.html rename to litellm/proxy/_experimental/out/model_hub/index.html diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/models-and-endpoints.html b/litellm/proxy/_experimental/out/models-and-endpoints/index.html similarity index 100% rename from litellm/proxy/_experimental/out/models-and-endpoints.html rename to litellm/proxy/_experimental/out/models-and-endpoints/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding/index.html similarity index 100% rename from litellm/proxy/_experimental/out/onboarding.html rename to litellm/proxy/_experimental/out/onboarding/index.html diff --git a/litellm/proxy/_experimental/out/organizations.html b/litellm/proxy/_experimental/out/organizations/index.html similarity index 100% rename from litellm/proxy/_experimental/out/organizations.html rename to litellm/proxy/_experimental/out/organizations/index.html diff --git a/litellm/proxy/_experimental/out/playground.html b/litellm/proxy/_experimental/out/playground/index.html similarity index 100% rename from litellm/proxy/_experimental/out/playground.html rename to litellm/proxy/_experimental/out/playground/index.html diff --git a/litellm/proxy/_experimental/out/policies.html b/litellm/proxy/_experimental/out/policies/index.html similarity index 100% rename from litellm/proxy/_experimental/out/policies.html rename to litellm/proxy/_experimental/out/policies/index.html diff --git a/litellm/proxy/_experimental/out/settings/admin-settings.html b/litellm/proxy/_experimental/out/settings/admin-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/admin-settings.html rename to litellm/proxy/_experimental/out/settings/admin-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/logging-and-alerts.html b/litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/logging-and-alerts.html rename to litellm/proxy/_experimental/out/settings/logging-and-alerts/index.html diff --git a/litellm/proxy/_experimental/out/settings/router-settings.html b/litellm/proxy/_experimental/out/settings/router-settings/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/router-settings.html rename to litellm/proxy/_experimental/out/settings/router-settings/index.html diff --git a/litellm/proxy/_experimental/out/settings/ui-theme.html b/litellm/proxy/_experimental/out/settings/ui-theme/index.html similarity index 100% rename from litellm/proxy/_experimental/out/settings/ui-theme.html rename to litellm/proxy/_experimental/out/settings/ui-theme/index.html diff --git a/litellm/proxy/_experimental/out/teams.html b/litellm/proxy/_experimental/out/teams/index.html similarity index 100% rename from litellm/proxy/_experimental/out/teams.html rename to litellm/proxy/_experimental/out/teams/index.html diff --git a/litellm/proxy/_experimental/out/test-key.html b/litellm/proxy/_experimental/out/test-key/index.html similarity index 100% rename from litellm/proxy/_experimental/out/test-key.html rename to litellm/proxy/_experimental/out/test-key/index.html diff --git a/litellm/proxy/_experimental/out/tools/mcp-servers.html b/litellm/proxy/_experimental/out/tools/mcp-servers/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/mcp-servers.html rename to litellm/proxy/_experimental/out/tools/mcp-servers/index.html diff --git a/litellm/proxy/_experimental/out/tools/vector-stores.html b/litellm/proxy/_experimental/out/tools/vector-stores/index.html similarity index 100% rename from litellm/proxy/_experimental/out/tools/vector-stores.html rename to litellm/proxy/_experimental/out/tools/vector-stores/index.html diff --git a/litellm/proxy/_experimental/out/usage.html b/litellm/proxy/_experimental/out/usage/index.html similarity index 100% rename from litellm/proxy/_experimental/out/usage.html rename to litellm/proxy/_experimental/out/usage/index.html diff --git a/litellm/proxy/_experimental/out/users.html b/litellm/proxy/_experimental/out/users/index.html similarity index 100% rename from litellm/proxy/_experimental/out/users.html rename to litellm/proxy/_experimental/out/users/index.html diff --git a/litellm/proxy/_experimental/out/virtual-keys.html b/litellm/proxy/_experimental/out/virtual-keys/index.html similarity index 100% rename from litellm/proxy/_experimental/out/virtual-keys.html rename to litellm/proxy/_experimental/out/virtual-keys/index.html diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index bfafc943c45..43eb1ea51ec 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -174,7 +174,8 @@ async def _handle_customer_object_permission_update( """ Handle object permission updates for customer endpoints. - Updates the update_end_user_table_data dict in place with the new object_permission_id. + Updates the update_end_user_table_data dict in place with the new object_permission_id + and removes the object_permission field to prevent it from being sent to the database. Args: non_default_values: Dictionary containing the update values including object_permission @@ -196,6 +197,10 @@ async def _handle_customer_object_permission_update( if object_permission_id is not None: update_end_user_table_data["object_permission_id"] = object_permission_id + # Remove object_permission from update_end_user_table_data to prevent DB write attempts + # object_permission is a read-only relationship field, not a writable column + update_end_user_table_data.pop("object_permission", None) + @router.post( "/end_user/new", @@ -612,14 +617,6 @@ async def update_end_user( ## Update user table, with update params + new budget id (if set) ## verbose_proxy_logger.debug("/customer/update: Received data = %s", data) - # Ensure object_permission is not in the update data - # It should have been converted to object_permission_id by handle_update_object_permission_common - if "object_permission" in update_end_user_table_data: - verbose_proxy_logger.warning( - f"object_permission still in update_end_user_table_data: {update_end_user_table_data.get('object_permission')}" - ) - update_end_user_table_data.pop("object_permission", None) - if data.user_id is not None and len(data.user_id) > 0: update_end_user_table_data["user_id"] = data.user_id # type: ignore verbose_proxy_logger.debug("In update customer, user_id condition block.") 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 c2dbc94f721..da5b06f92f5 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 @@ -15,9 +15,8 @@ sys.path.insert( from starlette.datastructures import Headers -from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCPRequestHandler, -) +from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import \ + MCPRequestHandler from litellm.proxy._types import SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -1176,7 +1175,8 @@ class TestMCPAccessGroupsE2E: @pytest.mark.asyncio def test_mcp_path_based_server_segregation(monkeypatch): # Import the MCP server FastAPI app and context getter - from litellm.proxy._experimental.mcp_server.server import app, get_auth_context + from litellm.proxy._experimental.mcp_server.server import ( + app, get_auth_context) captured_mcp_servers = {} @@ -1277,7 +1277,8 @@ async def test_get_team_object_permission_with_already_loaded_permission(): Test that _get_team_object_permission returns the already loaded object_permission from the team object without making an additional DB call. """ - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.proxy._types import (LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable) # Create mock object permission mock_object_permission = LiteLLM_ObjectPermissionTable( @@ -1340,7 +1341,8 @@ async def test_get_team_object_permission_with_core_auth_auto_loading(): the team object returned by get_team_object() should already have object_permission loaded when an object_permission_id exists. """ - from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable + from litellm.proxy._types import (LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable) # Create mock object permission mock_object_permission = LiteLLM_ObjectPermissionTable( @@ -1595,3 +1597,659 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission(): assert set(result) == {"direct-server", "group-server"} mock_get_perm.assert_not_called() mock_access_groups.assert_called_once_with(["grp-alpha"]) + + +@pytest.mark.asyncio +class TestEndUserMCPPermissions: + """Test suite for end_user MCP permission functionality""" + + @pytest.mark.parametrize( + "key_servers,team_servers,end_user_servers,require_flag,expected_result,scenario", + [ + # Test case 1: End user with permissions, key/team with permissions - intersection + ( + ["server1", "server2", "server3"], + ["server1", "server2", "server3", "server4"], + ["server2", "server3", "server5"], + False, + ["server2", "server3"], + "end_user_intersects_with_key_team", + ), + # Test case 2: End user with permissions, key/team with no permissions (empty) - use end_user permissions + ( + [], + [], + ["server1", "server2"], + False, + ["server1", "server2"], + "end_user_with_empty_key_team", + ), + # Test case 3: End user with no permissions, require_flag=True - fall back to key/team restrictions + ( + ["server1", "server2"], + ["server1", "server2"], + [], + True, + ["server1", "server2"], + "end_user_no_perms_with_require_flag_true_fallback_to_key_team", + ), + # Test case 4: End user with no permissions, require_flag=False - use key/team permissions + ( + ["server1", "server2"], + ["server1", "server2"], + [], + False, + ["server1", "server2"], + "end_user_no_perms_with_require_flag_false", + ), + # Test case 5: End user with permissions, require_flag=True - still intersect with key/team + ( + ["server1", "server2"], + ["server1", "server2", "server3"], + ["server2", "server3", "server4"], + True, + ["server2"], + "end_user_with_require_flag_true_still_intersects", + ), + # Test case 6: End user has permissions but no overlap with key/team + ( + ["server1", "server2"], + ["server1", "server2"], + ["server3", "server4"], + False, + [], + "end_user_no_overlap_with_key_team", + ), + # Test case 7: Key has servers, team empty, end_user has overlap with key + ( + ["server1", "server2", "server3"], + [], + ["server2", "server3", "server4"], + False, + ["server2", "server3"], + "key_only_intersects_with_end_user", + ), + # Test case 8: All have different servers + ( + ["server1"], + ["server1", "server2"], + ["server3"], + False, + [], + "all_different_servers", + ), + # Test case 9: require_flag=True, end_user has permissions but key/team empty + ( + [], + [], + ["server1", "server2"], + True, + ["server1", "server2"], + "require_flag_true_with_empty_key_team_uses_end_user", + ), + ], + ) + async def test_get_allowed_mcp_servers_with_end_user( + self, + key_servers, + team_servers, + end_user_servers, + require_flag, + expected_result, + scenario, + ): + """Test get_allowed_mcp_servers with end_user permissions""" + + # Create a mock user with end_user_id + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + ) + + # Mock the helper methods + 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: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user" + ) as mock_end_user_servers: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + # Set up return values + mock_key_servers.return_value = key_servers + mock_team_servers.return_value = team_servers + mock_end_user_servers.return_value = end_user_servers + mock_general_settings.get.return_value = require_flag + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth + ) + + # Assert the result (order-independent comparison) + assert sorted(result) == sorted(expected_result) + + # Verify helper methods were called + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) + mock_end_user_servers.assert_called_once_with(mock_user_auth) + + @pytest.mark.parametrize( + "key_servers,team_servers,expected_result,scenario", + [ + # Test case 1: No end_user_id - should use key/team logic only + ( + ["server1", "server2"], + ["server1", "server2", "server3"], + ["server1", "server2"], + "no_end_user_id_uses_key_team", + ), + # Test case 2: No end_user_id with empty key + ( + [], + ["server1", "server2"], + ["server1", "server2"], + "no_end_user_id_inherits_from_team", + ), + ], + ) + async def test_get_allowed_mcp_servers_without_end_user_id( + self, + key_servers, + team_servers, + expected_result, + scenario, + ): + """Test get_allowed_mcp_servers when end_user_id is not set""" + + # Create a mock user WITHOUT end_user_id + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id=None, + ) + + # Mock the helper methods + 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: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user" + ) as mock_end_user_servers: + # Set up return values + 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=mock_user_auth + ) + + # Assert the result (order-independent comparison) + assert sorted(result) == sorted(expected_result) + + # Verify helper methods were called correctly + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) + # end_user method should NOT be called since no end_user_id + mock_end_user_servers.assert_not_called() + + async def test_get_allowed_mcp_servers_for_end_user_with_valid_user(self): + """Test _get_allowed_mcp_servers_for_end_user with valid end_user""" + from litellm.proxy._types import (LiteLLM_EndUserTable, + LiteLLM_ObjectPermissionTable) + + # Create mock object permission + mock_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-123", + mcp_servers=["server1", "server2"], + mcp_access_groups=["group1"], + ) + + # Create mock end_user object + mock_end_user = LiteLLM_EndUserTable( + user_id="test-end-user", + object_permission=mock_object_permission, + object_permission_id="perm-123", + ) + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + end_user_id="test-end-user", + ) + + # Mock the dependencies + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user: + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups" + ) as mock_get_access_groups: + # Configure mocks + mock_get_end_user.return_value = mock_end_user + mock_get_access_groups.return_value = ["group-server1", "group-server2"] + + # Call the method + result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( + mock_user_auth + ) + + # Assert the result contains both direct and access group servers + assert set(result) == {"server1", "server2", "group-server1", "group-server2"} + + # Verify methods were called + mock_get_end_user.assert_called_once() + mock_get_access_groups.assert_called_once_with(["group1"]) + + async def test_get_allowed_mcp_servers_for_end_user_with_no_permission(self): + """Test _get_allowed_mcp_servers_for_end_user when end_user has no object_permission""" + from litellm.proxy._types import LiteLLM_EndUserTable + + # Create mock end_user object without object_permission + mock_end_user = LiteLLM_EndUserTable( + user_id="test-end-user", + object_permission=None, + ) + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + end_user_id="test-end-user", + ) + + # Mock the dependencies + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user: + # Configure mock + mock_get_end_user.return_value = mock_end_user + + # Call the method + result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( + mock_user_auth + ) + + # Assert empty list is returned + assert result == [] + + # Verify method was called + mock_get_end_user.assert_called_once() + + async def test_get_allowed_mcp_servers_for_end_user_without_end_user_id(self): + """Test _get_allowed_mcp_servers_for_end_user when no end_user_id is set""" + + # Create mock user auth without end_user_id + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + end_user_id=None, + ) + + # Call the method + result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( + mock_user_auth + ) + + # Assert empty list is returned + assert result == [] + + async def test_get_allowed_mcp_servers_for_end_user_without_prisma_client(self): + """Test _get_allowed_mcp_servers_for_end_user when prisma_client is None""" + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + end_user_id="test-end-user", + ) + + # Mock prisma_client as None + with patch("litellm.proxy.proxy_server.prisma_client", None): + # Call the method + result = await MCPRequestHandler._get_allowed_mcp_servers_for_end_user( + mock_user_auth + ) + + # Assert empty list is returned + assert result == [] + + +@pytest.mark.asyncio +class TestEndUserMCPPermissionsFallbackBehavior: + """Test suite for end_user MCP permission fallback behavior. + + Tests the updated behavior where when end_user has NO explicit MCP permissions defined, + the system falls back to key/team restrictions instead of blocking all access. + """ + + @pytest.mark.parametrize( + "key_servers,team_servers,end_user_servers,require_flag,expected_result,scenario", + [ + # Test case 1: End user with NO permissions, require_flag=True - should fall back to key/team + ( + ["server1", "server2"], + ["server1", "server2", "server3"], + [], + True, + ["server1", "server2"], + "no_end_user_perms_require_flag_true_fallback_to_key_team_intersection", + ), + # Test case 2: End user with NO permissions, require_flag=False - should fall back to key/team + ( + ["server1", "server2"], + ["server1", "server2", "server3"], + [], + False, + ["server1", "server2"], + "no_end_user_perms_require_flag_false_fallback_to_key_team", + ), + # Test case 3: End user with NO permissions, only key has permissions - should use key permissions + ( + ["server1", "server2", "server3"], + [], + [], + True, + ["server1", "server2", "server3"], + "no_end_user_perms_only_key_has_permissions", + ), + # Test case 4: End user with NO permissions, only team has permissions - should inherit from team + ( + [], + ["team_server1", "team_server2"], + [], + True, + ["team_server1", "team_server2"], + "no_end_user_perms_only_team_has_permissions", + ), + # Test case 5: End user with NO permissions, both key and team empty - should return empty + ( + [], + [], + [], + True, + [], + "no_permissions_anywhere_returns_empty", + ), + # Test case 6: End user with NO permissions, key and team have different servers + ( + ["server1", "server2"], + ["server3", "server4"], + [], + True, + [], + "no_end_user_perms_key_team_no_overlap_returns_empty", + ), + # Test case 7: End user HAS permissions, require_flag=True - should still intersect + ( + ["server1", "server2", "server3"], + ["server1", "server2", "server3"], + ["server2", "server3", "server4"], + True, + ["server2", "server3"], + "end_user_has_perms_require_flag_true_still_intersects", + ), + # Test case 8: End user HAS permissions but empty key/team - should use end_user permissions + ( + [], + [], + ["server1", "server2"], + True, + ["server1", "server2"], + "end_user_has_perms_empty_key_team_uses_end_user", + ), + # Test case 9: End user NO permissions, complex key/team intersection + ( + ["server1", "server2", "server3", "server4"], + ["server2", "server3", "server5"], + [], + True, + ["server2", "server3"], + "no_end_user_perms_complex_key_team_intersection", + ), + # Test case 10: End user HAS permissions with no overlap - should return empty + ( + ["server1", "server2"], + ["server1", "server2"], + ["server5", "server6"], + True, + [], + "end_user_has_perms_no_overlap_returns_empty", + ), + ], + ) + async def test_end_user_permission_fallback_to_key_team( + self, + key_servers, + team_servers, + end_user_servers, + require_flag, + expected_result, + scenario, + ): + """Test that end_user permissions fall back to key/team when not defined""" + + # Create a mock user with end_user_id + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + ) + + # Mock the helper methods + 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: + with patch.object( + MCPRequestHandler, "_get_allowed_mcp_servers_for_end_user" + ) as mock_end_user_servers: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + # Set up return values + mock_key_servers.return_value = key_servers + mock_team_servers.return_value = team_servers + mock_end_user_servers.return_value = end_user_servers + mock_general_settings.get.return_value = require_flag + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth + ) + + # Assert the result (order-independent comparison) + assert sorted(result) == sorted(expected_result), ( + f"Test scenario '{scenario}' failed: " + f"Expected {sorted(expected_result)}, got {sorted(result)}" + ) + + # Verify helper methods were called + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) + mock_end_user_servers.assert_called_once_with(mock_user_auth) + + async def test_end_user_no_permissions_with_require_flag_true_detailed(self): + """ + Detailed test: When require_end_user_mcp_access_defined=True and end_user has NO permissions, + should fall back to key/team restrictions (not block all access). + + This is the main behavior change being tested. + """ + from litellm.proxy._types import LiteLLM_EndUserTable + + # Create mock end_user object without object_permission + mock_end_user = LiteLLM_EndUserTable( + user_id="test-end-user", + object_permission=None, + ) + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + ) + + # Mock the dependencies + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user: + 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: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + # Configure mocks + mock_get_end_user.return_value = mock_end_user + mock_key_servers.return_value = ["server1", "server2"] + mock_team_servers.return_value = ["server1", "server2", "server3"] + mock_general_settings.get.return_value = True # require_flag=True + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth + ) + + # CRITICAL: Should fall back to key/team intersection, not block all + assert sorted(result) == ["server1", "server2"], ( + "When end_user has NO permissions and require_flag=True, " + "should fall back to key/team restrictions, not block all access" + ) + + # Verify methods were called + mock_get_end_user.assert_called_once() + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) + + async def test_end_user_no_permissions_without_require_flag(self): + """ + Test: When require_end_user_mcp_access_defined=False and end_user has NO permissions, + should use key/team restrictions (default behavior). + """ + from litellm.proxy._types import LiteLLM_EndUserTable + + # Create mock end_user object without object_permission + mock_end_user = LiteLLM_EndUserTable( + user_id="test-end-user", + object_permission=None, + ) + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + ) + + # Mock the dependencies + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user: + 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: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + # Configure mocks + mock_get_end_user.return_value = mock_end_user + mock_key_servers.return_value = ["server1", "server2", "server3"] + mock_team_servers.return_value = [] + mock_general_settings.get.return_value = False # require_flag=False + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth + ) + + # Should use key permissions when team is empty + assert sorted(result) == ["server1", "server2", "server3"] + + # Verify methods were called + mock_get_end_user.assert_called_once() + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth) + + async def test_end_user_with_permissions_still_enforces_intersection(self): + """ + Test: When end_user HAS explicit permissions, should still enforce intersection + with key/team restrictions (behavior unchanged). + """ + from litellm.proxy._types import (LiteLLM_EndUserTable, + LiteLLM_ObjectPermissionTable) + + # Create mock object permission for end_user + mock_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm-123", + mcp_servers=["server2", "server3", "server4"], + mcp_access_groups=[], + ) + + # Create mock end_user object WITH object_permission + mock_end_user = LiteLLM_EndUserTable( + user_id="test-end-user", + object_permission=mock_object_permission, + object_permission_id="perm-123", + ) + + # Create mock user auth + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + end_user_id="test-end-user", + ) + + # Mock the dependencies + mock_prisma = MagicMock() + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch("litellm.proxy.auth.auth_checks.get_end_user_object") as mock_get_end_user: + 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: + with patch.object( + MCPRequestHandler, "_get_mcp_servers_from_access_groups" + ) as mock_access_groups: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + # Configure mocks + mock_get_end_user.return_value = mock_end_user + mock_key_servers.return_value = ["server1", "server2", "server3"] + mock_team_servers.return_value = ["server1", "server2", "server3"] + mock_access_groups.return_value = [] # No access group servers + mock_general_settings.get.return_value = True + + # Call the method + result = await MCPRequestHandler.get_allowed_mcp_servers( + user_api_key_auth=mock_user_auth + ) + + # Should return intersection of key/team AND end_user + # key/team intersection = ["server1", "server2", "server3"] + # end_user = ["server2", "server3", "server4"] + # final intersection = ["server2", "server3"] + assert sorted(result) == ["server2", "server3"] + + # Verify methods were called + mock_get_end_user.assert_called_once() + mock_key_servers.assert_called_once_with(mock_user_auth) + mock_team_servers.assert_called_once_with(mock_user_auth)