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

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
             )

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:

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

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,

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:

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