diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 632743236c4..c8765eb0bd0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -84,7 +84,7 @@ "limit": 56 }, "reportPrivateUsage": { - "limit": 1825 + "limit": 1824 }, "reportRedeclaration": { "limit": 8 diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 6baeab3dfff..81586b4eef3 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -5,7 +5,8 @@ Handles agent permission checking for keys and teams using object_permission_id. Follows the same pattern as MCP permission handling. """ -from typing import Final +from dataclasses import dataclass +from typing import Final, TypeAlias from litellm._logging import verbose_logger from litellm.proxy._types import ( @@ -17,6 +18,27 @@ from litellm.proxy._types import ( from litellm.repositories.table_repositories import AgentsRepository +@dataclass(frozen=True, slots=True) +class UnrestrictedAgentAccess: + """No agent grant exists on the key or its team, so every agent is reachable.""" + + +@dataclass(frozen=True, slots=True) +class RestrictedAgentAccess: + """Only ``agent_ids`` are reachable. An empty set denies every agent.""" + + agent_ids: frozenset[str] + + +AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess + + +def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -32,38 +54,32 @@ class AgentRequestHandler: """ @staticmethod - async def get_allowed_agents( + async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ - Get list of allowed agent IDs for the given user/key based on permissions. + Resolve the agents the given user/key may reach. - Returns: - List[str]: List of allowed agent IDs. Empty list means no restrictions (allow all). + ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team + carries any grant. Grants that intersect to nothing stay restricted, so + narrowing a caller can never widen what it reaches. """ try: - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) + team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - raw_key_grants: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) - raw_team_grants: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - allowed_agents_for_key: Final = frozenset( - global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_key_grants - ) - allowed_agents_for_team: Final = frozenset( - global_agent_registry.stable_agent_id(agent_id) for agent_id in raw_team_grants - ) - - # If team has agent restrictions, handle inheritance and intersection logic - if allowed_agents_for_team and allowed_agents_for_key: - # Key has its own agent permissions - use intersection with team permissions - return sorted(allowed_agents_for_key & allowed_agents_for_team) - if allowed_agents_for_team: - # Key has no agent permissions - inherit from team - return sorted(allowed_agents_for_team) - return sorted(allowed_agents_for_key) + match (key_access, team_access): + case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()): + return UnrestrictedAgentAccess() + case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(team_ids)) + case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()): + return RestrictedAgentAccess(_to_stable_ids(key_ids)) + case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)): + return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids)) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def is_agent_allowed( @@ -82,14 +98,12 @@ class AgentRequestHandler: """ from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - allowed_agents: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth) - - # Empty list means no restrictions - allow all - if len(allowed_agents) == 0: - return True - - stable_id: Final = global_agent_registry.stable_agent_id(agent_id) - return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agents) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth): + case UnrestrictedAgentAccess(): + return True + case RestrictedAgentAccess(allowed_agent_ids): + stable_id: Final = global_agent_registry.stable_agent_id(agent_id) + return not global_agent_registry.ids_for_agent(stable_id).isdisjoint(allowed_agent_ids) @staticmethod def _get_key_object_permission( @@ -143,55 +157,58 @@ class AgentRequestHandler: @staticmethod async def _get_allowed_agents_for_key( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a key. 1. First checks native key-level agent permissions (object_permission) 2. Also includes agents from key's access_group_ids (unified access groups) + A key that declares agents or access groups is restricted even when those + declarations resolve to nothing, so an emptied or deleted access group denies + rather than opening the key up. Lookup failures still propagate to the caller, + which keeps them fail-open. + Note: object_permission is already loaded by get_key_object() in main auth flow. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() try: - all_agents: list[str] = [] - # 1. Get agents from object_permission (native permissions) key_object_permission: Final = AgentRequestHandler._get_key_object_permission(user_api_key_auth) - if key_object_permission is not None: - # Get direct agents - direct_agents: Final = key_object_permission.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - key_object_permission.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple( + key_object_permission.agents or () if key_object_permission is not None else () + ) + declared_access_groups: Final = tuple( + key_object_permission.agent_access_groups or () if key_object_permission is not None else () + ) # 2. Fallback: get agent IDs from key's access_group_ids (unified access groups) - key_access_group_ids: Final = 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, - ) + key_access_group_ids: Final = tuple(user_api_key_auth.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=key_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not key_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(key_access_group_ids))) + if key_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + access_group_agents + unified_agents)) except Exception as e: verbose_logger.warning("Failed to get allowed agents for key: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod async def _get_allowed_agents_for_team( user_api_key_auth: UserAPIKeyAuth | None = None, - ) -> list[str]: + ) -> AgentAccess: """ Get allowed agents for a team. @@ -199,12 +216,13 @@ class AgentRequestHandler: 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. + Declared-but-empty grants stay restricted; see `_get_allowed_agents_for_key`. """ if user_api_key_auth is None: - return [] + return UnrestrictedAgentAccess() if user_api_key_auth.team_id is None: - return [] + return UnrestrictedAgentAccess() try: from litellm.proxy.auth.auth_checks import get_team_object @@ -215,7 +233,7 @@ class AgentRequestHandler: ) if not prisma_client: - return [] + return UnrestrictedAgentAccess() # Fetch the team object once for both permission sources team_obj: Final = await get_team_object( @@ -227,42 +245,38 @@ class AgentRequestHandler: ) if team_obj is None: - return [] - - all_agents: list[str] = [] + return UnrestrictedAgentAccess() # 1. Get agents from object_permission (native permissions) object_permissions: Final = team_obj.object_permission - if object_permissions is not None: - # Get direct agents - direct_agents: Final = object_permissions.agents or [] - - # Get agents from access groups - access_group_agents: Final = await AgentRequestHandler._get_agents_from_access_groups( - object_permissions.agent_access_groups or [] - ) - - all_agents = direct_agents + access_group_agents - + direct_agents: Final = tuple(object_permissions.agents or () if object_permissions is not None else ()) + declared_access_groups: Final = tuple( + object_permissions.agent_access_groups or () if object_permissions is not None else () + ) # 2. Also include agents from team's access_group_ids (unified access groups) - team_access_group_ids: Final = team_obj.access_group_ids or [] - if team_access_group_ids: - from litellm.proxy.auth.auth_checks import ( - _get_agent_ids_from_access_groups, - ) + team_access_group_ids: Final = tuple(team_obj.access_group_ids or ()) - unified_agents: Final = await _get_agent_ids_from_access_groups( - access_group_ids=team_access_group_ids, - ) - all_agents.extend(unified_agents) + if not direct_agents and not declared_access_groups and not team_access_group_ids: + return UnrestrictedAgentAccess() - return list(set(all_agents)) + access_group_agents: Final = ( + tuple(await AgentRequestHandler._get_agents_from_access_groups(list(declared_access_groups))) + if declared_access_groups + else () + ) + unified_agents: Final = ( + tuple(await AgentRequestHandler._get_unified_access_group_agents(list(team_access_group_ids))) + if team_access_group_ids + else () + ) + + return RestrictedAgentAccess(frozenset(direct_agents + 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("Failed to get allowed agents for team: %s", e) - return [] + return UnrestrictedAgentAccess() @staticmethod def _get_config_agent_ids_for_access_groups(config_agents: list, access_groups: list[str]) -> set[str]: @@ -281,18 +295,26 @@ 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. + + Query failures propagate so the caller can tell "this group is empty" (deny) + apart from "the lookup failed" (fail-open). """ - agent_ids: Final[set[str]] = set() - if access_groups and prisma_client is not None: - try: - agents: Final = 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("Error getting agents from access groups: %s", e) - return agent_ids + if not access_groups or prisma_client is None: + return set() + + agents: Final = 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_unified_access_group_agents(access_group_ids: list[str]) -> list[str]: + """ + Resolve unified access group ids to agent IDs. + """ + from litellm.proxy.auth.auth_checks import _get_agent_ids_from_access_groups + + return await _get_agent_ids_from_access_groups(access_group_ids=access_group_ids) @staticmethod async def _get_agents_from_access_groups( @@ -304,20 +326,17 @@ 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: Final = AgentRequestHandler._get_config_agent_ids_for_access_groups( - global_agent_registry.agent_list, access_groups - ) + # Use the helper for config-loaded agents + config_agent_ids: Final = 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: Final = await AgentRequestHandler._get_db_agent_ids_for_access_groups( + prisma_client, access_groups + ) - return list(agent_ids) - except Exception as e: - verbose_logger.warning("Failed to get agents from access groups: %s", e) - return [] + return list(config_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 028885aa065..5bc5b5566f8 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -248,6 +248,8 @@ async def get_agents( from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) try: @@ -261,15 +263,14 @@ 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: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) + agent_access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict) + all_agents: Final = global_agent_registry.get_agent_list() - # 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: Final = global_agent_registry.get_agent_list() - returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] + match agent_access: + case UnrestrictedAgentAccess(): + returned_agents = all_agents + case RestrictedAgentAccess(allowed_agent_ids): + returned_agents = [agent for agent in all_agents if agent.agent_id in allowed_agent_ids] # Fetch current spend from DB for all returned agents from litellm.proxy.proxy_server import prisma_client @@ -1061,27 +1062,29 @@ 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, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view where_condition: Final[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. + permitted_agent_ids: list[str] = [] + # 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. - if not permitted_agent_ids: - if user_api_key_dict.user_id is None: - permitted_agent_ids = [] - else: - owned_records: Final = await agents_table(prisma_client).find_many( - where={"created_by": user_api_key_dict.user_id} - ) - permitted_agent_ids = [a.agent_id for a in owned_records] + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + permitted_agent_ids = list(allowed_agent_ids) + case UnrestrictedAgentAccess(): + if user_api_key_dict.user_id is not None: + owned_records: Final = await agents_table(prisma_client).find_many( + where={"created_by": user_api_key_dict.user_id} + ) + permitted_agent_ids = [a.agent_id for a in owned_records] if agent_ids_list: permitted_agent_id_set: Final = set(permitted_agent_ids) diff --git a/litellm/proxy/agent_endpoints/model_list_helpers.py b/litellm/proxy/agent_endpoints/model_list_helpers.py index 862b66d3860..4a88644bae5 100644 --- a/litellm/proxy/agent_endpoints/model_list_helpers.py +++ b/litellm/proxy/agent_endpoints/model_list_helpers.py @@ -4,8 +4,6 @@ Helper functions for appending A2A agents to model lists. Used by proxy model endpoints to make agents appear in UI alongside models. """ -from typing import Final - from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -27,20 +25,23 @@ 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, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - model_groups.append( - ModelGroupInfoProxy( - model_group=f"a2a/{agent.agent_name}", - mode="chat", - providers=["a2a"], - ) - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + model_groups.append( + ModelGroupInfoProxy( + model_group=f"a2a/{agent.agent_name}", + mode="chat", + providers=["a2a"], + ) + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to model_group/info: %s", e) @@ -61,30 +62,33 @@ 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, + RestrictedAgentAccess, ) - allowed_agent_ids: Final = await AgentRequestHandler.get_allowed_agents(user_api_key_auth=user_api_key_dict) - - for agent_id in allowed_agent_ids: - agent = global_agent_registry.get_agent_by_id(agent_id) - if agent is not None: - models.append( - { - "model_name": f"a2a/{agent.agent_name}", - "litellm_params": { - "model": f"a2a/{agent.agent_name}", - "custom_llm_provider": "a2a", - }, - "model_info": { - "id": agent.agent_id, - "mode": "chat", - "db_model": True, - "created_by": agent.created_by, - "created_at": agent.created_at, - "updated_at": agent.updated_at, - }, - } - ) + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict): + case RestrictedAgentAccess(allowed_agent_ids): + for agent_id in allowed_agent_ids: + agent = global_agent_registry.get_agent_by_id(agent_id) + if agent is not None: + models.append( + { + "model_name": f"a2a/{agent.agent_name}", + "litellm_params": { + "model": f"a2a/{agent.agent_name}", + "custom_llm_provider": "a2a", + }, + "model_info": { + "id": agent.agent_id, + "mode": "chat", + "db_model": True, + "created_by": agent.created_by, + "created_at": agent.created_at, + "updated_at": agent.updated_at, + }, + } + ) + case _: + pass except Exception as e: verbose_proxy_logger.debug("Error appending agents to v2/model/info: %s", e) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index afd8c0107ea..8ff4bfb36c0 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -9,7 +9,7 @@ "limit": 834 }, "ANN201": { - "limit": 2033 + "limit": 2032 }, "ANN202": { "limit": 865 @@ -60,7 +60,7 @@ "limit": 0 }, "BLE001": { - "limit": 2926 + "limit": 2924 }, "C401": { "limit": 8 @@ -135,7 +135,7 @@ "limit": 27 }, "PERF401": { - "limit": 13 + "limit": 12 }, "PERF402": { "limit": 0 diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py index 78cb1488533..842a635ecad 100644 --- a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -11,6 +11,10 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: @@ -64,8 +68,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_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ): 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 c57296dfc3f..066d33dc187 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 @@ -17,6 +17,8 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, + RestrictedAgentAccess, + UnrestrictedAgentAccess, ) @@ -26,11 +28,11 @@ class TestAgentRequestHandler: Test suite for AgentRequestHandler permission logic. """ - async def test_get_allowed_agents_intersection_logic(self): + async def test_resolve_agent_access_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). + Only a caller with no grant anywhere is unrestricted. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -45,13 +47,17 @@ class TestAgentRequestHandler: 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"] + mock_key.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2", "agent3"}) + ) + mock_team.return_value = RestrictedAgentAccess( + frozenset({"agent2", "agent4"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent2"] + assert result == RestrictedAgentAccess(frozenset({"agent2"})) # Case 2: Team has agents, key has none - inherit from team with patch.object( @@ -60,41 +66,124 @@ class TestAgentRequestHandler: with patch.object( AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: - mock_key.return_value = [] - mock_team.return_value = ["team_agent1", "team_agent2"] + mock_key.return_value = UnrestrictedAgentAccess() + mock_team.return_value = RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["team_agent1", "team_agent2"] + assert result == RestrictedAgentAccess( + frozenset({"team_agent1", "team_agent2"}) + ) - # Case 3: No restrictions - returns empty list (allow all) + # Case 3: Key has agents, team has none - key restrictions stand 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 = [] + mock_key.return_value = RestrictedAgentAccess(frozenset({"key_agent1"})) + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == RestrictedAgentAccess(frozenset({"key_agent1"})) + + # Case 4: No grant anywhere - unrestricted (documented open-by-default) + 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 = UnrestrictedAgentAccess() + mock_team.return_value = UnrestrictedAgentAccess() + + result = await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) + assert result == UnrestrictedAgentAccess() + + async def test_disjoint_key_and_team_grants_deny_every_agent(self): + """LIT-5143: a key restricted to one agent inside a team restricted to another + must reach nothing. The empty intersection used to read as "no restrictions", + so adding the team grant handed the key every agent on the proxy.""" + mock_user_auth: Final = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-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 = RestrictedAgentAccess(frozenset({"agent-alpha"})) + mock_team.return_value = RestrictedAgentAccess(frozenset({"agent-beta"})) + + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + for agent_id in ("agent-alpha", "agent-beta", "agent-secret"): + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id=agent_id, user_api_key_auth=mock_user_auth + ) + is False + ), agent_id + + async def test_empty_access_group_denies_every_agent(self): + """LIT-5143: a key restricted to an access group that resolves to no agents is + restricted to nothing, not unrestricted. A failed group lookup still fails open.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionTable + + mock_user_auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + mock_user_auth.object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="obj-1", + agents=[], + agent_access_groups=["group-with-no-agents"], + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.return_value = [] + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset()) + + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="agent-secret", user_api_key_auth=mock_user_auth + ) + is False + ) + + with patch.object( + AgentRequestHandler, "_get_agents_from_access_groups", new_callable=AsyncMock + ) as mock_groups: + mock_groups.side_effect = Exception("DB Error") + + assert await AgentRequestHandler._get_allowed_agents_for_key( + user_api_key_auth=mock_user_auth + ) == UnrestrictedAgentAccess() async def test_is_agent_allowed_respects_permissions(self): """ - Test is_agent_allowed: returns True if agent in allowed list or if no restrictions. + Test is_agent_allowed: returns True if agent in allowed list or if unrestricted. Returns False if agent not in allowed list. """ 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" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent1", user_api_key_auth=mock_user_auth @@ -104,9 +193,11 @@ class TestAgentRequestHandler: # Agent not in allowed list - should be denied with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = ["agent1", "agent2"] + mock_get_allowed.return_value = RestrictedAgentAccess( + frozenset({"agent1", "agent2"}) + ) assert ( await AgentRequestHandler.is_agent_allowed( agent_id="agent3", user_api_key_auth=mock_user_auth @@ -114,11 +205,23 @@ class TestAgentRequestHandler: is False ) - # Empty list means no restrictions - should allow any agent + # Restricted to nothing - should deny every agent with patch.object( - AgentRequestHandler, "get_allowed_agents" + AgentRequestHandler, "resolve_agent_access" ) as mock_get_allowed: - mock_get_allowed.return_value = [] + mock_get_allowed.return_value = RestrictedAgentAccess(frozenset()) + assert ( + await AgentRequestHandler.is_agent_allowed( + agent_id="any_agent", user_api_key_auth=mock_user_auth + ) + is False + ) + + # Unrestricted - should allow any agent + with patch.object( + AgentRequestHandler, "resolve_agent_access" + ) as mock_get_allowed: + mock_get_allowed.return_value = UnrestrictedAgentAccess() assert ( await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=mock_user_auth @@ -130,17 +233,19 @@ 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 == [] + result = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=None) + assert result == UnrestrictedAgentAccess() is_allowed = await AgentRequestHandler.is_agent_allowed( agent_id="any_agent", user_api_key_auth=None ) assert is_allowed is True - async def test_get_allowed_agents_handles_errors_gracefully(self): + async def test_resolve_agent_access_handles_errors_gracefully(self): """ - Test that errors during permission lookup are handled gracefully (returns empty list). + Test that errors during permission lookup are handled gracefully. This stays + fail-open for now to preserve existing availability behavior; fail-closed is + tracked separately. """ mock_user_auth = UserAPIKeyAuth( api_key="test-key", @@ -156,12 +261,12 @@ class TestAgentRequestHandler: AgentRequestHandler, "_get_allowed_agents_for_team" ) as mock_team: mock_key.side_effect = Exception("DB Error") - mock_team.return_value = [] + mock_team.return_value = UnrestrictedAgentAccess() - result = await AgentRequestHandler.get_allowed_agents( + result = await AgentRequestHandler.resolve_agent_access( user_api_key_auth=mock_user_auth ) - assert result == [] + assert result == UnrestrictedAgentAccess() async def test_get_allowed_agents_for_key_via_access_group_ids(self): """ @@ -185,7 +290,9 @@ class TestAgentRequestHandler: 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"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag-1", "agent-from-ag-2"}) + ) async def test_get_allowed_agents_for_key_combines_native_and_access_groups(self): """ @@ -215,7 +322,9 @@ class TestAgentRequestHandler: result = await AgentRequestHandler._get_allowed_agents_for_key( user_api_key_auth=mock_user_auth ) - assert sorted(result) == ["agent-from-ag", "native-agent-1"] + assert result == RestrictedAgentAccess( + frozenset({"agent-from-ag", "native-agent-1"}) + ) async def test_is_agent_allowed_accepts_legacy_config_agent_id_grants(self): """LIT-5144: object_permission grants stored under the pre-fix full-entry hash @@ -241,12 +350,13 @@ class TestAgentRequestHandler: "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", registry, ): - with patch.object(AgentRequestHandler, "get_allowed_agents") as mock_get_allowed: + with patch.object(AgentRequestHandler, "resolve_agent_access") as mock_get_allowed: for grant, expected in ( - ([legacy_id], True), - ([agent.agent_id], True), - (["unrelated-agent-id"], False), - ([], True), + (RestrictedAgentAccess(frozenset({legacy_id})), True), + (RestrictedAgentAccess(frozenset({agent.agent_id})), True), + (RestrictedAgentAccess(frozenset({"unrelated-agent-id"})), False), + (RestrictedAgentAccess(frozenset()), False), + (UnrestrictedAgentAccess(), True), ): mock_get_allowed.return_value = grant assert ( @@ -257,10 +367,10 @@ class TestAgentRequestHandler: is expected ), grant - async def test_get_allowed_agents_intersects_legacy_team_grant_with_stable_key_grant(self): + async def test_resolve_agent_access_intersects_legacy_team_grant_with_stable_key_grant(self): """LIT-5144: a team grant stored under the pre-fix full-entry hash and a key grant stored under the name-based id name the same agent; the intersection must resolve - to that agent instead of collapsing to the allow-all empty list.""" + to that agent instead of collapsing to an empty set.""" entry: Final = { "agent_name": "shared-agent", "agent_card_params": { @@ -283,12 +393,21 @@ class TestAgentRequestHandler: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_key") as mock_key: with patch.object(AgentRequestHandler, "_get_allowed_agents_for_team") as mock_team: for key_grant, team_grant in ( - ([agent.agent_id], [legacy_id]), - ([legacy_id], [agent.agent_id]), - ([legacy_id], []), + ( + RestrictedAgentAccess(frozenset({agent.agent_id})), + RestrictedAgentAccess(frozenset({legacy_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + RestrictedAgentAccess(frozenset({agent.agent_id})), + ), + ( + RestrictedAgentAccess(frozenset({legacy_id})), + UnrestrictedAgentAccess(), + ), ): mock_key.return_value = key_grant mock_team.return_value = team_grant - assert await AgentRequestHandler.get_allowed_agents(user_api_key_auth=mock_user_auth) == [ - agent.agent_id - ], (key_grant, team_grant) + assert await AgentRequestHandler.resolve_agent_access( + user_api_key_auth=mock_user_auth + ) == RestrictedAgentAccess(frozenset({agent.agent_id})), (key_grant, team_grant) diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index 6dbd7475be5..ea196bda529 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -6,6 +6,10 @@ from fastapi.testclient import TestClient from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints import endpoints as agent_endpoints +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.endpoints import ( _attach_keys_to_agents, _check_agent_management_permission, @@ -506,9 +510,11 @@ class TestAgentRBACProxyAdminViewOnly: self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id})) monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) - self.allowed_agents_spy = AsyncMock(return_value=["someone-elses-agent"]) + self.allowed_agents_spy = AsyncMock( + return_value=RestrictedAgentAccess(frozenset({"someone-elses-agent"})) + ) monkeypatch.setattr( - "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", self.allowed_agents_spy, ) @@ -537,9 +543,9 @@ class TestAgentRBACProxyAdminViewOnly: self.allowed_agents_spy.assert_awaited_once() def test_should_still_redact_secrets_for_view_only_admin(self): - """An unrestricted viewer (empty allowlist means no restrictions) sees the - same agents as an admin but with keys stripped and litellm_params masked.""" - self.allowed_agents_spy.return_value = [] + """An unrestricted viewer sees the same agents as an admin but with keys + stripped and litellm_params masked.""" + self.allowed_agents_spy.return_value = UnrestrictedAgentAccess() viewer_resp = self._list_agents(self.viewer_client) admin_resp = self._list_agents(self.admin_client) 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..ccf5942c89d 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 @@ -13,6 +13,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, +) from litellm.proxy.agent_endpoints.model_list_helpers import ( append_agents_to_model_group, append_agents_to_model_info, @@ -37,14 +40,14 @@ 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_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(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", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( @@ -80,14 +83,14 @@ async def test_append_agents_to_model_info(): ) # Mock AgentRequestHandler at its source location - mock_get_allowed_agents = AsyncMock(return_value=["agent-123"]) + mock_get_allowed_agents = AsyncMock(return_value=RestrictedAgentAccess(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", + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access", mock_get_allowed_agents, ): with patch( 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..61583d11dfa 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 @@ -13,6 +13,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( + RestrictedAgentAccess, + UnrestrictedAgentAccess, +) # --------------------------------------------------------------------------- @@ -227,8 +231,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_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), # no explicit agent permissions ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -275,8 +279,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_access", + new=AsyncMock(return_value=RestrictedAgentAccess(frozenset({"agent-permitted"}))), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -321,8 +325,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_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", @@ -366,8 +370,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_access", + new=AsyncMock(return_value=UnrestrictedAgentAccess()), ), patch( "litellm.proxy.agent_endpoints.endpoints.get_daily_activity", diff --git a/type-discipline-budget.json b/type-discipline-budget.json index b824f5bb550..991f8eaa934 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23250 + "limit": 23245 }, "LIT002": { - "limit": 27195 + "limit": 27179 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16777 + "limit": 16769 }, "LIT011": { "limit": 5602