From dd6a6e2fdfdc80d83ed4fd5023eefba6599d180d Mon Sep 17 00:00:00 2001 From: milan Date: Mon, 3 Aug 2026 22:12:12 +0000 Subject: [PATCH] fix(agents): stop treating an empty agent permission set as unrestricted A key restricted to one agent inside a team restricted to another intersected to an empty list, and every caller read an empty list as 'no restrictions set', so adding the second restriction handed the key every agent on the proxy. The same empty list came out of the except blocks around the permission lookups, so a database outage widened access the same way. Resolution now returns a tagged union of unrestricted, restricted to a concrete set, or unresolvable, and the callers match on it: a restricted caller with no agents left reaches none, and a grant that cannot be read denies rather than allows. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_permission_handler.py | 229 ++++++++-------- litellm/proxy/agent_endpoints/endpoints.py | 53 ++-- .../agent_endpoints/model_list_helpers.py | 8 +- .../proxy/agent_endpoints/test_agent_rbac.py | 7 +- .../auth/test_agent_permission_handler.py | 244 ++++++++++-------- .../test_model_list_helpers.py | 15 +- .../test_activity_tenant_scoping.py | 21 +- 7 files changed, 327 insertions(+), 250 deletions(-) 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",