From 92f0089c0dd966c99f36bc5adda66d4185cf5fb6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 10 Feb 2026 09:52:02 -0800 Subject: [PATCH] refactor: migrate get object permissions table logic to happen in user api key auth - allows functions to trust user api key object they receive has what they need --- .../mcp_server/auth/user_api_key_auth_mcp.py | 80 ++++------ .../mcp_server/rest_endpoints.py | 19 ++- .../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 .../auth/agent_permission_handler.py | 148 +++++++---------- litellm/proxy/auth/auth_checks.py | 46 ++++++ .../auth/test_user_api_key_auth_mcp.py | 129 ++++++--------- .../auth/test_object_permission_loading.py | 151 ++++++++++++++++++ 39 files changed, 347 insertions(+), 226 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%) create mode 100644 tests/test_litellm/proxy/auth/test_object_permission_loading.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 786cfbfb008..548e3bc3dbf 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 @@ -6,7 +6,12 @@ from starlette.requests import Request from starlette.types import Scope from litellm._logging import verbose_logger -from litellm.proxy._types import LiteLLM_TeamTable, ProxyException, SpecialHeaders, UserAPIKeyAuth +from litellm.proxy._types import ( + LiteLLM_TeamTable, + ProxyException, + SpecialHeaders, + UserAPIKeyAuth, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -372,45 +377,31 @@ class MCPRequestHandler: return [] @staticmethod - async def _get_key_object_permission( + def _get_key_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ): - """Helper to get key object_permission from cache or DB.""" - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + """ + Get key object_permission - already loaded by get_key_object() in main auth flow. + Note: object_permission is automatically populated when the key is fetched via + get_key_object() in litellm/proxy/auth/auth_checks.py + """ if not user_api_key_auth: return None - # Already loaded - if user_api_key_auth.object_permission: - return user_api_key_auth.object_permission - - # Need to fetch from DB - if user_api_key_auth.object_permission_id and prisma_client: - return await get_object_permission( - object_permission_id=user_api_key_auth.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - return None + return user_api_key_auth.object_permission @staticmethod async def _get_team_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ): - """Helper to get team object_permission from cache or DB.""" - from litellm.proxy.auth.auth_checks import ( - get_object_permission, - get_team_object, - ) + """ + Get team object_permission - automatically loaded by get_team_object() in main auth flow. + + Note: object_permission is automatically populated when the team is fetched via + get_team_object() in litellm/proxy/auth/auth_checks.py + """ + from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, @@ -423,7 +414,7 @@ class MCPRequestHandler: if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: return None - # First get the team object (which may have object_permission already loaded) + # Get the team object (which has object_permission already loaded) team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( team_id=user_api_key_auth.team_id, prisma_client=prisma_client, @@ -435,21 +426,7 @@ class MCPRequestHandler: if not team_obj: return None - # Already loaded - if team_obj.object_permission: - return team_obj.object_permission - - # Need to fetch from DB using object_permission_id - if team_obj.object_permission_id: - return await get_object_permission( - object_permission_id=team_obj.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - return None + return team_obj.object_permission @staticmethod async def get_allowed_tools_for_server( @@ -471,8 +448,8 @@ class MCPRequestHandler: return None try: - # Get key and team object permissions - key_obj_perm = await MCPRequestHandler._get_key_object_permission( + # Get key and team object permissions (already loaded in main auth flow) + key_obj_perm = MCPRequestHandler._get_key_object_permission( user_api_key_auth ) team_obj_perm = await MCPRequestHandler._get_team_object_permission( @@ -559,7 +536,8 @@ class MCPRequestHandler: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: try: - key_object_permission = await MCPRequestHandler._get_key_object_permission( + # Get key object permission (already loaded in main auth flow) + key_object_permission = MCPRequestHandler._get_key_object_permission( user_api_key_auth ) if key_object_permission is None: @@ -591,12 +569,10 @@ class MCPRequestHandler: """ Get allowed MCP servers for a team. - Uses the helper _get_team_object_permission which: - 1. First checks if object_permission is already loaded on the team - 2. If not, fetches from DB using object_permission_id if it exists + Note: object_permission is automatically loaded by get_team_object() in main auth flow. """ try: - # Use the helper method that properly handles fetching from DB if needed + # Get team object permission (already loaded in main auth flow) object_permissions = await MCPRequestHandler._get_team_object_permission( user_api_key_auth ) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 65d3a2caa16..ed558cf3c8e 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -159,6 +159,7 @@ if MCP_AVAILABLE: server, server_auth_header, raw_headers: Optional[Dict[str, str]] = None, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, ): """Helper function to get tools for a single server.""" tools = await global_mcp_server_manager._get_tools_from_server( @@ -197,9 +198,7 @@ if MCP_AVAILABLE: ) allowed_mcp_servers: List[MCPServer] = [] for allowed_server_id in allowed_server_ids_set: - server = global_mcp_server_manager.get_mcp_server_by_id( - allowed_server_id - ) + server = global_mcp_server_manager.get_mcp_server_by_id(allowed_server_id) if server is not None: allowed_mcp_servers.append(server) return allowed_mcp_servers @@ -276,9 +275,7 @@ if MCP_AVAILABLE: "message": f"The key is not allowed to access server {server_id}", }, ) - server = global_mcp_server_manager.get_mcp_server_by_id( - server_id - ) + server = global_mcp_server_manager.get_mcp_server_by_id(server_id) if server is None: return { "tools": [], @@ -292,7 +289,10 @@ if MCP_AVAILABLE: try: list_tools_result = await _get_tools_for_single_server( - server, server_auth_header, raw_headers_from_request + server, + server_auth_header, + raw_headers_from_request, + user_api_key_dict, ) except Exception as e: verbose_logger.exception( @@ -328,7 +328,10 @@ if MCP_AVAILABLE: try: tools_result = await _get_tools_for_single_server( - server, server_auth_header, raw_headers_from_request + server, + server_auth_header, + raw_headers_from_request, + user_api_key_dict, ) list_tools_result.extend(tools_result) 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/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 96e7a21cc33..bf3256cf47b 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -21,7 +21,7 @@ class AgentRequestHandler: 1. Key-level agent permissions 2. Team-level agent permissions 3. Agent access group resolution - + Follows the same inheritance logic as MCP: - If team has restrictions and key has restrictions: use intersection - If team has restrictions and key has none: inherit from team @@ -35,7 +35,7 @@ class AgentRequestHandler: ) -> List[str]: """ Get list of allowed agent IDs for the given user/key based on permissions. - + Returns: List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). """ @@ -45,7 +45,9 @@ class AgentRequestHandler: await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) ) allowed_agents_for_team = ( - await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) + await AgentRequestHandler._get_allowed_agents_for_team( + user_api_key_auth + ) ) # If team has agent restrictions, handle inheritance and intersection logic @@ -73,62 +75,48 @@ class AgentRequestHandler: ) -> bool: """ Check if a specific agent is allowed for the given user/key. - + Args: agent_id: The agent ID to check user_api_key_auth: User authentication info - + Returns: bool: True if agent is allowed, False otherwise """ allowed_agents = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) - + # Empty list means no restrictions - allow all if len(allowed_agents) == 0: return True - + return agent_id in allowed_agents @staticmethod - async def _get_key_object_permission( + def _get_key_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> Optional[LiteLLM_ObjectPermissionTable]: - """Helper to get key object_permission from cache or DB.""" - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + """ + Get key object_permission - already loaded by get_key_object() in main auth flow. + Note: object_permission is automatically populated when the key is fetched via + get_key_object() in litellm/proxy/auth/auth_checks.py + """ if not user_api_key_auth: return None - # Already loaded - if user_api_key_auth.object_permission: - return user_api_key_auth.object_permission - - # Need to fetch from DB - if user_api_key_auth.object_permission_id and prisma_client: - return await get_object_permission( - object_permission_id=user_api_key_auth.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - return None + return user_api_key_auth.object_permission @staticmethod async def _get_team_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> Optional[LiteLLM_ObjectPermissionTable]: - """Helper to get team object_permission from cache or DB.""" - from litellm.proxy.auth.auth_checks import ( - get_object_permission, - get_team_object, - ) + """ + Get team object_permission - automatically loaded by get_team_object() in main auth flow. + + Note: object_permission is automatically populated when the team is fetched via + get_team_object() in litellm/proxy/auth/auth_checks.py + """ + from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.proxy_server import ( prisma_client, proxy_logging_obj, @@ -138,7 +126,7 @@ class AgentRequestHandler: if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: return None - # First get the team object (which may have object_permission already loaded) + # Get the team object (which has object_permission already loaded) team_obj: Optional[LiteLLM_TeamTable] = await get_team_object( team_id=user_api_key_auth.team_id, prisma_client=prisma_client, @@ -150,21 +138,7 @@ class AgentRequestHandler: if not team_obj: return None - # Already loaded - if team_obj.object_permission: - return team_obj.object_permission - - # Need to fetch from DB using object_permission_id - if team_obj.object_permission_id: - return await get_object_permission( - object_permission_id=team_obj.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - - return None + return team_obj.object_permission @staticmethod async def _get_allowed_agents_for_key( @@ -172,31 +146,16 @@ class AgentRequestHandler: ) -> List[str]: """ Get allowed agents for a key from its object_permission. - """ - from litellm.proxy.auth.auth_checks import get_object_permission - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - user_api_key_cache, - ) + Note: object_permission is already loaded by get_key_object() in main auth flow. + """ if user_api_key_auth is None: return [] - if user_api_key_auth.object_permission_id is None: - return [] - - if prisma_client is None: - verbose_logger.debug("prisma_client is None") - return [] - try: - key_object_permission = await get_object_permission( - object_permission_id=user_api_key_auth.object_permission_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=user_api_key_auth.parent_otel_span, - proxy_logging_obj=proxy_logging_obj, + # Get key object permission (already loaded in main auth flow) + key_object_permission = AgentRequestHandler._get_key_object_permission( + user_api_key_auth ) if key_object_permission is None: return [] @@ -205,8 +164,10 @@ class AgentRequestHandler: direct_agents = key_object_permission.agents or [] # Get agents from access groups - access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] + access_group_agents = ( + await AgentRequestHandler._get_agents_from_access_groups( + key_object_permission.agent_access_groups or [] + ) ) # Combine both lists @@ -222,6 +183,8 @@ class AgentRequestHandler: ) -> List[str]: """ Get allowed agents for a team from its object_permission. + + Note: object_permission is already loaded by get_team_object() in main auth flow. """ if user_api_key_auth is None: return [] @@ -230,7 +193,7 @@ class AgentRequestHandler: return [] try: - # Use the helper method that properly handles fetching from DB if needed + # Get team object permission (already loaded in main auth flow) object_permissions = await AgentRequestHandler._get_team_object_permission( user_api_key_auth ) @@ -242,8 +205,10 @@ class AgentRequestHandler: direct_agents = object_permissions.agents or [] # Get agents from access groups - access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] + access_group_agents = ( + await AgentRequestHandler._get_agents_from_access_groups( + object_permissions.agent_access_groups or [] + ) ) # Combine both lists @@ -284,9 +249,7 @@ class AgentRequestHandler: for agent in agents: agent_ids.add(agent.agent_id) except Exception as e: - verbose_logger.debug( - f"Error getting agents from access groups: {e}" - ) + verbose_logger.debug(f"Error getting agents from access groups: {e}") return agent_ids @staticmethod @@ -306,16 +269,16 @@ class AgentRequestHandler: ) # Use the helper for DB agents - db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups( - prisma_client, access_groups + db_agent_ids = ( + await AgentRequestHandler._get_db_agent_ids_for_access_groups( + prisma_client, access_groups + ) ) agent_ids.update(db_agent_ids) return list(agent_ids) except Exception as e: - verbose_logger.warning( - f"Failed to get agents from access groups: {str(e)}" - ) + verbose_logger.warning(f"Failed to get agents from access groups: {str(e)}") return [] @staticmethod @@ -326,11 +289,15 @@ class AgentRequestHandler: Get list of agent access groups for the given user/key based on permissions. """ access_groups: List[str] = [] - access_groups_for_key = await AgentRequestHandler._get_agent_access_groups_for_key( - user_api_key_auth + access_groups_for_key = ( + await AgentRequestHandler._get_agent_access_groups_for_key( + user_api_key_auth + ) ) - access_groups_for_team = await AgentRequestHandler._get_agent_access_groups_for_team( - user_api_key_auth + access_groups_for_team = ( + await AgentRequestHandler._get_agent_access_groups_for_team( + user_api_key_auth + ) ) # If team has access groups, then key must have a subset of the team's access groups @@ -378,7 +345,9 @@ class AgentRequestHandler: return key_object_permission.agent_access_groups or [] except Exception as e: - verbose_logger.warning(f"Failed to get agent access groups for key: {str(e)}") + verbose_logger.warning( + f"Failed to get agent access groups for key: {str(e)}" + ) return [] @staticmethod @@ -425,4 +394,3 @@ class AgentRequestHandler: f"Failed to get agent access groups for team: {str(e)}" ) return [] - diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c6093172932..76ec67ab10e 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1368,6 +1368,22 @@ async def _get_team_object_from_user_api_key_cache( raise Exception _response = LiteLLM_TeamTableCachedObj(**response.dict()) + + # Load object_permission if object_permission_id exists but object_permission is not loaded + if _response.object_permission_id and not _response.object_permission: + try: + _response.object_permission = await get_object_permission( + object_permission_id=_response.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for team {team_id} with object_permission_id={_response.object_permission_id}: {e}" + ) + # save the team object to cache await _cache_team_object( team_id=team_id, @@ -1550,6 +1566,21 @@ async def get_team_object_by_alias( team = teams[0] team_obj = LiteLLM_TeamTableCachedObj(**team.model_dump()) + # Load object_permission if object_permission_id exists but object_permission is not loaded + if team_obj.object_permission_id and not team_obj.object_permission: + try: + team_obj.object_permission = await get_object_permission( + object_permission_id=team_obj.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for team {team_obj.team_id} with object_permission_id={team_obj.object_permission_id}: {e}" + ) + # Cache the result by both alias and team_id await user_api_key_cache.async_set_cache( key=cache_key, @@ -1838,6 +1869,21 @@ async def get_key_object( _response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True)) + # Load object_permission if object_permission_id exists but object_permission is not loaded + if _response.object_permission_id and not _response.object_permission: + try: + _response.object_permission = await get_object_permission( + object_permission_id=_response.object_permission_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except Exception as e: + verbose_proxy_logger.debug( + f"Failed to load object_permission for key with object_permission_id={_response.object_permission_id}: {e}" + ) + # save the key object to cache await _cache_key_object( hashed_token=hashed_token, 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 1ed21b07bb7..c2dbc94f721 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 @@ -331,11 +331,7 @@ class TestMCPRequestHandler: # Create an async mock for user_api_key_auth async def mock_user_api_key_auth(api_key, request): return UserAPIKeyAuth( - token=( - "test-token-sha256-empty-hash" - if api_key - else None - ), + token=("test-token-sha256-empty-hash" if api_key else None), api_key=api_key, user_id="test-user-id" if api_key else None, team_id="test-team-id" if api_key else None, @@ -632,8 +628,7 @@ class TestMCPOAuth2AuthFlow: # OAuth2 headers should still contain the Authorization token assert ( - oauth2_headers.get("Authorization") - == "Bearer atlassian-oauth2-token" + oauth2_headers.get("Authorization") == "Bearer atlassian-oauth2-token" ) async def test_litellm_key_in_authorization_backward_compat(self): @@ -1291,21 +1286,21 @@ async def test_get_team_object_permission_with_already_loaded_permission(): mcp_access_groups=["group1"], vector_stores=["store1"], ) - + # Create mock team object with object_permission already loaded mock_team_obj = LiteLLM_TeamTable( team_id="team-123", 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="team-123", ) - + # Mock get_team_object to return our team with loaded permission # Also need to mock prisma_client from proxy_server mock_prisma = MagicMock() @@ -1313,96 +1308,81 @@ async def test_get_team_object_permission_with_already_loaded_permission(): "litellm.proxy.proxy_server.prisma_client", mock_prisma, ): - with patch( - "litellm.proxy.auth.auth_checks.get_team_object" - ) as mock_get_team: + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: with patch( "litellm.proxy.auth.auth_checks.get_object_permission" ) as mock_get_perm: mock_get_team.return_value = mock_team_obj - + # Call the method result = await MCPRequestHandler._get_team_object_permission( mock_user_auth ) - + # Assert we got the object permission assert result == mock_object_permission assert result.mcp_servers == ["server1", "server2"] - + # Verify get_team_object was called mock_get_team.assert_called_once() - + # Verify get_object_permission was NOT called (since it was already loaded) mock_get_perm.assert_not_called() @pytest.mark.asyncio -async def test_get_team_object_permission_fetches_from_db_when_not_loaded(): +async def test_get_team_object_permission_with_core_auth_auto_loading(): """ - Test that _get_team_object_permission fetches from DB when object_permission - is not loaded but object_permission_id exists. + Test that _get_team_object_permission returns the object_permission that was + automatically loaded by get_team_object() in the core auth flow. + + Note: After migrating permission loading to core auth (get_team_object in auth_checks.py), + 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 - # Create mock object permission (to be returned from DB) + # Create mock object permission mock_object_permission = LiteLLM_ObjectPermissionTable( object_permission_id="perm-456", mcp_servers=["server3", "server4"], mcp_access_groups=["group2"], vector_stores=["store2"], ) - - # Create mock team object WITHOUT object_permission loaded (but has ID) + + # Create mock team object WITH object_permission already loaded + # (This is what get_team_object() returns after the core auth migration) mock_team_obj = LiteLLM_TeamTable( team_id="team-456", - object_permission=None, + object_permission=mock_object_permission, # Already loaded by core auth object_permission_id="perm-456", ) - + # Create mock user auth mock_user_auth = UserAPIKeyAuth( api_key="test-key", user_id="test-user", team_id="team-456", ) - + # Mock the methods - # Also need to mock prisma_client from proxy_server mock_prisma = MagicMock() with patch( "litellm.proxy.proxy_server.prisma_client", mock_prisma, ): - with patch( - "litellm.proxy.auth.auth_checks.get_team_object" - ) as mock_get_team: - with patch( - "litellm.proxy.auth.auth_checks.get_object_permission" - ) as mock_get_perm: - mock_get_team.return_value = mock_team_obj - mock_get_perm.return_value = mock_object_permission - - # Call the method - result = await MCPRequestHandler._get_team_object_permission( - mock_user_auth - ) - - # Assert we got the object permission - assert result == mock_object_permission - assert result.mcp_servers == ["server3", "server4"] - - # Verify get_team_object was called - mock_get_team.assert_called_once() - - # Verify get_object_permission WAS called (since it wasn't loaded) - mock_get_perm.assert_called_once_with( - object_permission_id="perm-456", - prisma_client=mock.ANY, - user_api_key_cache=mock.ANY, - parent_otel_span=mock_user_auth.parent_otel_span, - proxy_logging_obj=mock.ANY, - ) + with patch("litellm.proxy.auth.auth_checks.get_team_object") as mock_get_team: + mock_get_team.return_value = mock_team_obj + + # Call the method + result = await MCPRequestHandler._get_team_object_permission(mock_user_auth) + + # Assert we got the object permission (already loaded by core auth) + assert result == mock_object_permission + assert result.mcp_servers == ["server3", "server4"] + + # Verify get_team_object was called + mock_get_team.assert_called_once() @pytest.mark.asyncio @@ -1420,14 +1400,14 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): mcp_access_groups=["dev-group"], vector_stores=[], ) - + # Create mock user auth mock_user_auth = UserAPIKeyAuth( api_key="test-key", user_id="test-user", team_id="team-789", ) - + # Mock the helper methods with patch.object( MCPRequestHandler, "_get_team_object_permission" @@ -1437,13 +1417,16 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): ) as mock_get_access_group_servers: # Configure mocks mock_get_team_perm.return_value = mock_object_permission - mock_get_access_group_servers.return_value = ["group-server1", "group-server2"] - + mock_get_access_group_servers.return_value = [ + "group-server1", + "group-server2", + ] + # Call the method result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( mock_user_auth ) - + # Assert the result contains both direct and access group servers assert set(result) == { "direct-server1", @@ -1451,10 +1434,10 @@ async def test_get_allowed_mcp_servers_for_team_uses_helper(): "group-server1", "group-server2", } - + # Verify _get_team_object_permission was called (the helper we fixed) mock_get_team_perm.assert_called_once_with(mock_user_auth) - + # Verify access groups were resolved mock_get_access_group_servers.assert_called_once_with(["dev-group"]) @@ -1471,21 +1454,21 @@ async def test_get_allowed_mcp_servers_for_team_with_no_object_permission(): user_id="test-user", team_id="team-no-perm", ) - + # Mock the helper to return None (no object permission) with patch.object( MCPRequestHandler, "_get_team_object_permission" ) as mock_get_team_perm: mock_get_team_perm.return_value = None - + # Call the method result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( mock_user_auth ) - + # Assert empty list is returned assert result == [] - + # Verify the helper was called mock_get_team_perm.assert_called_once_with(mock_user_auth) @@ -1509,9 +1492,7 @@ async def test_get_allowed_mcp_servers_for_team_without_team_id_returns_empty(): team_id=None, ) - result = await MCPRequestHandler._get_allowed_mcp_servers_for_team( - mock_user_auth - ) + result = await MCPRequestHandler._get_allowed_mcp_servers_for_team(mock_user_auth) assert result == [] @@ -1546,9 +1527,7 @@ async def test_get_allowed_mcp_servers_for_key_guard_conditions( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, ) as mock_get_perm: - with patch( - "litellm.proxy.proxy_server.prisma_client", prisma_client_value - ): + with patch("litellm.proxy.proxy_server.prisma_client", prisma_client_value): result = await MCPRequestHandler._get_allowed_mcp_servers_for_key( user_api_key_auth ) @@ -1569,9 +1548,7 @@ async def test_get_allowed_mcp_servers_for_key_returns_empty_when_db_returns_non mock_prisma = object() - with patch( - "litellm.proxy.proxy_server.prisma_client", mock_prisma - ), patch( + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), patch( "litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock, ) as mock_get_perm: diff --git a/tests/test_litellm/proxy/auth/test_object_permission_loading.py b/tests/test_litellm/proxy/auth/test_object_permission_loading.py new file mode 100644 index 00000000000..54e4c82471e --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_object_permission_loading.py @@ -0,0 +1,151 @@ +""" +Test that object_permission is automatically loaded when fetching keys and teams. +""" +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.proxy._types import ( + LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTableCachedObj, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import get_key_object, get_team_object + + +@pytest.mark.asyncio +async def test_get_key_object_loads_object_permission(): + """ + Test that get_key_object automatically loads object_permission when object_permission_id exists. + """ + # Mock prisma client + mock_prisma_client = MagicMock() + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) # Not in cache + + # Mock the DB response with object_permission_id but no object_permission + mock_token_data = MagicMock() + mock_token_data.model_dump.return_value = { + "token": "test_token_hash", + "user_id": "test_user", + "object_permission_id": "test_perm_id", + "object_permission": None, + } + mock_prisma_client.get_data = AsyncMock(return_value=mock_token_data) + + # Mock the object_permission that should be loaded + mock_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="test_perm_id", + mcp_servers=["server1", "server2"], + vector_stores=["store1"], + ) + + # Mock get_object_permission to return the permission + with patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + AsyncMock(return_value=mock_object_permission) + ), patch( + "litellm.proxy.auth.auth_checks._cache_key_object", + AsyncMock() + ): + result = await get_key_object( + hashed_token="test_token_hash", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + # Verify that object_permission was loaded + assert result.object_permission is not None + assert result.object_permission.object_permission_id == "test_perm_id" + assert result.object_permission.mcp_servers == ["server1", "server2"] + + +@pytest.mark.asyncio +async def test_get_key_object_no_permission_id(): + """ + Test that get_key_object works correctly when no object_permission_id exists. + """ + # Mock prisma client + mock_prisma_client = MagicMock() + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) # Not in cache + + # Mock the DB response without object_permission_id + mock_token_data = MagicMock() + mock_token_data.model_dump.return_value = { + "token": "test_token_hash", + "user_id": "test_user", + "object_permission_id": None, + "object_permission": None, + } + mock_prisma_client.get_data = AsyncMock(return_value=mock_token_data) + + with patch( + "litellm.proxy.auth.auth_checks._cache_key_object", + AsyncMock() + ): + result = await get_key_object( + hashed_token="test_token_hash", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + # Verify that object_permission is None + assert result.object_permission is None + + +@pytest.mark.asyncio +async def test_get_team_object_loads_object_permission(): + """ + Test that get_team_object automatically loads object_permission when object_permission_id exists. + """ + # Mock prisma client + mock_prisma_client = MagicMock() + mock_cache = MagicMock() + mock_cache.async_get_cache = AsyncMock(return_value=None) # Not in cache + + # Mock team data with object_permission_id + mock_team = MagicMock() + mock_team.dict.return_value = { + "team_id": "test_team", + "team_alias": "Test Team", + "object_permission_id": "test_perm_id", + "object_permission": None, + } + + # Mock the object_permission that should be loaded + mock_object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="test_perm_id", + mcp_servers=["team_server1"], + vector_stores=["team_store1"], + ) + + with patch( + "litellm.proxy.auth.auth_checks._get_team_db_check", + AsyncMock(return_value=mock_team) + ), patch( + "litellm.proxy.auth.auth_checks.get_object_permission", + AsyncMock(return_value=mock_object_permission) + ), patch( + "litellm.proxy.auth.auth_checks._cache_team_object", + AsyncMock() + ), patch( + "litellm.proxy.auth.auth_checks._should_check_db", + return_value=True + ), patch( + "litellm.proxy.auth.auth_checks._update_last_db_access_time" + ): + result = await get_team_object( + team_id="test_team", + prisma_client=mock_prisma_client, + user_api_key_cache=mock_cache, + ) + + # Verify that object_permission was loaded + assert result.object_permission is not None + assert result.object_permission.object_permission_id == "test_perm_id" + assert result.object_permission.mcp_servers == ["team_server1"]