diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 8acff11b009..15a74a34664 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,6 +5,9 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ +from dataclasses import dataclass +from typing import assert_never + from litellm._logging import verbose_logger from litellm.proxy._types import ( UI_TEAM_ID, @@ -15,6 +18,51 @@ from litellm.proxy._types import ( from litellm.repositories.table_repositories import AgentsRepository +@dataclass(frozen=True, slots=True) +class UnrestrictedAgents: + """The level declared no agent restriction, so it does not narrow the caller.""" + + +@dataclass(frozen=True, slots=True) +class RestrictedAgents: + """The exact agents the caller may reach. An empty set reaches none of them.""" + + agent_ids: frozenset[str] + + +@dataclass(frozen=True, slots=True) +class UnresolvableAgents: + """A restriction is declared but could not be read, so it cannot be enforced.""" + + reason: str + + +AgentScope = UnrestrictedAgents | RestrictedAgents | UnresolvableAgents + + +def _combine_key_and_team_scope(key_scope: AgentScope, team_scope: AgentScope) -> AgentScope: + """Resolve the key and team levels into the scope a request is authorized against. + + A level that declares nothing does not narrow the other one. Two levels that both + declare narrow each other, and a pair that shares no agent leaves the caller with + none rather than with all: "restricted to nothing" and "not restricted" are + different answers, and only the second one may widen access. + """ + match (key_scope, team_scope): + case (UnresolvableAgents(), _): + return key_scope + case (_, UnresolvableAgents()): + return team_scope + case (UnrestrictedAgents(), _): + return team_scope + case (RestrictedAgents(), UnrestrictedAgents()): + return key_scope + case (RestrictedAgents(key_ids), RestrictedAgents(team_ids)): + return RestrictedAgents(key_ids & team_ids) + case _: + assert_never(key_scope) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -27,40 +75,23 @@ 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 + - If the intersection is empty, or a declared restriction cannot be read: allow none """ @staticmethod - async def get_allowed_agents( + async def resolve_agent_scope( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentScope: """ - Get list of allowed agent IDs for the given user/key based on permissions. + Resolve the agents the given key/team may reach. Returns: - List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). + AgentScope: whether the caller is unrestricted, restricted to a specific set + (possibly empty), or carries a restriction that could not be read. """ - try: - allowed_agents: list[str] = [] - allowed_agents_for_key = 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) - - # If team has agent restrictions, handle inheritance and intersection logic - if len(allowed_agents_for_team) > 0: - if len(allowed_agents_for_key) > 0: - # Key has its own agent permissions - use intersection with team permissions - for agent_id in allowed_agents_for_key: - if agent_id in allowed_agents_for_team: - allowed_agents.append(agent_id) - else: - # Key has no agent permissions - inherit from team - allowed_agents = allowed_agents_for_team - else: - allowed_agents = allowed_agents_for_key - - return list(set(allowed_agents)) - except Exception as e: - verbose_logger.warning(f"Failed to get allowed agents: {e!s}") - return [] + key_scope = await AgentRequestHandler._resolve_agents_for_key(user_api_key_auth) + team_scope = await AgentRequestHandler._resolve_agents_for_team(user_api_key_auth) + return _combine_key_and_team_scope(key_scope, team_scope) @staticmethod async def is_agent_allowed( @@ -77,13 +108,17 @@ class AgentRequestHandler: 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 + scope = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth) + match scope: + case UnrestrictedAgents(): + return True + case RestrictedAgents(agent_ids): + return agent_id in agent_ids + case UnresolvableAgents(reason): + verbose_logger.warning(f"Denying agent access, permissions unreadable: {reason}") + return False + case _: + assert_never(scope) @staticmethod def _get_key_object_permission( @@ -135,70 +170,64 @@ class AgentRequestHandler: return team_obj.object_permission @staticmethod - async def _get_allowed_agents_for_key( + async def _resolve_agents_for_key( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentScope: """ - Get allowed agents for a key. + Resolve the key level. 1. First checks native key-level agent permissions (object_permission) 2. Also includes agents from key's access_group_ids (unified access groups) Note: object_permission is already loaded by get_key_object() in main auth flow. + A key that names agents or access groups is restricted even when those names + resolve to nothing, so a group that is empty (or whose agents were deleted) + cannot hand the key every agent. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgents() try: - all_agents: list[str] = [] - - # 1. Get agents from object_permission (native permissions) key_object_permission = AgentRequestHandler._get_key_object_permission(user_api_key_auth) - if key_object_permission is not None: - # Get direct agents - direct_agents = key_object_permission.agents or [] + direct_agents = tuple(key_object_permission.agents or ()) if key_object_permission else () + agent_access_groups = ( + tuple(key_object_permission.agent_access_groups or ()) if key_object_permission else () + ) + key_access_group_ids = tuple(user_api_key_auth.access_group_ids or ()) - # Get agents from access groups - access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] - ) + if not (direct_agents or agent_access_groups or key_access_group_ids): + return UnrestrictedAgents() - all_agents = direct_agents + access_group_agents + access_group_agents = await AgentRequestHandler._get_agents_from_access_groups(list(agent_access_groups)) - # 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 [] + unified_agents: tuple[str, ...] = () if key_access_group_ids: 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, + unified_agents = tuple( + await _get_agent_ids_from_access_groups(access_group_ids=list(key_access_group_ids)) ) - all_agents.extend(unified_agents) - return list(set(all_agents)) + return RestrictedAgents(frozenset(direct_agents + tuple(access_group_agents) + unified_agents)) except Exception as e: - verbose_logger.warning(f"Failed to get allowed agents for key: {e!s}") - return [] + return UnresolvableAgents(f"key agent permissions: {e!s}") @staticmethod - async def _get_allowed_agents_for_team( + async def _resolve_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentScope: """ - Get allowed agents for a team. + Resolve the team level. 1. First checks native team-level agent permissions (object_permission) 2. Also includes agents from team's access_group_ids (unified access groups) Fetches the team object once and reuses it for both permission sources. """ - if user_api_key_auth is None: - return [] - - if user_api_key_auth.team_id is None: - return [] + if user_api_key_auth is None or user_api_key_auth.team_id is None: + return UnrestrictedAgents() try: from litellm.proxy.auth.auth_checks import get_team_object @@ -209,7 +238,7 @@ class AgentRequestHandler: ) if not prisma_client: - return [] + return UnrestrictedAgents() # Fetch the team object once for both permission sources team_obj = await get_team_object( @@ -221,42 +250,35 @@ class AgentRequestHandler: ) if team_obj is None: - return [] + return UnrestrictedAgents() - all_agents: list[str] = [] - - # 1. Get agents from object_permission (native permissions) object_permissions = team_obj.object_permission - if object_permissions is not None: - # Get direct agents - direct_agents = object_permissions.agents or [] + direct_agents = tuple(object_permissions.agents or ()) if object_permissions else () + agent_access_groups = tuple(object_permissions.agent_access_groups or ()) if object_permissions else () + team_access_group_ids = tuple(team_obj.access_group_ids or ()) - # Get agents from access groups - access_group_agents = await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] - ) + if not (direct_agents or agent_access_groups or team_access_group_ids): + return UnrestrictedAgents() - all_agents = direct_agents + access_group_agents + access_group_agents = await AgentRequestHandler._get_agents_from_access_groups(list(agent_access_groups)) - # 2. Also include agents from team's access_group_ids (unified access groups) - team_access_group_ids = team_obj.access_group_ids or [] + unified_agents: tuple[str, ...] = () if team_access_group_ids: 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, + unified_agents = tuple( + await _get_agent_ids_from_access_groups(access_group_ids=list(team_access_group_ids)) ) - all_agents.extend(unified_agents) - return list(set(all_agents)) + return RestrictedAgents(frozenset(direct_agents + tuple(access_group_agents) + unified_agents)) except Exception as e: # litellm-dashboard is the default UI team and will never have agents; # skip noisy warnings for it. if user_api_key_auth.team_id != UI_TEAM_ID: verbose_logger.warning(f"Failed to get allowed agents for team: {e!s}") - return [] + return UnresolvableAgents(f"team agent permissions: {e!s}") @staticmethod def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: @@ -275,18 +297,16 @@ class AgentRequestHandler: async def _get_db_agent_ids_for_access_groups(prisma_client, access_groups: list[str]) -> set[str]: """ Helper to get agent_ids from DB agents that match any of the given access groups. + + A failed query raises: an unreadable grant is not an empty grant, and swallowing + it here would report a restricted caller as carrying no restriction at all. """ - agent_ids: set[str] = set() - if access_groups and prisma_client is not None: - try: - agents = await AgentsRepository(prisma_client).table.find_many( - where={"agent_access_groups": {"hasSome": access_groups}} - ) - 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}") - return agent_ids + if not access_groups or prisma_client is None: + return set() + agents = await AgentsRepository(prisma_client).table.find_many( + where={"agent_access_groups": {"hasSome": access_groups}} + ) + return {agent.agent_id for agent in agents} @staticmethod async def _get_agents_from_access_groups( @@ -298,20 +318,15 @@ class AgentRequestHandler: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import prisma_client - try: - # Use the helper for config-loaded agents - agent_ids = AgentRequestHandler._get_config_agent_ids_for_access_groups( - global_agent_registry.agent_list, access_groups - ) + # Use the helper for config-loaded agents + agent_ids = AgentRequestHandler._get_config_agent_ids_for_access_groups( + global_agent_registry.agent_list, access_groups + ) - # Use the helper for DB agents - db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups) - agent_ids.update(db_agent_ids) + # Use the helper for DB agents + db_agent_ids = await AgentRequestHandler._get_db_agent_ids_for_access_groups(prisma_client, access_groups) - return list(agent_ids) - except Exception as e: - verbose_logger.warning(f"Failed to get agents from access groups: {e!s}") - return [] + return list(agent_ids | db_agent_ids) @staticmethod async def get_agent_access_groups( diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 1efbdeb0132..77b9b8d26df 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -12,7 +12,7 @@ import asyncio import os import uuid from collections.abc import Mapping, Sequence -from typing import TypedDict +from typing import TYPE_CHECKING, TypedDict from fastapi import APIRouter, Depends, HTTPException, Query, Request from typing_extensions import Required @@ -46,6 +46,9 @@ from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) +if TYPE_CHECKING: + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import AgentScope + def _proxy_base_url(http_request: Request) -> str: """Return the proxy's public base URL, preferring PROXY_BASE_URL when set.""" @@ -197,6 +200,30 @@ async def _check_agent_url_health( } +def _agents_visible_to(scope: "AgentScope") -> list[AgentResponse]: + """ + Narrow the registry to the agents a caller may see, denying when its grants are unreadable. + """ + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgents, + UnresolvableAgents, + UnrestrictedAgents, + ) + + all_agents = global_agent_registry.get_agent_list() + match scope: + case UnrestrictedAgents(): + return all_agents + case RestrictedAgents(allowed_agent_ids): + return [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] + case UnresolvableAgents(reason): + raise HTTPException( + status_code=503, + detail=f"Agent permissions are currently unreadable, refusing to list agents: {reason}", + ) + + @router.get( "/v1/agents", tags=["[beta] A2A Agents"], @@ -246,16 +273,9 @@ async def get_agents( ): returned_agents = global_agent_registry.get_agent_list() else: - # Get allowed agents from object_permission (key/team level) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - # If no restrictions (empty list), return all agents - if len(allowed_agent_ids) == 0: - returned_agents = global_agent_registry.get_agent_list() - else: - # Filter agents by allowed IDs - all_agents = global_agent_registry.get_agent_list() - returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] + returned_agents = _agents_visible_to( + await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=user_api_key_dict) + ) # Fetch current spend from DB for all returned agents from litellm.proxy.proxy_server import prisma_client @@ -1044,16 +1064,17 @@ async def get_agent_daily_activity( # intersect their explicit `agent_ids` filter with the same allowlist. from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgents, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view where_condition: dict[str, object] = {} if not _user_has_admin_view(user_api_key_dict): - permitted_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - # `get_allowed_agents` returns an empty list when the caller's key - # and team carry no agent restrictions. For activity scoping that's - # not "see everything" — fall back to the agents the caller - # created so they cannot enumerate other tenants' agents. + scope = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=user_api_key_dict) + permitted_agent_ids = list(scope.agent_ids) if isinstance(scope, RestrictedAgents) else [] + # An unrestricted caller is not "see everything" for activity scoping; + # fall back to the agents the caller created so they cannot + # enumerate other tenants' agents. # Guard against `user_id is None`: a literal None in Prisma # `where={"created_by": None}` resolves to ``created_by IS NULL`` # and would expose every ownerless agent's rows. diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 56053f59c85..12950f98dcf 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -25,9 +25,11 @@ async def append_agents_to_model_group( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgents, ) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) + scope = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=user_api_key_dict) + allowed_agent_ids = scope.agent_ids if isinstance(scope, RestrictedAgents) else frozenset() for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) @@ -59,9 +61,11 @@ async def append_agents_to_model_info( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgents, ) - allowed_agent_ids = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) + scope = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=user_api_key_dict) + allowed_agent_ids = scope.agent_ids if isinstance(scope, RestrictedAgents) else frozenset() for agent_id in allowed_agent_ids: agent = global_agent_registry.get_agent_by_id(agent_id) diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py index 78cb1488533..a8f46126542 100644 --- a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -11,6 +11,9 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + UnrestrictedAgents, +) def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: @@ -64,8 +67,8 @@ async def test_get_agents_allowed_when_not_disabled(): MagicMock(get_agent_list=MagicMock(return_value=[])), ): with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + new=AsyncMock(return_value=UnrestrictedAgents()), ): result = await get_agents(request=request_mock, user_api_key_dict=user) assert result == [] 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 6bdea9c2615..5b18f07f17d 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 @@ -10,9 +10,12 @@ import pytest sys.path.insert(0, os.path.abspath("../../../..")) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgents, + UnresolvableAgents, + UnrestrictedAgents, ) @@ -22,11 +25,11 @@ class TestAgentRequestHandler: Test suite for AgentRequestHandler permission logic. """ - async def test_get_allowed_agents_intersection_logic(self): + async def test_resolve_agent_scope_intersection_logic(self): """ Test key/team intersection: when both have restrictions, only common agents are allowed. When team has restrictions but key has none, key inherits from team. - When neither has restrictions, returns empty list (meaning allow all). + When neither has restrictions, the caller is unrestricted. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -34,91 +37,83 @@ class TestAgentRequestHandler: team_id="test-team", ) - # Case 1: Both key and team have agents - intersection - 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 = ["agent1", "agent2", "agent3"] - mock_team.return_value = ["agent2", "agent4"] + with patch.object(AgentRequestHandler, "_resolve_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_resolve_agents_for_team") as mock_team: + mock_key.return_value = RestrictedAgents(frozenset({"agent1", "agent2", "agent3"})) + mock_team.return_value = RestrictedAgents(frozenset({"agent2", "agent4"})) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgents(frozenset({"agent2"})) + + with patch.object(AgentRequestHandler, "_resolve_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_resolve_agents_for_team") as mock_team: + mock_key.return_value = UnrestrictedAgents() + mock_team.return_value = RestrictedAgents(frozenset({"team_agent1", "team_agent2"})) + + result = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgents(frozenset({"team_agent1", "team_agent2"})) + + with patch.object(AgentRequestHandler, "_resolve_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_resolve_agents_for_team") as mock_team: + mock_key.return_value = UnrestrictedAgents() + mock_team.return_value = UnrestrictedAgents() + + result = await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=mock_user_auth) + assert result == UnrestrictedAgents() + + async def test_disjoint_key_and_team_permissions_deny_every_agent(self): + """ + Regression: a key restricted to agent1 inside a team restricted to agent2 shares no + agent with its team, so it must reach neither of them (and certainly not a third + agent nobody granted it). The intersection used to collapse to an empty list, which + the caller read as "no restrictions" and turned into access to every agent. + """ + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + ) + + with patch.object(AgentRequestHandler, "_resolve_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_resolve_agents_for_team") as mock_team: + mock_key.return_value = RestrictedAgents(frozenset({"agent1"})) + mock_team.return_value = RestrictedAgents(frozenset({"agent2"})) + + assert await AgentRequestHandler.resolve_agent_scope( user_api_key_auth=mock_user_auth - ) - assert sorted(result) == ["agent2"] + ) == RestrictedAgents(frozenset()) - # Case 2: Team has agents, key has none - inherit from team - 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 = [] - mock_team.return_value = ["team_agent1", "team_agent2"] - - result = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=mock_user_auth - ) - assert sorted(result) == ["team_agent1", "team_agent2"] - - # Case 3: No restrictions - returns empty list (allow all) - 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 = [] - mock_team.return_value = [] - - result = await AgentRequestHandler.get_allowed_agents( - user_api_key_auth=mock_user_auth - ) - assert result == [] + for agent_id in ("agent1", "agent2", "agent-nobody-granted"): + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent_id, user_api_key_auth=mock_user_auth + ) + is False + ) async def test_is_agent_allowed_respects_permissions(self): """ - Test is_agent_allowed: returns True if agent in allowed list or if no restrictions. - Returns False if agent not in allowed list. + Test is_agent_allowed: returns True if agent in allowed set or if no restrictions. + Returns False if agent not in allowed set. """ mock_user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user") - # Agent in allowed list - should be allowed - with patch.object( - AgentRequestHandler, "get_allowed_agents" - ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + with patch.object(AgentRequestHandler, "resolve_agent_scope") as mock_scope: + mock_scope.return_value = RestrictedAgents(frozenset({"agent1", "agent2"})) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="agent1", user_api_key_auth=mock_user_auth - ) - is True + await AgentRequestHandler.is_agent_allowed(agent_id="agent1", user_api_key_auth=mock_user_auth) is True ) - # Agent not in allowed list - should be denied - with patch.object( - AgentRequestHandler, "get_allowed_agents" - ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + with patch.object(AgentRequestHandler, "resolve_agent_scope") as mock_scope: + mock_scope.return_value = RestrictedAgents(frozenset({"agent1", "agent2"})) assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="agent3", user_api_key_auth=mock_user_auth - ) - is False + await AgentRequestHandler.is_agent_allowed(agent_id="agent3", user_api_key_auth=mock_user_auth) is False ) - # Empty list means no restrictions - should allow any agent - with patch.object( - AgentRequestHandler, "get_allowed_agents" - ) as mock_get_allowed: - mock_get_allowed.return_value = [] + with patch.object(AgentRequestHandler, "resolve_agent_scope") as mock_scope: + mock_scope.return_value = UnrestrictedAgents() assert ( - await AgentRequestHandler.is_agent_allowed( - agent_id="any_agent", user_api_key_auth=mock_user_auth - ) + await AgentRequestHandler.is_agent_allowed(agent_id="any_agent", user_api_key_auth=mock_user_auth) is True ) @@ -126,17 +121,15 @@ class TestAgentRequestHandler: """ Test that when user_api_key_auth is None, all agents are allowed (no restrictions). """ - result = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=None) - assert result == [] + assert await AgentRequestHandler.resolve_agent_scope(user_api_key_auth=None) == UnrestrictedAgents() - is_allowed = await AgentRequestHandler.is_agent_allowed( - agent_id="any_agent", user_api_key_auth=None - ) - assert is_allowed is True + assert await AgentRequestHandler.is_agent_allowed(agent_id="any_agent", user_api_key_auth=None) is True - async def test_get_allowed_agents_handles_errors_gracefully(self): + async def test_unreadable_permissions_deny_instead_of_allowing_everything(self): """ - Test that errors during permission lookup are handled gracefully (returns empty list). + Regression: a lookup that raises (a database outage, say) leaves the caller's + restriction unknown. That used to surface as an empty list, i.e. "unrestricted", + so an outage handed every restricted key every agent. It must deny instead. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -145,23 +138,65 @@ class TestAgentRequestHandler: object_permission_id="test-permission", ) - 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.side_effect = Exception("DB Error") - mock_team.return_value = [] + with patch.object(AgentRequestHandler, "_resolve_agents_for_key") as mock_key: + with patch.object(AgentRequestHandler, "_resolve_agents_for_team") as mock_team: + mock_key.return_value = UnresolvableAgents("DB Error") + mock_team.return_value = UnrestrictedAgents() - result = await AgentRequestHandler.get_allowed_agents( + assert await AgentRequestHandler.resolve_agent_scope( user_api_key_auth=mock_user_auth - ) - assert result == [] + ) == UnresolvableAgents("DB Error") - async def test_get_allowed_agents_for_key_via_access_group_ids(self): + assert ( + await AgentRequestHandler.is_agent_allowed(agent_id="agent1", user_api_key_auth=mock_user_auth) + is False + ) + + async def test_key_access_group_lookup_failure_is_unresolvable(self): """ - Test that _get_allowed_agents_for_key includes agents from key's access_group_ids + Regression: the access group query used to swallow its own exception and return no + agents, which made a restricted key look unrestricted. The failure has to reach the + caller as UnresolvableAgents so the request is denied. + """ + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + access_group_ids=["ag-1"], + ) + + with patch( + "litellm.proxy.auth.auth_checks._get_agent_ids_from_access_groups", + new_callable=AsyncMock, + side_effect=Exception("connection refused"), + ): + scope = await AgentRequestHandler._resolve_agents_for_key(user_api_key_auth=mock_user_auth) + + assert isinstance(scope, UnresolvableAgents) + assert "connection refused" in scope.reason + + async def test_access_group_that_resolves_to_nothing_still_restricts(self): + """ + Regression: a key whose only grant is an access group with no agents left is + restricted to nothing, not unrestricted. + """ + mock_user_auth = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + access_group_ids=["ag-empty"], + ) + + with patch( + "litellm.proxy.auth.auth_checks._get_agent_ids_from_access_groups", + new_callable=AsyncMock, + return_value=[], + ): + scope = await AgentRequestHandler._resolve_agents_for_key(user_api_key_auth=mock_user_auth) + + assert scope == RestrictedAgents(frozenset()) + + async def test_resolve_agents_for_key_via_access_group_ids(self): + """ + Test that _resolve_agents_for_key includes agents from key's access_group_ids (unified access groups) when key has no native object_permission. """ mock_user_auth = UserAPIKeyAuth( @@ -170,26 +205,20 @@ class TestAgentRequestHandler: access_group_ids=["ag-with-agents"], ) - with patch.object( - AgentRequestHandler, "_get_key_object_permission", return_value=None - ): + with patch.object(AgentRequestHandler, "_get_key_object_permission", return_value=None): with patch( "litellm.proxy.auth.auth_checks._get_agent_ids_from_access_groups", new_callable=AsyncMock, return_value=["agent-from-ag-1", "agent-from-ag-2"], ): - result = await AgentRequestHandler._get_allowed_agents_for_key( - user_api_key_auth=mock_user_auth - ) - assert sorted(result) == ["agent-from-ag-1", "agent-from-ag-2"] + result = await AgentRequestHandler._resolve_agents_for_key(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgents(frozenset({"agent-from-ag-1", "agent-from-ag-2"})) - async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self): + async def test_resolve_agents_for_key_combines_native_and_access_groups(self): """ - Test that _get_allowed_agents_for_key combines agents from native object_permission + Test that _resolve_agents_for_key combines agents from native object_permission and key's access_group_ids (unified access groups). """ - from litellm.proxy._types import LiteLLM_ObjectPermissionTable - mock_permission = LiteLLM_ObjectPermissionTable( object_permission_id="obj-1", agents=["native-agent-1"], @@ -200,7 +229,6 @@ class TestAgentRequestHandler: user_id="test-user", access_group_ids=["ag-1"], ) - # Attach object_permission so _get_key_object_permission returns it mock_user_auth.object_permission = mock_permission with patch( @@ -208,7 +236,5 @@ class TestAgentRequestHandler: new_callable=AsyncMock, return_value=["agent-from-ag"], ): - result = await AgentRequestHandler._get_allowed_agents_for_key( - user_api_key_auth=mock_user_auth - ) - assert sorted(result) == ["agent-from-ag", "native-agent-1"] + result = await AgentRequestHandler._resolve_agents_for_key(user_api_key_auth=mock_user_auth) + assert result == RestrictedAgents(frozenset({"agent-from-ag", "native-agent-1"})) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py index 86eb7a83079..37047eb5045 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py @@ -19,6 +19,9 @@ from litellm.proxy.agent_endpoints.model_list_helpers import ( ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.agents import AgentResponse +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgents, +) from litellm.types.proxy.management_endpoints.model_management_endpoints import ( ModelGroupInfoProxy, ) @@ -37,15 +40,15 @@ async def test_append_agents_to_model_group(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["test-agent-id"]) + mock_resolve_agent_scope = AsyncMock(return_value=RestrictedAgents(frozenset({"test-agent-id"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - mock_get_allowed_agents, + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + mock_resolve_agent_scope, ): with patch( "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", @@ -80,15 +83,15 @@ async def test_append_agents_to_model_info(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["agent-123"]) + mock_resolve_agent_scope = AsyncMock(return_value=RestrictedAgents(frozenset({"agent-123"}))) # Mock global_agent_registry mock_registry = Mock() mock_registry.get_agent_by_id = Mock(return_value=mock_agent) with patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - mock_get_allowed_agents, + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + mock_resolve_agent_scope, ): with patch( "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", diff --git a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py index 0855c201945..9c8ea1749d6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py +++ b/tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py @@ -14,6 +14,11 @@ import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgents, + UnrestrictedAgents, +) + # --------------------------------------------------------------------------- # /team/daily/activity — per-team admin/permission requirement @@ -227,8 +232,8 @@ async def test_agent_activity_non_admin_no_perms_falls_back_to_owned(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), # no explicit agent permissions + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + new=AsyncMock(return_value=UnrestrictedAgents()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -275,8 +280,8 @@ async def test_agent_activity_non_admin_intersects_explicit_agent_ids(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=["agent-permitted"]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + new=AsyncMock(return_value=RestrictedAgents(frozenset({"agent-permitted"}))), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -321,8 +326,8 @@ async def test_agent_activity_keyless_caller_does_not_query_created_by_null(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + new=AsyncMock(return_value=UnrestrictedAgents()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -366,8 +371,8 @@ async def test_agent_activity_non_admin_no_access_returns_empty_page(): new=AsyncMock(return_value=None), ), patch( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", - new=AsyncMock(return_value=[]), + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_scope", + new=AsyncMock(return_value=UnrestrictedAgents()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity",