mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
fix(agents): hide agents from non-admins who were never granted them (#39636)
Listing agents (GET /v1/agents and MCP agent_search) treated the absence of any agent grant on the key or team as permission to see every agent. Non-admin keys now list only the union of explicit grants, and dashboard sessions resolve that union through the user's real teams and user row instead of the shared dashboard team. Proxy admins still see everything and direct access to a named agent is unchanged. Resolves LIT-6862 Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
918ada8d57
commit
959e730d55
3 changed files with 146 additions and 13 deletions
|
|
@ -5,10 +5,13 @@ Handles agent permission checking for keys and teams using object_permission_id.
|
|||
Follows the same pattern as MCP permission handling.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts
|
||||
from litellm.proxy._types import (
|
||||
UI_TEAM_ID,
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
|
|
@ -443,15 +446,47 @@ class AgentRequestHandler:
|
|||
return []
|
||||
|
||||
|
||||
async def accessible_agents(user_api_key_auth: UserAPIKeyAuth) -> tuple[AgentResponse, ...]:
|
||||
"""Every registry agent for proxy admins, else the agents the key's and team's grants reach."""
|
||||
def _granted_ids(access: AgentAccess) -> frozenset[str]:
|
||||
match access:
|
||||
case UnrestrictedAgentAccess():
|
||||
return frozenset()
|
||||
case RestrictedAgentAccess(agent_ids):
|
||||
return agent_ids
|
||||
|
||||
|
||||
ResolveAgentAccess: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[AgentAccess]]
|
||||
EffectiveAuthContexts: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[Sequence[UserAPIKeyAuth]]]
|
||||
|
||||
|
||||
async def _granted_agent_ids(
|
||||
user_api_key_auth: UserAPIKeyAuth,
|
||||
resolve_access: ResolveAgentAccess,
|
||||
effective_contexts: EffectiveAuthContexts,
|
||||
) -> frozenset[str]:
|
||||
"""Union of the explicit grants reachable from the key, its team, or (for a dashboard session)
|
||||
the user's real teams and user row. No grant anywhere yields the empty set, unlike the
|
||||
open-by-default ``resolve_agent_access`` that guards direct access."""
|
||||
accesses: Final = await asyncio.gather(
|
||||
*(resolve_access(auth_context) for auth_context in await effective_contexts(user_api_key_auth))
|
||||
)
|
||||
return frozenset().union(*(_granted_ids(access) for access in accesses))
|
||||
|
||||
|
||||
async def accessible_agents(
|
||||
user_api_key_auth: UserAPIKeyAuth,
|
||||
all_agents: tuple[AgentResponse, ...] | None = None,
|
||||
resolve_access: ResolveAgentAccess | None = None,
|
||||
effective_contexts: EffectiveAuthContexts = build_effective_auth_contexts,
|
||||
) -> tuple[AgentResponse, ...]:
|
||||
"""Every registry agent for proxy admins, else only the agents the caller was granted."""
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
all_agents: Final = global_agent_registry.get_agent_list()
|
||||
agents: Final = global_agent_registry.get_agent_list() if all_agents is None else all_agents
|
||||
if user_api_key_auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN.value):
|
||||
return all_agents
|
||||
match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_auth):
|
||||
case UnrestrictedAgentAccess():
|
||||
return all_agents
|
||||
case RestrictedAgentAccess(allowed_agent_ids):
|
||||
return tuple(agent for agent in all_agents if agent.agent_id in allowed_agent_ids)
|
||||
return agents
|
||||
allowed_agent_ids: Final = await _granted_agent_ids(
|
||||
user_api_key_auth,
|
||||
AgentRequestHandler.resolve_agent_access if resolve_access is None else resolve_access,
|
||||
effective_contexts,
|
||||
)
|
||||
return tuple(agent for agent in agents if agent.agent_id in allowed_agent_ids)
|
||||
|
|
|
|||
|
|
@ -10,15 +10,42 @@ from unittest.mock import AsyncMock, patch
|
|||
import pytest
|
||||
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
AgentAccess,
|
||||
AgentRequestHandler,
|
||||
RestrictedAgentAccess,
|
||||
UnrestrictedAgentAccess,
|
||||
accessible_agents,
|
||||
)
|
||||
|
||||
|
||||
def _registry_with(*agent_names: str) -> AgentRegistry:
|
||||
registry: Final = AgentRegistry()
|
||||
registry.load_agents_from_config(
|
||||
[
|
||||
{
|
||||
"agent_name": name,
|
||||
"agent_card_params": {"name": name, "url": "http://localhost", "version": "1.0.0"},
|
||||
}
|
||||
for name in agent_names
|
||||
]
|
||||
)
|
||||
return registry
|
||||
|
||||
|
||||
def _agent_id(registry: AgentRegistry, agent_name: str) -> str:
|
||||
agent: Final = registry.get_agent_by_name(agent_name)
|
||||
assert agent is not None
|
||||
return agent.agent_id
|
||||
|
||||
|
||||
async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
|
||||
return [user_api_key_auth]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestAgentRequestHandler:
|
||||
"""
|
||||
|
|
@ -265,6 +292,78 @@ class TestAgentRequestHandler:
|
|||
)
|
||||
assert result == UnrestrictedAgentAccess()
|
||||
|
||||
async def test_accessible_agents_hides_ungranted_agents_from_non_admins(self):
|
||||
"""LIT-6862: a key with no agent grant on itself or its team must list nothing,
|
||||
while a proxy admin with the same lack of grants still lists every agent."""
|
||||
registry: Final = _registry_with("alpha", "beta")
|
||||
internal_user: Final = UserAPIKeyAuth(
|
||||
api_key="test-key", user_id="alice", team_id="team-no-perms", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
)
|
||||
proxy_admin: Final = UserAPIKeyAuth(
|
||||
api_key="admin-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
|
||||
)
|
||||
|
||||
async def no_grant_anywhere(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess:
|
||||
return UnrestrictedAgentAccess()
|
||||
|
||||
assert (
|
||||
await accessible_agents(internal_user, registry.get_agent_list(), no_grant_anywhere, _single_context) == ()
|
||||
)
|
||||
assert {
|
||||
agent.agent_name
|
||||
for agent in await accessible_agents(
|
||||
proxy_admin, registry.get_agent_list(), no_grant_anywhere, _single_context
|
||||
)
|
||||
} == {"alpha", "beta"}
|
||||
|
||||
async def test_accessible_agents_lists_only_granted_agents(self):
|
||||
"""A grant for one agent lists that agent and hides the ungranted one."""
|
||||
registry: Final = _registry_with("alpha", "beta")
|
||||
granted_user: Final = UserAPIKeyAuth(
|
||||
api_key="test-key", user_id="bob", team_id="team-granted", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
)
|
||||
|
||||
async def alpha_only(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess:
|
||||
return RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")}))
|
||||
|
||||
listed: Final = await accessible_agents(granted_user, registry.get_agent_list(), alpha_only, _single_context)
|
||||
assert [agent.agent_name for agent in listed] == ["alpha"]
|
||||
|
||||
async def test_accessible_agents_resolves_dashboard_session_through_real_teams_and_user(self):
|
||||
"""LIT-6862: a dashboard session carries the shared litellm-dashboard team id, which holds no
|
||||
grants. Listing must union the grants of the user's real teams and of the user row instead
|
||||
of treating the session as ungranted or as unrestricted."""
|
||||
registry: Final = _registry_with("alpha", "beta", "gamma")
|
||||
session: Final = UserAPIKeyAuth(
|
||||
api_key="session-key",
|
||||
user_id="alice",
|
||||
team_id=UI_SESSION_TOKEN_TEAM_ID,
|
||||
user_role=LitellmUserRoles.INTERNAL_USER,
|
||||
)
|
||||
admitted_user: Final = UserAPIKeyAuth(user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
grants: Final = {
|
||||
"team-granted": RestrictedAgentAccess(frozenset({_agent_id(registry, "alpha")})),
|
||||
"team-no-perms": UnrestrictedAgentAccess(),
|
||||
UI_SESSION_TOKEN_TEAM_ID: UnrestrictedAgentAccess(),
|
||||
}
|
||||
|
||||
async def effective_contexts(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]:
|
||||
assert user_api_key_auth is session
|
||||
return [
|
||||
session.model_copy(update={"team_id": "team-granted"}),
|
||||
session.model_copy(update={"team_id": "team-no-perms"}),
|
||||
admitted_user,
|
||||
]
|
||||
|
||||
async def resolve_access(user_api_key_auth: UserAPIKeyAuth) -> AgentAccess:
|
||||
if user_api_key_auth is admitted_user:
|
||||
return RestrictedAgentAccess(frozenset({_agent_id(registry, "beta")}))
|
||||
assert user_api_key_auth.team_id is not None
|
||||
return grants[user_api_key_auth.team_id]
|
||||
|
||||
listed: Final = await accessible_agents(session, registry.get_agent_list(), resolve_access, effective_contexts)
|
||||
assert {agent.agent_name for agent in listed} == {"alpha", "beta"}
|
||||
|
||||
async def test_get_allowed_agents_for_key_via_access_group_ids(self):
|
||||
"""
|
||||
Test that _get_allowed_agents_for_key includes agents from key's access_group_ids
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ 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,
|
||||
|
|
@ -550,9 +549,9 @@ class TestAgentRBACProxyAdminViewOnly:
|
|||
self.allowed_agents_spy.assert_awaited_once()
|
||||
|
||||
def test_should_still_redact_secrets_for_view_only_admin(self):
|
||||
"""An unrestricted viewer sees the same agents as an admin but with keys
|
||||
"""A viewer granted every agent sees the same agents as an admin but with keys
|
||||
stripped; litellm_params secrets never appear in either response."""
|
||||
self.allowed_agents_spy.return_value = UnrestrictedAgentAccess()
|
||||
self.allowed_agents_spy.return_value = RestrictedAgentAccess(frozenset({"agent-1", "agent-2"}))
|
||||
viewer_resp = self._list_agents(self.viewer_client)
|
||||
admin_resp = self._list_agents(self.admin_client)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue