From dbf6c99e96878029c3bf64233ac7095d53aabc5c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 14 Mar 2026 17:00:23 -0700 Subject: [PATCH] fix(agents/): enforce sub-agent caller logic --- .../out/{404.html => 404/index.html} | 0 .../index.html} | 0 .../index.html} | 0 .../out/{chat.html => chat/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 | 149 +++++++++++--- litellm/types/agents.py | 1 + .../auth/test_agent_permission_handler.py | 190 +++++++++++++++++- 37 files changed, 303 insertions(+), 37 deletions(-) rename litellm/proxy/_experimental/out/{404.html => 404/index.html} (100%) rename litellm/proxy/_experimental/out/{_not-found.html => _not-found/index.html} (100%) rename litellm/proxy/_experimental/out/{api-reference.html => api-reference/index.html} (100%) rename litellm/proxy/_experimental/out/{chat.html => chat/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{api-playground.html => api-playground/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{budgets.html => budgets/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{caching.html => caching/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{claude-code-plugins.html => claude-code-plugins/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{old-usage.html => old-usage/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{prompts.html => prompts/index.html} (100%) rename litellm/proxy/_experimental/out/experimental/{tag-management.html => tag-management/index.html} (100%) rename litellm/proxy/_experimental/out/{guardrails.html => guardrails/index.html} (100%) rename litellm/proxy/_experimental/out/{login.html => login/index.html} (100%) rename litellm/proxy/_experimental/out/{logs.html => logs/index.html} (100%) rename litellm/proxy/_experimental/out/mcp/oauth/{callback.html => callback/index.html} (100%) rename litellm/proxy/_experimental/out/{model-hub.html => model-hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub.html => model_hub/index.html} (100%) rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) rename litellm/proxy/_experimental/out/{models-and-endpoints.html => models-and-endpoints/index.html} (100%) rename litellm/proxy/_experimental/out/{onboarding.html => onboarding/index.html} (100%) rename litellm/proxy/_experimental/out/{organizations.html => organizations/index.html} (100%) rename litellm/proxy/_experimental/out/{playground.html => playground/index.html} (100%) rename litellm/proxy/_experimental/out/{policies.html => policies/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{admin-settings.html => admin-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{logging-and-alerts.html => logging-and-alerts/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{router-settings.html => router-settings/index.html} (100%) rename litellm/proxy/_experimental/out/settings/{ui-theme.html => ui-theme/index.html} (100%) rename litellm/proxy/_experimental/out/{teams.html => teams/index.html} (100%) rename litellm/proxy/_experimental/out/{test-key.html => test-key/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{mcp-servers.html => mcp-servers/index.html} (100%) rename litellm/proxy/_experimental/out/tools/{vector-stores.html => vector-stores/index.html} (100%) rename litellm/proxy/_experimental/out/{usage.html => usage/index.html} (100%) rename litellm/proxy/_experimental/out/{users.html => users/index.html} (100%) rename litellm/proxy/_experimental/out/{virtual-keys.html => virtual-keys/index.html} (100%) diff --git a/litellm/proxy/_experimental/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/chat.html b/litellm/proxy/_experimental/out/chat/index.html similarity index 100% rename from litellm/proxy/_experimental/out/chat.html rename to litellm/proxy/_experimental/out/chat/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 42cf31e1e2b..aa0a5350634 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -8,12 +8,8 @@ Follows the same pattern as MCP permission handling. from typing import List, Optional, Set from litellm._logging import verbose_logger -from litellm.proxy._types import ( - LiteLLM_ObjectPermissionTable, - LiteLLM_TeamTable, - UI_TEAM_ID, - UserAPIKeyAuth, -) +from litellm.proxy._types import (UI_TEAM_ID, LiteLLM_ObjectPermissionTable, + LiteLLM_TeamTable, UserAPIKeyAuth) class AgentRequestHandler: @@ -28,6 +24,11 @@ class AgentRequestHandler: - If team has restrictions and key has none: inherit from team - If team has no restrictions: use key restrictions - If no restrictions: allow all agents + + Additionally supports agent-level sub-agent restrictions: + - If the key is scoped to an agent (agent_id on the key), the calling agent's + own object_permission is checked for allowed sub-agents. + - The agent-level restriction is intersected with the key/team result. """ @staticmethod @@ -37,6 +38,11 @@ class AgentRequestHandler: """ Get list of allowed agent IDs for the given user/key based on permissions. + Checks three layers: + 1. Key-level permissions + 2. Team-level permissions (intersected/inherited with key) + 3. Calling-agent-level permissions (intersected with key/team result) + Returns: List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). """ @@ -64,6 +70,11 @@ class AgentRequestHandler: else: allowed_agents = allowed_agents_for_key + # Apply calling-agent-level sub-agent restrictions + allowed_agents = await AgentRequestHandler._apply_calling_agent_restrictions( + allowed_agents, user_api_key_auth + ) + return list(set(allowed_agents)) except Exception as e: verbose_logger.warning(f"Failed to get allowed agents: {str(e)}") @@ -118,11 +129,9 @@ class AgentRequestHandler: 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) if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client: return None @@ -179,9 +188,8 @@ class AgentRequestHandler: # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) key_access_group_ids = user_api_key_auth.access_group_ids or [] if key_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + from litellm.proxy.auth.auth_checks import \ + _get_agent_ids_from_access_groups unified_agents = await _get_agent_ids_from_access_groups( access_group_ids=key_access_group_ids, @@ -213,11 +221,9 @@ class AgentRequestHandler: try: 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 not prisma_client: return [] @@ -254,9 +260,8 @@ class AgentRequestHandler: # 2. Also include agents from team's access_group_ids (unified access groups) team_access_group_ids = team_obj.access_group_ids or [] if team_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + from litellm.proxy.auth.auth_checks import \ + _get_agent_ids_from_access_groups unified_agents = await _get_agent_ids_from_access_groups( access_group_ids=team_access_group_ids, @@ -273,6 +278,87 @@ class AgentRequestHandler: ) return [] + @staticmethod + async def _apply_calling_agent_restrictions( + key_team_allowed: List[str], + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[str]: + """ + If the key is scoped to an agent (user_api_key_auth.agent_id), look up + that agent's object_permission to get its allowed sub-agents, then + intersect with the key/team result. + + Returns the narrowed list, or the original list if the calling agent + has no sub-agent restrictions. + """ + agent_level_allowed = ( + await AgentRequestHandler._get_allowed_agents_for_calling_agent( + user_api_key_auth + ) + ) + if len(agent_level_allowed) == 0: + return key_team_allowed + + if len(key_team_allowed) == 0: + return agent_level_allowed + + agent_level_set = set(agent_level_allowed) + return [a for a in key_team_allowed if a in agent_level_set] + + @staticmethod + async def _get_allowed_agents_for_calling_agent( + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[str]: + """ + Get allowed sub-agents for the calling agent. + + When a key is scoped to an agent (agent_id on the key), the agent + itself may have an object_permission restricting which other agents + (sub-agents) it can invoke. + + Reads the agent's object_permission from the registry and resolves + both direct agent IDs and agent_access_groups. + """ + if user_api_key_auth is None: + return [] + + calling_agent_id = getattr(user_api_key_auth, "agent_id", None) + if not calling_agent_id: + return [] + + try: + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + + agent = global_agent_registry.get_agent_by_id( + agent_id=calling_agent_id + ) + if agent is None: + return [] + + obj_perm = agent.object_permission + if not obj_perm or not isinstance(obj_perm, dict): + return [] + + all_agents: List[str] = [] + + direct_agents = obj_perm.get("agents") or [] + all_agents.extend(direct_agents) + + access_groups = obj_perm.get("agent_access_groups") or [] + if access_groups: + resolved = await AgentRequestHandler._get_agents_from_access_groups( + access_groups + ) + all_agents.extend(resolved) + + return list(set(all_agents)) + except Exception as e: + verbose_logger.warning( + f"Failed to get allowed agents for calling agent: {str(e)}" + ) + return [] + @staticmethod def _get_config_agent_ids_for_access_groups( config_agents: List, access_groups: List[str] @@ -314,7 +400,8 @@ class AgentRequestHandler: """ Resolve agent access groups to agent IDs by querying BOTH the agent table (DB) AND config-loaded agents. """ - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry from litellm.proxy.proxy_server import prisma_client try: @@ -371,11 +458,9 @@ class AgentRequestHandler: ) -> List[str]: """Get agent access groups for the key.""" 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 [] @@ -411,11 +496,9 @@ class AgentRequestHandler: ) -> List[str]: """Get agent 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/types/agents.py b/litellm/types/agents.py index efb2e73bfb5..974de8cbc12 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -174,6 +174,7 @@ class AgentObjectPermission(TypedDict, total=False): mcp_tool_permissions: Optional[Dict[str, List[str]]] models: Optional[List[str]] agents: Optional[List[str]] + agent_access_groups: Optional[List[str]] class AgentConfig(TypedDict, total=False): diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 533dc0557b5..063f2e93bf6 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -4,16 +4,15 @@ Unit tests for AgentRequestHandler - Agent permission management for keys and te import os import sys -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest sys.path.insert(0, os.path.abspath("../../../..")) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, -) +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler @pytest.mark.asyncio @@ -165,3 +164,186 @@ class TestAgentRequestHandler: user_api_key_auth=mock_user_auth ) assert sorted(result) == ["agent-from-ag", "native-agent-1"] + + async def test_calling_agent_restricts_allowed_sub_agents(self): + """ + When a key is scoped to an agent that has sub-agent restrictions, + the result is intersected with key/team permissions. + """ + from litellm.types.agents import AgentResponse + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="parent-agent", + ) + + parent_agent = AgentResponse( + agent_id="parent-agent", + agent_name="parent", + agent_card_params={"name": "parent"}, + object_permission={ + "agents": ["sub-agent-1", "sub-agent-2"], + }, + ) + + mock_registry = MagicMock() + mock_registry.get_agent_by_id.return_value = parent_agent + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + mock_registry, + ): + mock_key.return_value = [] + mock_team.return_value = [] + + result = await AgentRequestHandler.get_allowed_agents( + user_api_key_auth=mock_user_auth + ) + assert sorted(result) == ["sub-agent-1", "sub-agent-2"] + + async def test_calling_agent_intersects_with_key_team(self): + """ + When both key/team and calling agent have restrictions, + the final result is the intersection. + """ + from litellm.types.agents import AgentResponse + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="parent-agent", + ) + + parent_agent = AgentResponse( + agent_id="parent-agent", + agent_name="parent", + agent_card_params={"name": "parent"}, + object_permission={ + "agents": ["sub-agent-1", "sub-agent-2", "sub-agent-3"], + }, + ) + + mock_registry = MagicMock() + mock_registry.get_agent_by_id.return_value = parent_agent + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + mock_registry, + ): + mock_key.return_value = ["sub-agent-1", "sub-agent-4"] + mock_team.return_value = [] + + result = await AgentRequestHandler.get_allowed_agents( + user_api_key_auth=mock_user_auth + ) + assert result == ["sub-agent-1"] + + async def test_calling_agent_no_restrictions_passes_through(self): + """ + When the calling agent has no object_permission (no sub-agent restrictions), + key/team permissions pass through unchanged. + """ + from litellm.types.agents import AgentResponse + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="parent-agent", + ) + + parent_agent = AgentResponse( + agent_id="parent-agent", + agent_name="parent", + agent_card_params={"name": "parent"}, + object_permission=None, + ) + + mock_registry = MagicMock() + mock_registry.get_agent_by_id.return_value = parent_agent + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + mock_registry, + ): + mock_key.return_value = ["agent-a", "agent-b"] + mock_team.return_value = [] + + result = await AgentRequestHandler.get_allowed_agents( + user_api_key_auth=mock_user_auth + ) + assert sorted(result) == ["agent-a", "agent-b"] + + async def test_calling_agent_with_access_groups(self): + """ + When the calling agent's object_permission uses agent_access_groups, + those are resolved to agent IDs and used for intersection. + """ + from litellm.types.agents import AgentResponse + + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="parent-agent", + ) + + parent_agent = AgentResponse( + agent_id="parent-agent", + agent_name="parent", + agent_card_params={"name": "parent"}, + object_permission={ + "agents": ["sub-agent-1"], + "agent_access_groups": ["group-a"], + }, + ) + + mock_registry = MagicMock() + mock_registry.get_agent_by_id.return_value = parent_agent + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + mock_registry, + ): + with patch.object( + AgentRequestHandler, + "_get_agents_from_access_groups", + new_callable=AsyncMock, + return_value=["sub-agent-from-group"], + ): + mock_key.return_value = [] + mock_team.return_value = [] + + result = await AgentRequestHandler.get_allowed_agents( + user_api_key_auth=mock_user_auth + ) + assert sorted(result) == [ + "sub-agent-1", + "sub-agent-from-group", + ] + + async def test_no_agent_id_on_key_skips_agent_level_check(self): + """ + When the key has no agent_id, the agent-level check is skipped entirely. + """ + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + ) + + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: + mock_key.return_value = ["agent-x"] + mock_team.return_value = [] + + result = await AgentRequestHandler.get_allowed_agents( + user_api_key_auth=mock_user_auth + ) + assert result == ["agent-x"]