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 dd987372df8..860569d24cb 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,8 +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 @@ -459,9 +463,11 @@ class MCPRequestHandler: 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, - user_api_key_cache) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + 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}" @@ -537,9 +543,14 @@ class MCPRequestHandler: # Intersect with agent's tool permissions if agent_id is set if user_api_key_auth.agent_id: + # Pre-fetch agent object_permission once to avoid duplicate DB query + agent_obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) agent_tools = await MCPRequestHandler._get_agent_tool_permissions_for_server( server_id=server_id, user_api_key_auth=user_api_key_auth, + agent_object_permission=agent_obj_perm, ) if agent_tools is not None: if allowed_tools is not None: @@ -611,11 +622,12 @@ class MCPRequestHandler: user_api_key_auth ) if key_object_permission is None and user_api_key_auth and user_api_key_auth.object_permission_id: - 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) + 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, + ) if prisma_client is not None: key_object_permission = await get_object_permission( object_permission_id=user_api_key_auth.object_permission_id, @@ -693,9 +705,11 @@ class MCPRequestHandler: Returns the MCP servers from the end_user's object_permission. """ from litellm.proxy.auth.auth_checks import get_end_user_object - from litellm.proxy.proxy_server import (prisma_client, - proxy_logging_obj, - user_api_key_cache) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if not user_api_key_auth or not user_api_key_auth.end_user_id: return [] @@ -742,36 +756,66 @@ class MCPRequestHandler: return [] @staticmethod - async def _get_allowed_mcp_servers_for_agent( + async def _get_agent_object_permission( user_api_key_auth: Optional[UserAPIKeyAuth] = None, - ) -> List[str]: + ): """ - Get allowed MCP servers for an agent (from the agent's object_permission). + Fetch the agent's object_permission from the DB (single query). - Returns the MCP servers from the agent's object_permission. - If agent has no object_permission, returns [] (no extra restriction). + Returns the object_permission object or None. """ from litellm.proxy.proxy_server import prisma_client if not user_api_key_auth or not user_api_key_auth.agent_id: - return [] + return None if prisma_client is None: verbose_logger.debug("prisma_client is None") - return [] + return None try: agent_row = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": user_api_key_auth.agent_id}, include={"object_permission": True}, ) - if ( - agent_row is None - or agent_row.object_permission is None - ): + if agent_row is None or agent_row.object_permission is None: + return None + + return agent_row.object_permission + except Exception as e: + verbose_logger.warning( + f"Failed to get agent object permission: {str(e)}" + ) + return None + + @staticmethod + async def _get_allowed_mcp_servers_for_agent( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + agent_object_permission=None, + ) -> List[str]: + """ + Get allowed MCP servers for an agent (from the agent's object_permission). + + Returns the MCP servers from the agent's object_permission. + If agent has no object_permission, returns [] (no extra restriction). + + Args: + user_api_key_auth: User auth with agent_id + agent_object_permission: Pre-fetched object_permission to avoid duplicate DB query. + If None, will be fetched from DB. + """ + if not user_api_key_auth or not user_api_key_auth.agent_id: + return [] + + try: + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + if obj_perm is None: return [] - obj_perm = agent_row.object_permission direct_mcp_servers = getattr(obj_perm, "mcp_servers", None) or [] if isinstance(direct_mcp_servers, str): direct_mcp_servers = [] @@ -796,28 +840,30 @@ class MCPRequestHandler: async def _get_agent_tool_permissions_for_server( server_id: str, user_api_key_auth: Optional[UserAPIKeyAuth] = None, + agent_object_permission=None, ) -> Optional[List[str]]: """ Get allowed tool names for a server from the agent's object_permission. Returns None if agent has no tool restrictions for this server. - """ - from litellm.proxy.proxy_server import prisma_client - if not user_api_key_auth or not user_api_key_auth.agent_id or not prisma_client: + Args: + server_id: Server ID to check permissions for + user_api_key_auth: User auth with agent_id + agent_object_permission: Pre-fetched object_permission to avoid duplicate DB query. + If None, will be fetched from DB. + """ + if not user_api_key_auth or not user_api_key_auth.agent_id: return None try: - agent_row = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": user_api_key_auth.agent_id}, - include={"object_permission": True}, - ) - if ( - agent_row is None - or agent_row.object_permission is None - ): + obj_perm = agent_object_permission + if obj_perm is None: + obj_perm = await MCPRequestHandler._get_agent_object_permission( + user_api_key_auth + ) + if obj_perm is None: return None - obj_perm = agent_row.object_permission mcp_tool_permissions = getattr( obj_perm, "mcp_tool_permissions", None ) @@ -880,8 +926,9 @@ class MCPRequestHandler: try: # Import here to avoid circular import - from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ - global_mcp_server_manager + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) # Use the new helper for config-loaded servers server_ids = MCPRequestHandler._get_config_server_ids_for_access_groups( @@ -935,9 +982,11 @@ class MCPRequestHandler: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: 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) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] @@ -973,9 +1022,11 @@ class MCPRequestHandler: Get MCP access groups for the team """ from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.proxy_server import (prisma_client, - proxy_logging_obj, - user_api_key_cache) + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if user_api_key_auth is None: return [] diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 4feddd67dd7..159c9fb93d9 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -6,7 +6,8 @@ from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.management_helpers.object_permission_utils import ( - attach_object_permission_to_dict, handle_update_object_permission_common) + handle_update_object_permission_common, +) from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -141,13 +142,16 @@ class AgentRegistry: # Create agent in DB created_agent = await prisma_client.db.litellm_agentstable.create( - data=create_data + data=create_data, + include={"object_permission": True}, ) created_agent_dict = created_agent.model_dump() - await attach_object_permission_to_dict( - created_agent_dict, prisma_client - ) + if created_agent.object_permission is not None: + try: + created_agent_dict["object_permission"] = created_agent.object_permission.model_dump() + except Exception: + created_agent_dict["object_permission"] = created_agent.object_permission.dict() return AgentResponse(**created_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error adding agent to DB: {str(e)}") @@ -232,11 +236,14 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), }, + include={"object_permission": True}, ) patched_agent_dict = patched_agent.model_dump() - await attach_object_permission_to_dict( - patched_agent_dict, prisma_client - ) + if patched_agent.object_permission is not None: + try: + patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump() + except Exception: + patched_agent_dict["object_permission"] = patched_agent.object_permission.dict() return AgentResponse(**patched_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error patching agent in DB: {str(e)}") @@ -305,12 +312,15 @@ class AgentRegistry: updated_agent = await prisma_client.db.litellm_agentstable.update( where={"agent_id": agent_id}, data=update_data, + include={"object_permission": True}, ) updated_agent_dict = updated_agent.model_dump() - await attach_object_permission_to_dict( - updated_agent_dict, prisma_client - ) + if updated_agent.object_permission is not None: + try: + updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump() + except Exception: + updated_agent_dict["object_permission"] = updated_agent.object_permission.dict() return AgentResponse(**updated_agent_dict) # type: ignore except Exception as e: raise Exception(f"Error updating agent in DB: {str(e)}") @@ -325,14 +335,18 @@ class AgentRegistry: try: agents_from_db = await prisma_client.db.litellm_agentstable.find_many( order={"created_at": "desc"}, + include={"object_permission": True}, ) agents: List[Dict[str, Any]] = [] for agent in agents_from_db: agent_dict = dict(agent) - await attach_object_permission_to_dict( - agent_dict, prisma_client - ) + # object_permission is eagerly loaded via include above + if agent.object_permission is not None: + try: + agent_dict["object_permission"] = agent.object_permission.model_dump() + except Exception: + agent_dict["object_permission"] = agent.object_permission.dict() agents.append(agent_dict) return agents diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 5200a1fc94d..b411b81b434 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -14,16 +14,19 @@ from fastapi import APIRouter, Depends, HTTPException, Request import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import (CommonProxyErrors, LitellmUserRoles, - UserAPIKeyAuth) +from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.management_endpoints.common_daily_activity import \ - get_daily_activity -from litellm.types.agents import (AgentConfig, AgentMakePublicResponse, - AgentResponse, MakeAgentsPublicRequest, - PatchAgentRequest) -from litellm.types.proxy.management_endpoints.common_daily_activity import \ - SpendAnalyticsPaginatedResponse +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.types.agents import ( + AgentConfig, + AgentMakePublicResponse, + AgentResponse, + MakeAgentsPublicRequest, + PatchAgentRequest, +) +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) router = APIRouter() @@ -49,10 +52,10 @@ async def get_agents( Returns: List[AgentResponse] """ - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ - AgentRequestHandler + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + AgentRequestHandler, + ) try: returned_agents: List[AgentResponse] = [] @@ -105,8 +108,9 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### -from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY +from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, +) @router.post( @@ -226,19 +230,19 @@ async def get_agent_by_id(agent_id: str): raise HTTPException(status_code=500, detail="Prisma client not initialized") try: - from litellm.proxy.management_helpers.object_permission_utils import \ - attach_object_permission_to_dict - agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id) if agent is None: agent_row = await prisma_client.db.litellm_agentstable.find_unique( - where={"agent_id": agent_id} + where={"agent_id": agent_id}, + include={"object_permission": True}, ) if agent_row is not None: agent_dict = agent_row.model_dump() - await attach_object_permission_to_dict( - agent_dict, prisma_client - ) + if agent_row.object_permission is not None: + try: + agent_dict["object_permission"] = agent_row.object_permission.model_dump() + except Exception: + agent_dict["object_permission"] = agent_row.object_permission.dict() agent = AgentResponse(**agent_dict) # type: ignore if agent is None: @@ -532,8 +536,9 @@ async def make_agent_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, + ) from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions @@ -648,8 +653,9 @@ async def make_agents_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import \ - global_agent_registry as AGENT_REGISTRY + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry as AGENT_REGISTRY, + ) from litellm.proxy.proxy_server import proxy_config # Load existing config 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 6232f7974da..b7ae33d1f80 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,8 +15,9 @@ 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 @@ -1175,8 +1176,7 @@ 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,8 +1277,7 @@ 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( @@ -1341,8 +1340,7 @@ 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( @@ -1707,9 +1705,10 @@ class TestAgentMCPPermissions: user_api_key_auth=user_api_key_auth, ) assert result == ["tool_a"] - mock_agent_tools.assert_called_once_with( - "server_1", user_api_key_auth - ) + mock_agent_tools.assert_called_once() + call_kwargs = mock_agent_tools.call_args.kwargs + assert call_kwargs["server_id"] == "server_1" + assert call_kwargs["user_api_key_auth"] == user_api_key_auth async def test_get_allowed_tools_for_server_agent_no_restriction(self): """Agent has no tool permissions for server; key/team result is unchanged."""