fix: eliminate duplicate DB queries in MCP agent auth and N+1 in agent listing

- Extract _get_agent_object_permission helper so _get_allowed_mcp_servers_for_agent
  and _get_agent_tool_permissions_for_server share a single DB fetch instead of
  each independently querying the same agent row (was 1+N queries per MCP request)
- Use include={"object_permission": True} on find_many in get_all_agents_from_db
  to eagerly load permissions in one query instead of N+1
- Use include={"object_permission": True} on create/update/find_unique in all
  agent CRUD operations, removing attach_object_permission_to_dict follow-up calls

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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 []

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

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

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."""