scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): deny agent access when key and team grants resolve to nothing (#36221)

Browse source 
* fix(proxy): deny when agent grants resolve to nothing

`get_allowed_agents` returned a plain list where the empty value meant both
"this caller was never restricted" and "this caller's grants resolved to
nothing". Downstream read either as allow-all, so a key restricted to one
agent inside a team restricted to another reached every agent on the proxy,
and an access group that resolved to no agents did the same.

Replace it with `resolve_agent_access`, returning a tagged
UnrestrictedAgentAccess | RestrictedAgentAccess. Only a caller with no grant
anywhere is unrestricted; an empty restricted set denies. Access group lookup
failures now propagate to the key/team resolvers so a DB error still fails
open exactly as before, while a group that genuinely resolves to nothing
denies.

* style(proxy): drop redundant comments from the agent access match
This commit is contained in:
ryan-crabbe-berri 2026-08-07 13:44:11 -07:00 • committed by GitHub
parent 90f8e1f472
commit 78addb230b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 401 additions and 239 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -84,7 +84,7 @@
"limit": 56
},
"reportPrivateUsage": {
"limit": 1825
"limit": 1824
},
"reportRedeclaration": {
"limit": 8

241
litellm/proxy/agent_endpoints/auth/agent_permission_handler.py
View file

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

45
litellm/proxy/agent_endpoints/endpoints.py
View file

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

76
litellm/proxy/agent_endpoints/model_list_helpers.py
View file

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

6
ruff-strict-budget.json
View file

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

8
tests/litellm/proxy/agent_endpoints/test_agent_rbac.py
View file

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

209
tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py
View file

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

16
tests/test_litellm/proxy/agent_endpoints/test_endpoints.py
View file

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

11
tests/test_litellm/proxy/agent_endpoints/test_model_list_helpers.py
View file

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

20
tests/test_litellm/proxy/management_endpoints/test_activity_tenant_scoping.py
View file

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

6
type-discipline-budget.json
View file

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