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>
This commit is contained in:
milan 2026-08-03 22:12:12 +00:00
parent 7c8364c991
commit dd6a6e2fdf
7 changed files with 327 additions and 250 deletions

View file

@ -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(

View file

@ -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.

View file

@ -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)

View file

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

View file

@ -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"}))

View file

@ -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",

View file

@ -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",