mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
feat(agents): attach access groups to agents and enforce them for models, MCP servers and agent calls
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
acf75a525c
commit
b84f8b6a77
30 changed files with 1138 additions and 161 deletions
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -71,6 +71,7 @@ model LiteLLM_AgentsTable {
|
|||
static_headers Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
agent_access_groups String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
object_permission_id String?
|
||||
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
|
||||
spend Float @default(0.0)
|
||||
|
|
|
|||
|
|
@ -1549,10 +1549,18 @@ class MCPRequestHandler:
|
|||
allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(
|
||||
user_api_key_auth
|
||||
)
|
||||
if len(allowed_mcp_servers_for_agent) > 0:
|
||||
agent_access_group_servers: Final = await MCPRequestHandler._get_agent_access_group_server_ceiling(
|
||||
user_api_key_auth
|
||||
)
|
||||
if len(allowed_mcp_servers_for_agent) > 0 or agent_access_group_servers is not None:
|
||||
has_lower_level_mcp_restrictions = True
|
||||
# Intersect: agent can only use servers allowed by BOTH key/team AND agent config
|
||||
allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent]
|
||||
# Intersect: agent can only use servers allowed by key/team AND agent config AND agent access groups
|
||||
allowed_mcp_servers = [
|
||||
s
|
||||
for s in allowed_mcp_servers
|
||||
if (len(allowed_mcp_servers_for_agent) == 0 or s in allowed_mcp_servers_for_agent)
|
||||
and (agent_access_group_servers is None or s in agent_access_group_servers)
|
||||
]
|
||||
verbose_logger.debug(
|
||||
"Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers
|
||||
)
|
||||
|
|
@ -3137,6 +3145,29 @@ class MCPRequestHandler:
|
|||
verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
async def _get_agent_access_group_server_ceiling(
|
||||
user_api_key_auth: UserAPIKeyAuth,
|
||||
) -> frozenset[str] | None:
|
||||
"""
|
||||
Server IDs the agent's attached unified access groups (``LiteLLM_AgentsTable.access_group_ids``)
|
||||
allow, or None when the agent has none attached. Unlike the object_permission path above, an
|
||||
attached group set that names no servers is an empty ceiling and denies every server.
|
||||
"""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
|
||||
if not user_api_key_auth.agent_id:
|
||||
return None
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id)
|
||||
if ceiling is None:
|
||||
return None
|
||||
return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
|
||||
|
||||
@staticmethod
|
||||
async def _get_agent_tool_permissions_for_server(
|
||||
server_id: str,
|
||||
|
|
|
|||
|
|
@ -2357,6 +2357,20 @@
|
|||
},
|
||||
"AgentConfig": {
|
||||
"properties": {
|
||||
"access_group_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Access Group Ids"
|
||||
},
|
||||
"agent_card_params": {
|
||||
"$ref": "#/components/schemas/AgentCard"
|
||||
},
|
||||
|
|
@ -2683,6 +2697,20 @@
|
|||
},
|
||||
"AgentResponse": {
|
||||
"properties": {
|
||||
"access_group_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Access Group Ids"
|
||||
},
|
||||
"agent_card_params": {
|
||||
"additionalProperties": true,
|
||||
"title": "Agent Card Params",
|
||||
|
|
@ -3471,6 +3499,20 @@
|
|||
},
|
||||
"PatchAgentRequest": {
|
||||
"properties": {
|
||||
"access_group_ids": {
|
||||
"anyOf": [
|
||||
{
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"type": "array"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Access Group Ids"
|
||||
},
|
||||
"agent_card_params": {
|
||||
"$ref": "#/components/schemas/AgentCard"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -4127,6 +4127,11 @@ class ProxyErrorTypes(str, enum.Enum):
|
|||
Project does not have access to the model
|
||||
"""
|
||||
|
||||
agent_model_access_denied = "agent_model_access_denied"
|
||||
"""
|
||||
The agent behind the key does not have access to the model
|
||||
"""
|
||||
|
||||
model_cost_map_missing = "model_cost_map_missing"
|
||||
|
||||
expired_key = "expired_key"
|
||||
|
|
@ -4201,7 +4206,7 @@ class ProxyErrorTypes(str, enum.Enum):
|
|||
|
||||
@classmethod
|
||||
def get_model_access_error_type_for_object(
|
||||
cls, object_type: Literal["key", "user", "team", "org", "project"]
|
||||
cls, object_type: Literal["key", "user", "team", "org", "project", "agent"]
|
||||
) -> "ProxyErrorTypes":
|
||||
"""
|
||||
Get the model access error type for object_type
|
||||
|
|
@ -4216,6 +4221,8 @@ class ProxyErrorTypes(str, enum.Enum):
|
|||
return cls.org_model_access_denied
|
||||
elif object_type == "project":
|
||||
return cls.project_model_access_denied
|
||||
elif object_type == "agent":
|
||||
return cls.agent_model_access_denied
|
||||
|
||||
@classmethod
|
||||
def get_vector_store_access_error_type_for_object(
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from types import MappingProxyType
|
|||
from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm.constants import REDACTED_BY_LITELM_STRING
|
||||
|
|
@ -37,6 +38,7 @@ class AgentRecordDump(TypedDict):
|
|||
agent_card_params: dict[str, object]
|
||||
static_headers: dict[str, str] | None
|
||||
extra_headers: list[str] | None
|
||||
access_group_ids: ReadOnly[Sequence[str] | None]
|
||||
object_permission: dict[str, object] | None
|
||||
spend: float
|
||||
tpm_limit: int | None
|
||||
|
|
@ -284,6 +286,12 @@ def _resolved_agent_param_value(
|
|||
return _MISSING_AGENT_PARAM
|
||||
|
||||
|
||||
def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]:
|
||||
if "access_group_ids" not in agent:
|
||||
return MappingProxyType({})
|
||||
return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))})
|
||||
|
||||
|
||||
def _restore_redacted_litellm_params(
|
||||
incoming: Mapping[str, object],
|
||||
existing: Mapping[str, object],
|
||||
|
|
@ -516,6 +524,7 @@ class AgentRegistry:
|
|||
static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None
|
||||
|
||||
extra_headers_val: Final = agent.get("extra_headers")
|
||||
access_group_ids_val: Final = agent.get("access_group_ids")
|
||||
|
||||
create_data: Final[dict[str, object]] = {
|
||||
"agent_name": agent_name,
|
||||
|
|
@ -532,6 +541,8 @@ class AgentRegistry:
|
|||
create_data["static_headers"] = static_headers_val
|
||||
if extra_headers_val is not None:
|
||||
create_data["extra_headers"] = extra_headers_val
|
||||
if access_group_ids_val is not None:
|
||||
create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val))
|
||||
if object_permission_id is not None:
|
||||
create_data["object_permission_id"] = object_permission_id
|
||||
|
||||
|
|
@ -601,7 +612,7 @@ class AgentRegistry:
|
|||
existing_agent: Final[Mapping[str, object]] = dict(existing_record)
|
||||
|
||||
augment_agent: Final = {**existing_agent, **agent}
|
||||
update_data: Final[dict[str, object]] = {}
|
||||
update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)}
|
||||
if augment_agent.get("agent_name"):
|
||||
update_data["agent_name"] = augment_agent.get("agent_name")
|
||||
if "litellm_params" in agent:
|
||||
|
|
@ -703,6 +714,7 @@ class AgentRegistry:
|
|||
safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({})
|
||||
)
|
||||
extra_headers_val_u: Final = agent.get("extra_headers") or []
|
||||
access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ()))
|
||||
|
||||
update_data: Final[dict[str, object]] = {
|
||||
"agent_name": agent_name,
|
||||
|
|
@ -710,6 +722,7 @@ class AgentRegistry:
|
|||
"agent_card_params": agent_card_params,
|
||||
"static_headers": static_headers_val_u,
|
||||
"extra_headers": extra_headers_val_u,
|
||||
"access_group_ids": access_group_ids_val_u,
|
||||
"updated_by": updated_by,
|
||||
"updated_at": datetime.now(timezone.utc),
|
||||
}
|
||||
|
|
|
|||
80
litellm/proxy/agent_endpoints/auth/agent_access_groups.py
Normal file
80
litellm/proxy/agent_endpoints/auth/agent_access_groups.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
"""
|
||||
Ceiling that an agent's attached access groups place on requests made with that agent's key.
|
||||
|
||||
Keys and teams use access groups as grants. An agent uses them the way it already uses its
|
||||
``object_permission``: the union of the attached groups caps what the agent's key can reach,
|
||||
on top of whatever the key and team allow. A group that cannot be loaded contributes nothing,
|
||||
so a missing or unreadable group can only narrow the agent, never widen it.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import LiteLLM_AccessGroupTable
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]]
|
||||
AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LiteLLM_AccessGroupTable | None]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentAccessGroupCeiling:
|
||||
"""Everything the agent's attached access groups allow. An empty set denies that resource kind."""
|
||||
|
||||
access_group_ids: tuple[str, ...]
|
||||
models: frozenset[str]
|
||||
mcp_server_ids: frozenset[str]
|
||||
agent_ids: frozenset[str]
|
||||
|
||||
|
||||
async def _load_agent(agent_id: str) -> AgentResponse | None:
|
||||
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
|
||||
|
||||
return await get_agent_with_read_through(agent_id)
|
||||
|
||||
|
||||
async def _load_access_group(access_group_id: str) -> LiteLLM_AccessGroupTable | None:
|
||||
from litellm.proxy.auth.auth_checks import get_access_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id)
|
||||
return None
|
||||
try:
|
||||
return await get_access_object(
|
||||
access_group_id=access_group_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def resolve_agent_access_group_ceiling(
|
||||
agent_id: str,
|
||||
load_agent: AgentLoader = _load_agent,
|
||||
load_access_group: AccessGroupLoader = _load_access_group,
|
||||
) -> AgentAccessGroupCeiling | None:
|
||||
"""``None`` when the agent has no access groups attached, so nothing is capped."""
|
||||
agent: Final = await load_agent(agent_id)
|
||||
access_group_ids: Final = tuple(agent.access_group_ids or ()) if agent is not None else ()
|
||||
if not access_group_ids:
|
||||
return None
|
||||
|
||||
loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids))
|
||||
groups: Final = tuple(group for group in loaded if group is not None)
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=access_group_ids,
|
||||
models=frozenset(model for group in groups for model in group.access_model_names),
|
||||
mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids),
|
||||
agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids),
|
||||
)
|
||||
|
|
@ -66,9 +66,24 @@ class AgentRequestHandler:
|
|||
Resolve the agents the given user/key may reach.
|
||||
|
||||
``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.
|
||||
carries any grant and the agent behind the key has no access groups attached.
|
||||
Grants that intersect to nothing stay restricted, so narrowing a caller can
|
||||
never widen what it reaches.
|
||||
"""
|
||||
key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth)
|
||||
agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth)
|
||||
if agent_ceiling is None:
|
||||
return key_team_access
|
||||
match key_team_access:
|
||||
case UnrestrictedAgentAccess():
|
||||
return RestrictedAgentAccess(agent_ceiling)
|
||||
case RestrictedAgentAccess(key_team_ids):
|
||||
return RestrictedAgentAccess(key_team_ids & agent_ceiling)
|
||||
|
||||
@staticmethod
|
||||
async def _resolve_key_team_agent_access(
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
) -> AgentAccess:
|
||||
try:
|
||||
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)
|
||||
|
|
@ -86,6 +101,20 @@ class AgentRequestHandler:
|
|||
verbose_logger.warning("Failed to get allowed agents: %s", e)
|
||||
return UnrestrictedAgentAccess()
|
||||
|
||||
@staticmethod
|
||||
async def _agent_access_group_ceiling(
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
) -> frozenset[str] | None:
|
||||
"""Stable IDs of the agents the calling agent's attached access groups allow; None when none attached."""
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling
|
||||
|
||||
if user_api_key_auth is None or not user_api_key_auth.agent_id:
|
||||
return None
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id)
|
||||
if ceiling is None:
|
||||
return None
|
||||
return _to_stable_ids(ceiling.agent_ids)
|
||||
|
||||
@staticmethod
|
||||
async def is_agent_allowed(
|
||||
agent_id: str,
|
||||
|
|
|
|||
|
|
@ -1003,6 +1003,9 @@ async def common_checks(
|
|||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
# 2.4 If the agent behind the key has access groups attached, they cap the models it can call
|
||||
await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router)
|
||||
|
||||
## 2.1 If user can call model (if personal key)
|
||||
if _model and team_object is None and user_object is not None:
|
||||
with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"):
|
||||
|
|
@ -4126,7 +4129,7 @@ def _can_object_call_model(
|
|||
models: list[str],
|
||||
team_model_aliases: dict[str, str] | None = None,
|
||||
team_id: str | None = None,
|
||||
object_type: Literal["user", "team", "key", "org", "project"] = "user",
|
||||
object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user",
|
||||
fallback_depth: int = 0,
|
||||
) -> Literal[True]:
|
||||
"""
|
||||
|
|
@ -4192,6 +4195,38 @@ def _can_object_call_model(
|
|||
)
|
||||
|
||||
|
||||
async def _check_agent_access_group_model_access(
|
||||
model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str]
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
llm_router: Router | None,
|
||||
) -> Literal[True]:
|
||||
"""Raises when the key's agent has access groups attached and none of them names the model.
|
||||
Attached groups that name no model deny every model; ``_can_object_call_model`` would read
|
||||
an empty allowlist as unrestricted."""
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling
|
||||
|
||||
if not model or valid_token is None or not valid_token.agent_id:
|
||||
return True
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling(valid_token.agent_id)
|
||||
if ceiling is None:
|
||||
return True
|
||||
if not ceiling.models:
|
||||
raise ModelAccessDeniedProxyException(
|
||||
message=model_access_denied_client_message(model=model),
|
||||
internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models",
|
||||
type=ProxyErrorTypes.agent_model_access_denied,
|
||||
param="model",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
return _can_object_call_model(
|
||||
model=model,
|
||||
llm_router=llm_router,
|
||||
models=sorted(ceiling.models),
|
||||
team_id=valid_token.team_id,
|
||||
object_type="agent",
|
||||
)
|
||||
|
||||
|
||||
def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool:
|
||||
"""
|
||||
Returns True if `model` being accessed is an alias of a team model
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ from types import MappingProxyType
|
|||
from typing import Final, Protocol
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
|
|
@ -109,6 +110,36 @@ class _KeyTable(Protocol):
|
|||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
|
||||
|
||||
|
||||
class _AgentRecord(Protocol):
|
||||
@property
|
||||
def agent_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def access_group_ids(self) -> Sequence[str] | None: ...
|
||||
|
||||
|
||||
class _AgentTable(Protocol):
|
||||
async def find_many(self, where: Mapping[str, object]) -> Sequence[_AgentRecord]: ...
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
|
||||
|
||||
|
||||
class _HasSomeFilter(TypedDict):
|
||||
hasSome: ReadOnly[Sequence[str]]
|
||||
|
||||
|
||||
class _AgentAccessGroupsWhere(TypedDict):
|
||||
access_group_ids: ReadOnly[_HasSomeFilter]
|
||||
|
||||
|
||||
class _AgentIdWhere(TypedDict):
|
||||
agent_id: ReadOnly[str]
|
||||
|
||||
|
||||
class _AgentAccessGroupsData(TypedDict):
|
||||
access_group_ids: ReadOnly[Sequence[str]]
|
||||
|
||||
|
||||
class _AccessGroupTx(Protocol):
|
||||
@property
|
||||
def litellm_accessgrouptable(self) -> _AccessGroupTable: ...
|
||||
|
|
@ -119,6 +150,9 @@ class _AccessGroupTx(Protocol):
|
|||
@property
|
||||
def litellm_verificationtoken(self) -> _KeyTable: ...
|
||||
|
||||
@property
|
||||
def litellm_agentstable(self) -> _AgentTable: ...
|
||||
|
||||
|
||||
def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None:
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
|
|
@ -324,6 +358,41 @@ async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: li
|
|||
)
|
||||
|
||||
|
||||
def _without_access_group(access_group_ids: Sequence[str] | None, access_group_id: str) -> tuple[str, ...]:
|
||||
return tuple(ag for ag in (access_group_ids or ()) if ag != access_group_id)
|
||||
|
||||
|
||||
async def _detach_access_group_from_agents(tx: _AccessGroupTx, access_group_id: str) -> tuple[str, ...]:
|
||||
agents_with_group: Final = await tx.litellm_agentstable.find_many(
|
||||
where=_AgentAccessGroupsWhere(access_group_ids=_HasSomeFilter(hasSome=(access_group_id,)))
|
||||
)
|
||||
for agent in agents_with_group:
|
||||
await tx.litellm_agentstable.update(
|
||||
where=_AgentIdWhere(agent_id=agent.agent_id),
|
||||
data=_AgentAccessGroupsData(
|
||||
access_group_ids=_without_access_group(agent.access_group_ids, access_group_id)
|
||||
),
|
||||
)
|
||||
return tuple(agent.agent_id for agent in agents_with_group)
|
||||
|
||||
|
||||
def _detach_access_group_from_agent_registry(agent_ids: Sequence[str], access_group_id: str) -> None:
|
||||
registered: Final = tuple(
|
||||
agent
|
||||
for agent in (global_agent_registry.get_agent_by_id(agent_id) for agent_id in agent_ids)
|
||||
if agent is not None
|
||||
)
|
||||
for agent in registered:
|
||||
global_agent_registry.deregister_agent(agent_name=agent.agent_name)
|
||||
global_agent_registry.register_agent(
|
||||
agent_config=agent.model_copy(
|
||||
update=_AgentAccessGroupsData(
|
||||
access_group_ids=_without_access_group(agent.access_group_ids, access_group_id)
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cache patch helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -705,11 +774,14 @@ async def delete_access_group(
|
|||
out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group}
|
||||
await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id)
|
||||
|
||||
detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id)
|
||||
|
||||
await tx.litellm_accessgrouptable.delete(where={"access_group_id": access_group_id})
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
|
||||
|
||||
await invalidate_access_group_cache(access_group_id)
|
||||
_detach_access_group_from_agent_registry(detached_agent_ids, access_group_id)
|
||||
await _patch_team_caches_remove_access_group(
|
||||
affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj
|
||||
)
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ model LiteLLM_AgentsTable {
|
|||
static_headers Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
agent_access_groups String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
object_permission_id String?
|
||||
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
|
||||
spend Float @default(0.0)
|
||||
|
|
|
|||
|
|
@ -189,6 +189,7 @@ class AgentConfig(TypedDict, total=False):
|
|||
session_rpm_limit: int | None
|
||||
static_headers: dict[str, str] | None
|
||||
extra_headers: list[str] | None
|
||||
access_group_ids: ReadOnly[Sequence[str] | None]
|
||||
|
||||
|
||||
class PatchAgentRequest(TypedDict, total=False):
|
||||
|
|
@ -202,6 +203,7 @@ class PatchAgentRequest(TypedDict, total=False):
|
|||
session_rpm_limit: int | None
|
||||
static_headers: dict[str, str] | None
|
||||
extra_headers: list[str] | None
|
||||
access_group_ids: ReadOnly[Sequence[str] | None]
|
||||
|
||||
|
||||
# Request/Response models for CRUD endpoints
|
||||
|
|
@ -226,6 +228,7 @@ class AgentResponse(BaseModel):
|
|||
session_rpm_limit: int | None = None
|
||||
static_headers: dict[str, str] | None = None
|
||||
extra_headers: list[str] | None = None
|
||||
access_group_ids: Sequence[str] | None = None
|
||||
keys: list[AgentKeySummary] | None = None
|
||||
search_score: float | None = None
|
||||
created_at: datetime | None = None
|
||||
|
|
|
|||
|
|
@ -71,6 +71,7 @@ model LiteLLM_AgentsTable {
|
|||
static_headers Json? @default("{}")
|
||||
extra_headers String[] @default([])
|
||||
agent_access_groups String[] @default([])
|
||||
access_group_ids String[] @default([])
|
||||
object_permission_id String?
|
||||
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
|
||||
spend Float @default(0.0)
|
||||
|
|
|
|||
|
|
@ -4208,6 +4208,65 @@ class TestAgentMCPPermissions:
|
|||
assert sorted(result) == ["server_1", "server_2"]
|
||||
mock_agent.assert_called_once_with(user_api_key_auth)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("group_ceiling", "expected"),
|
||||
[
|
||||
(frozenset({"server_1"}), ["server_1"]),
|
||||
(frozenset({"server_1", "server_2", "server_3"}), ["server_1", "server_2"]),
|
||||
(frozenset(), []),
|
||||
],
|
||||
)
|
||||
async def test_get_allowed_mcp_servers_agent_access_group_ceiling(self, group_ceiling, expected):
|
||||
"""The agent's attached access groups cap the key/team servers; groups naming no server deny all."""
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag")
|
||||
with (
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]),
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]),
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]),
|
||||
patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=group_ceiling),
|
||||
):
|
||||
access = await MCPRequestHandler.get_mcp_server_access(user_api_key_auth=user_api_key_auth)
|
||||
assert sorted(access.server_ids) == expected
|
||||
assert access.scope == "scoped"
|
||||
|
||||
async def test_get_allowed_mcp_servers_agent_without_access_groups_is_uncapped(self):
|
||||
user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag")
|
||||
with (
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]),
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]),
|
||||
patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]),
|
||||
patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=None),
|
||||
):
|
||||
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=user_api_key_auth)
|
||||
assert sorted(result) == ["server_1", "server_2"]
|
||||
|
||||
async def test_agent_access_group_server_ceiling_expands_group_servers(self):
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
|
||||
|
||||
ceiling = AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",),
|
||||
models=frozenset(),
|
||||
mcp_server_ids=frozenset({"server_1"}),
|
||||
agent_ids=frozenset(),
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=ceiling),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager"
|
||||
) as mock_manager,
|
||||
):
|
||||
mock_manager.expand_permission_list.return_value = ["server_1"]
|
||||
result = await MCPRequestHandler._get_agent_access_group_server_ceiling(
|
||||
UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag")
|
||||
)
|
||||
assert result == frozenset({"server_1"})
|
||||
mock_manager.expand_permission_list.assert_called_once_with(["server_1"])
|
||||
|
||||
assert await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k")) is None
|
||||
|
||||
async def test_get_allowed_mcp_servers_key_team_agent_intersection(self):
|
||||
"""Key allows [1, 2], agent allows [2, 3]. Result = [2]."""
|
||||
user_api_key_auth = UserAPIKeyAuth(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,130 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.models.access_group import LiteLLM_AccessGroupTable
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
AgentAccessGroupCeiling,
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
_CARD: Final = {"name": "agent", "url": "http://localhost:9999", "version": "1.0.0"}
|
||||
|
||||
|
||||
def _agent(access_group_ids: list[str] | None) -> AgentResponse:
|
||||
return AgentResponse(
|
||||
agent_id="agent-1", agent_name="agent", agent_card_params=_CARD, access_group_ids=access_group_ids
|
||||
)
|
||||
|
||||
|
||||
def _group(
|
||||
group_id: str,
|
||||
models: tuple[str, ...] = (),
|
||||
mcp_servers: tuple[str, ...] = (),
|
||||
agents: tuple[str, ...] = (),
|
||||
) -> LiteLLM_AccessGroupTable:
|
||||
return LiteLLM_AccessGroupTable(
|
||||
access_group_id=group_id,
|
||||
access_group_name=group_id,
|
||||
access_model_names=list(models),
|
||||
access_mcp_server_ids=list(mcp_servers),
|
||||
access_agent_ids=list(agents),
|
||||
)
|
||||
|
||||
|
||||
def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]):
|
||||
async def load_agent(agent_id: str) -> AgentResponse | None:
|
||||
return agent
|
||||
|
||||
async def load_group(group_id: str) -> LiteLLM_AccessGroupTable | None:
|
||||
return groups.get(group_id)
|
||||
|
||||
return load_agent, load_group
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("access_group_ids", [None, []])
|
||||
async def test_agent_without_access_groups_has_no_ceiling(access_group_ids: list[str] | None):
|
||||
load_agent, load_group = _loaders(_agent(access_group_ids), {"g1": _group("g1", models=("gpt-5",))})
|
||||
|
||||
assert await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_agent_has_no_ceiling():
|
||||
load_agent, load_group = _loaders(None, {})
|
||||
|
||||
assert await resolve_agent_access_group_ceiling("missing", load_agent, load_group) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ceiling_is_the_union_of_every_attached_group():
|
||||
load_agent, load_group = _loaders(
|
||||
_agent(["g1", "g2"]),
|
||||
{
|
||||
"g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)),
|
||||
"g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)),
|
||||
},
|
||||
)
|
||||
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group)
|
||||
|
||||
assert ceiling == AgentAccessGroupCeiling(
|
||||
access_group_ids=("g1", "g2"),
|
||||
models=frozenset({"gpt-5", "claude-sonnet"}),
|
||||
mcp_server_ids=frozenset({"mcp-a", "mcp-b"}),
|
||||
agent_ids=frozenset({"agent-b", "agent-c"}),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies():
|
||||
load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))})
|
||||
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group)
|
||||
|
||||
assert ceiling == AgentAccessGroupCeiling(
|
||||
access_group_ids=("g1", "gone"),
|
||||
models=frozenset({"gpt-5"}),
|
||||
mcp_server_ids=frozenset(),
|
||||
agent_ids=frozenset(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted():
|
||||
load_agent, load_group = _loaders(_agent(["gone"]), {})
|
||||
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group)
|
||||
|
||||
assert ceiling is not None
|
||||
assert ceiling.models == frozenset()
|
||||
assert ceiling.mcp_server_ids == frozenset()
|
||||
assert ceiling.agent_ids == frozenset()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group
|
||||
from litellm.proxy.auth import auth_checks
|
||||
|
||||
async def missing_group(**_: object) -> LiteLLM_AccessGroupTable:
|
||||
raise HTTPException(status_code=404, detail={"error": "Access group doesn't exist in db."})
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", object())
|
||||
monkeypatch.setattr(auth_checks, "get_access_object", missing_group)
|
||||
|
||||
assert await _load_access_group("gone") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch):
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
|
||||
assert await _load_access_group("ag-1") is None
|
||||
|
|
@ -13,6 +13,7 @@ import pytest
|
|||
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_access_groups import AgentAccessGroupCeiling
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
AgentAccess,
|
||||
AgentRequestHandler,
|
||||
|
|
@ -157,6 +158,88 @@ class TestAgentRequestHandler:
|
|||
is False
|
||||
), agent_id
|
||||
|
||||
@staticmethod
|
||||
def _ceiling(agent_ids: frozenset[str]) -> AgentAccessGroupCeiling:
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",),
|
||||
models=frozenset(),
|
||||
mcp_server_ids=frozenset(),
|
||||
agent_ids=agent_ids,
|
||||
)
|
||||
|
||||
async def test_agent_access_groups_cap_an_otherwise_unrestricted_key(self):
|
||||
"""A key with no agent grant of its own may still only reach the agents its
|
||||
agent's attached access groups name."""
|
||||
agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent")
|
||||
|
||||
with (
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()),
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()),
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta"}))),
|
||||
) as mock_ceiling,
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess(
|
||||
frozenset({"agent-beta"})
|
||||
)
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key) is True
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False
|
||||
mock_ceiling.assert_called_with("caller-agent")
|
||||
|
||||
async def test_agent_access_groups_intersect_with_key_and_team_grants(self):
|
||||
agent_key: Final = UserAPIKeyAuth(
|
||||
api_key="test-key", user_id="test-user", team_id="test-team", agent_id="caller-agent"
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
AgentRequestHandler,
|
||||
"_get_allowed_agents_for_key",
|
||||
return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta"})),
|
||||
),
|
||||
patch.object(
|
||||
AgentRequestHandler,
|
||||
"_get_allowed_agents_for_team",
|
||||
return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta", "agent-gamma"})),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta", "agent-gamma"}))),
|
||||
),
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess(
|
||||
frozenset({"agent-beta"})
|
||||
)
|
||||
|
||||
async def test_agent_access_groups_naming_no_agent_deny_every_agent(self):
|
||||
agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent")
|
||||
|
||||
with (
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()),
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()),
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=self._ceiling(frozenset())),
|
||||
),
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess(frozenset())
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False
|
||||
|
||||
async def test_key_without_agent_never_consults_agent_access_groups(self):
|
||||
plain_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
|
||||
|
||||
with (
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()),
|
||||
patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()),
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=self._ceiling(frozenset())),
|
||||
) as mock_ceiling,
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(plain_key) == UnrestrictedAgentAccess()
|
||||
mock_ceiling.assert_not_called()
|
||||
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -990,3 +990,136 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|||
stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"])
|
||||
assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY
|
||||
assert stored_params["is_public"] is True
|
||||
|
||||
|
||||
def _agent_row_mock(access_group_ids: list[str]) -> MagicMock:
|
||||
row: Final = MagicMock()
|
||||
row.model_dump.return_value = {
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
"litellm_params": {},
|
||||
"object_permission": None,
|
||||
"access_group_ids": access_group_ids,
|
||||
}
|
||||
row.object_permission = None
|
||||
return row
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_agent_to_db_persists_deduplicated_access_group_ids():
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_create = AsyncMock(return_value=_agent_row_mock(["ag-1", "ag-2"]))
|
||||
mock_prisma.db.litellm_agentstable.create = mock_create
|
||||
|
||||
result: Final = await registry.add_agent_to_db(
|
||||
agent={
|
||||
"agent_name": "Test Agent",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
"access_group_ids": ["ag-1", "ag-2", "ag-1"],
|
||||
},
|
||||
prisma_client=mock_prisma,
|
||||
created_by="test-user",
|
||||
)
|
||||
|
||||
assert tuple(mock_create.call_args.kwargs["data"]["access_group_ids"]) == ("ag-1", "ag-2")
|
||||
assert result.access_group_ids == ["ag-1", "ag-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_default():
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_create = AsyncMock(return_value=_agent_row_mock([]))
|
||||
mock_prisma.db.litellm_agentstable.create = mock_create
|
||||
|
||||
await registry.add_agent_to_db(
|
||||
agent={"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params()},
|
||||
prisma_client=mock_prisma,
|
||||
created_by="test-user",
|
||||
)
|
||||
|
||||
assert "access_group_ids" not in mock_create.call_args.kwargs["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("patch_body", "expected"),
|
||||
[
|
||||
({"access_group_ids": ["ag-2", "ag-3"]}, ["ag-2", "ag-3"]),
|
||||
({"access_group_ids": []}, []),
|
||||
({"access_group_ids": None}, []),
|
||||
],
|
||||
)
|
||||
async def test_patch_agent_in_db_replaces_access_group_ids_when_provided(patch_body: dict, expected: list[str]):
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
||||
await registry.patch_agent_in_db(
|
||||
agent_id="agent-123", agent=patch_body, prisma_client=mock_prisma, updated_by="test-user"
|
||||
)
|
||||
|
||||
assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted():
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old Name",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"]))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
||||
await registry.patch_agent_in_db(
|
||||
agent_id="agent-123", agent={"agent_name": "New Name"}, prisma_client=mock_prisma, updated_by="test-user"
|
||||
)
|
||||
|
||||
assert "access_group_ids" not in mock_update.call_args.kwargs["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("body_access_group_ids", "expected"),
|
||||
[(["ag-9", "ag-9"], ["ag-9"]), (None, []), ("omitted", [])],
|
||||
)
|
||||
async def test_update_agent_in_db_always_writes_access_group_ids(body_access_group_ids, expected: list[str]):
|
||||
"""PUT is a full replacement: omitting the field clears any previously attached groups."""
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"])
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
body: Final = {
|
||||
"agent_name": "Test Agent",
|
||||
"agent_card_params": _sample_agent_card_params(),
|
||||
"litellm_params": {"model": "bedrock/agentcore/my-agent"},
|
||||
**({} if body_access_group_ids == "omitted" else {"access_group_ids": body_access_group_ids}),
|
||||
}
|
||||
|
||||
await registry.update_agent_in_db(
|
||||
agent_id="agent-123", agent=body, prisma_client=mock_prisma, updated_by="test-user"
|
||||
)
|
||||
|
||||
assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected)
|
||||
|
|
|
|||
|
|
@ -8461,3 +8461,94 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
|
|||
def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None:
|
||||
assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True
|
||||
assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False
|
||||
|
||||
|
||||
# Agent access group model ceiling
|
||||
|
||||
|
||||
def _agent_model_ceiling(models: frozenset[str]):
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
|
||||
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset()
|
||||
)
|
||||
|
||||
|
||||
async def _run_common_checks_for_agent_key(model: str, valid_token: UserAPIKeyAuth):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.auth.auth_checks import common_checks
|
||||
|
||||
return await common_checks(
|
||||
request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]},
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
end_user_object=None,
|
||||
global_proxy_spend=None,
|
||||
general_settings={},
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
valid_token=valid_token,
|
||||
request=MagicMock(spec=Request),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_agent_access_groups_cap_models_even_when_key_allows_them():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=_agent_model_ceiling(frozenset({"gpt-5"}))),
|
||||
):
|
||||
assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _run_common_checks_for_agent_key("claude-sonnet", agent_key)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied
|
||||
assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_agent_access_groups_naming_no_model_deny_every_model():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[])
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=_agent_model_ceiling(frozenset())),
|
||||
),
|
||||
pytest.raises(ProxyException) as exc_info,
|
||||
):
|
||||
await _run_common_checks_for_agent_key("gpt-5", agent_key)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_agent_without_access_groups_adds_no_model_ceiling():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as mock_ceiling:
|
||||
assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True
|
||||
assert await _run_common_checks_for_agent_key("claude-sonnet", agent_key) is True
|
||||
|
||||
mock_ceiling.assert_called_with("agent-1")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_common_checks_key_without_agent_never_consults_agent_access_groups():
|
||||
plain_key: Final = UserAPIKeyAuth(token="plain-token", models=["gpt-5"])
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling",
|
||||
new=AsyncMock(return_value=_agent_model_ceiling(frozenset())),
|
||||
) as mock_ceiling:
|
||||
assert await _run_common_checks_for_agent_key("gpt-5", plain_key) is True
|
||||
|
||||
mock_ceiling.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -17,7 +17,9 @@ from litellm.proxy._types import (
|
|||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
from litellm.proxy.proxy_server import app
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
|
||||
def _make_access_group_record(
|
||||
|
|
@ -126,6 +128,7 @@ def client_and_mocks(monkeypatch):
|
|||
|
||||
mock_agents_table = MagicMock()
|
||||
mock_agents_table.find_many = AsyncMock(return_value=[])
|
||||
mock_agents_table.update = AsyncMock(return_value=None)
|
||||
|
||||
@asynccontextmanager
|
||||
async def mock_tx():
|
||||
|
|
@ -133,6 +136,7 @@ def client_and_mocks(monkeypatch):
|
|||
litellm_accessgrouptable=mock_access_group_table,
|
||||
litellm_teamtable=mock_team_table,
|
||||
litellm_verificationtoken=mock_key_table,
|
||||
litellm_agentstable=mock_agents_table,
|
||||
)
|
||||
yield tx
|
||||
|
||||
|
|
@ -158,15 +162,9 @@ def client_and_mocks(monkeypatch):
|
|||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.internal_usage_cache = MagicMock()
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache = MagicMock()
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(return_value=None)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock(return_value=None)
|
||||
monkeypatch.setattr(ps, "proxy_logging_obj", mock_proxy_logging)
|
||||
|
||||
admin_user = UserAPIKeyAuth(
|
||||
|
|
@ -239,9 +237,7 @@ def test_create_access_group_duplicate_name_conflict(client_and_mocks):
|
|||
"unique constraint violation",
|
||||
],
|
||||
)
|
||||
def test_create_access_group_race_condition_returns_409(
|
||||
client_and_mocks, error_message
|
||||
):
|
||||
def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message):
|
||||
"""Create race condition: Prisma unique constraint surfaces as 409, not 500."""
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
|
|
@ -288,9 +284,7 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks
|
|||
|
||||
# Use raise_server_exceptions=False so unhandled exceptions become 500 responses
|
||||
test_client = TestClient(app, raise_server_exceptions=False)
|
||||
resp = test_client.post(
|
||||
"/v1/access_group", json={"access_group_name": "test-group"}
|
||||
)
|
||||
resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"})
|
||||
assert resp.status_code == 500
|
||||
|
||||
|
||||
|
|
@ -558,9 +552,7 @@ def test_update_access_group_empty_body(client_and_mocks):
|
|||
"""Update with empty body succeeds; only updated_by is set."""
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", access_group_name="unchanged"
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged")
|
||||
mock_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
resp = client.put("/v1/access_group/ag-update", json={})
|
||||
|
|
@ -576,14 +568,10 @@ def test_update_access_group_name_success(client_and_mocks):
|
|||
"""Update access_group_name succeeds when new name is unique."""
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", access_group_name="old-name"
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name")
|
||||
mock_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
resp = client.put(
|
||||
"/v1/access_group/ag-update", json={"access_group_name": "new-name"}
|
||||
)
|
||||
resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "new-name"})
|
||||
assert resp.status_code == 200
|
||||
mock_table.update.assert_awaited_once()
|
||||
call_kwargs = mock_table.update.call_args.kwargs
|
||||
|
|
@ -594,19 +582,13 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks):
|
|||
"""Update access_group_name to existing name returns 409 (unique constraint)."""
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", access_group_name="old-name"
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name")
|
||||
mock_table.find_unique = AsyncMock(return_value=existing)
|
||||
mock_table.update = AsyncMock(
|
||||
side_effect=Exception(
|
||||
"Unique constraint failed on the fields: (`access_group_name`)"
|
||||
)
|
||||
side_effect=Exception("Unique constraint failed on the fields: (`access_group_name`)")
|
||||
)
|
||||
|
||||
resp = client.put(
|
||||
"/v1/access_group/ag-update", json={"access_group_name": "taken-name"}
|
||||
)
|
||||
resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "taken-name"})
|
||||
assert resp.status_code == 409
|
||||
assert "already exists" in resp.json()["detail"]
|
||||
mock_table.update.assert_awaited_once()
|
||||
|
|
@ -620,21 +602,15 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks):
|
|||
"unique constraint violation",
|
||||
],
|
||||
)
|
||||
def test_update_access_group_name_unique_constraint_returns_409(
|
||||
client_and_mocks, error_message
|
||||
):
|
||||
def test_update_access_group_name_unique_constraint_returns_409(client_and_mocks, error_message):
|
||||
"""Update access_group_name: Prisma unique constraint surfaces as 409."""
|
||||
client, _, mock_table, *_ = client_and_mocks
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", access_group_name="old-name"
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name")
|
||||
mock_table.find_unique = AsyncMock(return_value=existing)
|
||||
mock_table.update = AsyncMock(side_effect=Exception(error_message))
|
||||
|
||||
resp = client.put(
|
||||
"/v1/access_group/ag-update", json={"access_group_name": "race-name"}
|
||||
)
|
||||
resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "race-name"})
|
||||
assert resp.status_code == 409
|
||||
assert "already exists" in resp.json()["detail"]
|
||||
|
||||
|
|
@ -690,9 +666,7 @@ def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role):
|
|||
|
||||
def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks):
|
||||
"""Delete removes access_group_id from teams and keys before deleting the group."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
|
|
@ -722,10 +696,61 @@ def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks):
|
|||
where={"token": "key-token-1"},
|
||||
data={"access_group_ids": []},
|
||||
)
|
||||
mock_access_group_table.delete.assert_awaited_once_with(
|
||||
where={"access_group_id": "ag-to-delete"}
|
||||
mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"})
|
||||
|
||||
|
||||
def test_delete_access_group_detaches_group_from_agents(client_and_mocks):
|
||||
"""Delete strips the group from every agent that had it attached, so agents are not left
|
||||
pointing at a group that no longer exists (which would deny them every model, server and agent)."""
|
||||
client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks
|
||||
mock_agents_table = mock_prisma.db.litellm_agentstable
|
||||
|
||||
existing = _make_access_group_record(access_group_id="ag-to-delete")
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
agent_with_group = MagicMock()
|
||||
agent_with_group.agent_id = "agent-1"
|
||||
agent_with_group.access_group_ids = ["ag-keep", "ag-to-delete"]
|
||||
mock_agents_table.find_many = AsyncMock(return_value=[agent_with_group])
|
||||
global_agent_registry.register_agent(
|
||||
AgentResponse(
|
||||
agent_id="agent-1",
|
||||
agent_name="detach-test-agent",
|
||||
agent_card_params={"name": "detach-test-agent", "url": "http://localhost:9", "version": "1"},
|
||||
access_group_ids=["ag-keep", "ag-to-delete"],
|
||||
)
|
||||
)
|
||||
|
||||
try:
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 204
|
||||
|
||||
mock_agents_table.update.assert_awaited_once_with(
|
||||
where={"agent_id": "agent-1"},
|
||||
data={"access_group_ids": ("ag-keep",)},
|
||||
)
|
||||
mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"})
|
||||
registered = global_agent_registry.get_agent_by_id("agent-1")
|
||||
assert registered is not None
|
||||
assert tuple(registered.access_group_ids or ()) == ("ag-keep",)
|
||||
finally:
|
||||
global_agent_registry.deregister_agent("detach-test-agent")
|
||||
|
||||
|
||||
def test_delete_access_group_without_attached_agents_leaves_agents_untouched(client_and_mocks):
|
||||
client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks
|
||||
mock_agents_table = mock_prisma.db.litellm_agentstable
|
||||
|
||||
mock_access_group_table.find_unique = AsyncMock(
|
||||
return_value=_make_access_group_record(access_group_id="ag-to-delete")
|
||||
)
|
||||
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 204
|
||||
|
||||
mock_agents_table.find_many.assert_awaited_once_with(where={"access_group_ids": {"hasSome": ("ag-to-delete",)}})
|
||||
mock_agents_table.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"team_cache_group_ids,key_cache_group_ids,expected_team_ids_after,expected_key_ids_after",
|
||||
|
|
@ -792,9 +817,7 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
"""Delete patches cached team/key objects to remove the deleted access_group_id."""
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj
|
||||
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
|
|
@ -820,13 +843,9 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
team_id="team-1",
|
||||
access_group_ids=list(team_cache_group_ids),
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(
|
||||
return_value=cached_team
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=cached_team)
|
||||
else:
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
||||
# user_api_key_cache is queried both for teams (fallback after dual_cache) and
|
||||
# hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key)
|
||||
|
|
@ -834,9 +853,7 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
# Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects
|
||||
# inconsistently across Python/unittest versions; sync returns are awaited as immediate results.
|
||||
def user_cache_get_side_effect(*args, **kwargs):
|
||||
cache_key = (
|
||||
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
)
|
||||
cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
if cache_key == "team_id:team-1":
|
||||
if team_cache_group_ids is None:
|
||||
return None
|
||||
|
|
@ -868,14 +885,11 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
team_set_calls = [
|
||||
c
|
||||
for c in mock_cache.async_set_cache.call_args_list
|
||||
if c.kwargs.get("key", "") == "team_id:team-1"
|
||||
or (len(c.args) >= 1 and c.args[0] == "team_id:team-1")
|
||||
if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1")
|
||||
]
|
||||
assert len(team_set_calls) >= 1, "Expected team cache to be patched"
|
||||
# The cached team object should have the updated access_group_ids
|
||||
written_team = (
|
||||
team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1]
|
||||
)
|
||||
written_team = team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1]
|
||||
if isinstance(written_team, LiteLLM_TeamTableCachedObj):
|
||||
assert written_team.access_group_ids == expected_team_ids_after
|
||||
else:
|
||||
|
|
@ -883,8 +897,7 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
team_set_calls = [
|
||||
c
|
||||
for c in mock_cache.async_set_cache.call_args_list
|
||||
if c.kwargs.get("key", "") == "team_id:team-1"
|
||||
or (len(c.args) >= 1 and c.args[0] == "team_id:team-1")
|
||||
if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1")
|
||||
]
|
||||
assert len(team_set_calls) == 0, "Should not patch team cache when not cached"
|
||||
|
||||
|
|
@ -892,8 +905,7 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
key_set_calls = [
|
||||
c
|
||||
for c in mock_cache.async_set_cache.call_args_list
|
||||
if c.kwargs.get("key", "") == "hashed-key-1"
|
||||
or (len(c.args) >= 1 and c.args[0] == "hashed-key-1")
|
||||
if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1")
|
||||
]
|
||||
assert len(key_set_calls) >= 1, "Expected key cache to be patched"
|
||||
written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1]
|
||||
|
|
@ -903,17 +915,14 @@ def test_delete_access_group_patches_cached_team_and_key(
|
|||
key_set_calls = [
|
||||
c
|
||||
for c in mock_cache.async_set_cache.call_args_list
|
||||
if c.kwargs.get("key", "") == "hashed-key-1"
|
||||
or (len(c.args) >= 1 and c.args[0] == "hashed-key-1")
|
||||
if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1")
|
||||
]
|
||||
assert len(key_set_calls) == 0, "Should not patch key cache when not cached"
|
||||
|
||||
|
||||
def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
||||
"""Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize)."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
|
|
@ -929,9 +938,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
|||
mock_key_table.find_unique = AsyncMock(return_value=key_with_group)
|
||||
|
||||
# No team in cache
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
||||
# Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that.
|
||||
cached_key_payload = {
|
||||
|
|
@ -940,18 +947,14 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
|||
}
|
||||
|
||||
def user_cache_get_dict_when_key_matches(*args, **kwargs):
|
||||
cache_key = (
|
||||
kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
)
|
||||
cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None)
|
||||
if cache_key == "team_id:team-1":
|
||||
return None
|
||||
if cache_key == "hashed-key-dict":
|
||||
return UserAPIKeyAuth.model_validate(cached_key_payload)
|
||||
return None
|
||||
|
||||
mock_cache.async_get_cache = AsyncMock(
|
||||
side_effect=user_cache_get_dict_when_key_matches
|
||||
)
|
||||
mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_dict_when_key_matches)
|
||||
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 204
|
||||
|
|
@ -960,8 +963,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks):
|
|||
key_set_calls = [
|
||||
c
|
||||
for c in mock_cache.async_set_cache.call_args_list
|
||||
if c.kwargs.get("key", "") == "hashed-key-dict"
|
||||
or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict")
|
||||
if c.kwargs.get("key", "") == "hashed-key-dict" or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict")
|
||||
]
|
||||
assert len(key_set_calls) >= 1, "Expected key cache to be patched"
|
||||
written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1]
|
||||
|
|
@ -988,9 +990,7 @@ def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks):
|
|||
|
||||
existing = _make_access_group_record(access_group_id="ag-to-delete")
|
||||
mock_table.find_unique = AsyncMock(return_value=existing)
|
||||
mock_table.delete = AsyncMock(
|
||||
side_effect=Exception("P2025: Record to delete does not exist")
|
||||
)
|
||||
mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist"))
|
||||
|
||||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 404
|
||||
|
|
@ -1039,9 +1039,7 @@ def test_delete_access_group_500_on_generic_exception(client_and_mocks):
|
|||
("delete", "/v1/unified_access_group/ag-123", lambda: {}),
|
||||
],
|
||||
)
|
||||
def test_access_group_endpoints_db_not_connected(
|
||||
client_and_mocks, monkeypatch, method, url, factory
|
||||
):
|
||||
def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory):
|
||||
"""All endpoints return 500 when DB is not connected."""
|
||||
client, *_ = client_and_mocks
|
||||
|
||||
|
|
@ -1049,9 +1047,7 @@ def test_access_group_endpoints_db_not_connected(
|
|||
|
||||
resp = getattr(client, method)(url, **factory())
|
||||
assert resp.status_code == 500
|
||||
assert (
|
||||
resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value
|
||||
)
|
||||
assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1107,9 +1103,7 @@ def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_t
|
|||
|
||||
def test_create_access_group_syncs_assigned_teams(client_and_mocks):
|
||||
"""Create adds access_group_id to each assigned team's access_group_ids in DB."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
|
||||
team_record = _make_team_record("team-1")
|
||||
|
|
@ -1132,9 +1126,7 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks):
|
|||
|
||||
def test_create_access_group_syncs_assigned_keys(client_and_mocks):
|
||||
"""Create adds access_group_id to each assigned key's access_group_ids in DB."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
key_record = MagicMock()
|
||||
|
|
@ -1148,9 +1140,7 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks):
|
|||
)
|
||||
assert resp.status_code == 201
|
||||
|
||||
mock_key_table.find_unique.assert_awaited_once_with(
|
||||
where={"token": "hashed-token-1"}
|
||||
)
|
||||
mock_key_table.find_unique.assert_awaited_once_with(where={"token": "hashed-token-1"})
|
||||
mock_key_table.update.assert_awaited_once()
|
||||
call_kwargs = mock_key_table.update.call_args.kwargs
|
||||
assert call_kwargs["where"] == {"token": "hashed-token-1"}
|
||||
|
|
@ -1200,14 +1190,10 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks):
|
|||
|
||||
def test_update_access_group_syncs_added_teams(client_and_mocks):
|
||||
"""Update adds access_group_id to newly assigned teams."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", assigned_team_ids=["team-existing"]
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-existing"])
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
team_record = _make_team_record("team-new")
|
||||
|
|
@ -1248,14 +1234,10 @@ def test_update_access_group_rejects_nonexistent_team(client_and_mocks):
|
|||
|
||||
def test_update_access_group_syncs_removed_teams(client_and_mocks):
|
||||
"""Update removes access_group_id from de-assigned teams."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"]
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"])
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
team_to_remove = _make_team_record("team-remove", ["ag-update"])
|
||||
|
|
@ -1268,9 +1250,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
|
|||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
mock_team_table.find_unique.assert_awaited_once_with(
|
||||
where={"team_id": "team-remove"}
|
||||
)
|
||||
mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-remove"})
|
||||
mock_team_table.update.assert_awaited_once()
|
||||
call_kwargs = mock_team_table.update.call_args.kwargs
|
||||
assert call_kwargs["where"] == {"team_id": "team-remove"}
|
||||
|
|
@ -1296,19 +1276,15 @@ def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks):
|
|||
mock_team_table.update.assert_awaited_once()
|
||||
call_kwargs = mock_team_table.update.call_args.kwargs
|
||||
assert call_kwargs["where"] == {"team_id": "team-unmirrored"}
|
||||
assert call_kwargs["data"]["access_group_ids"] == ["ag-other"]
|
||||
assert tuple(call_kwargs["data"]["access_group_ids"]) == ("ag-other",)
|
||||
|
||||
|
||||
def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks):
|
||||
"""Update does not sync teams when assigned_team_ids is absent from the payload."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", assigned_team_ids=["team-1"]
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-1"])
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"})
|
||||
|
|
@ -1320,14 +1296,10 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc
|
|||
|
||||
def test_update_access_group_syncs_added_keys(client_and_mocks):
|
||||
"""Update adds access_group_id to newly assigned keys."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", assigned_key_ids=["old-token"]
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["old-token"])
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
key_record = MagicMock()
|
||||
|
|
@ -1350,14 +1322,10 @@ def test_update_access_group_syncs_added_keys(client_and_mocks):
|
|||
|
||||
def test_update_access_group_syncs_removed_keys(client_and_mocks):
|
||||
"""Update removes access_group_id from de-assigned keys."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
existing = _make_access_group_record(
|
||||
access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"]
|
||||
)
|
||||
existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"])
|
||||
mock_access_group_table.find_unique = AsyncMock(return_value=existing)
|
||||
|
||||
key_to_remove = MagicMock()
|
||||
|
|
@ -1385,9 +1353,7 @@ def test_update_access_group_syncs_removed_keys(client_and_mocks):
|
|||
|
||||
def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks):
|
||||
"""Delete includes teams from assigned_team_ids even when not found by hasSome query."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_team_table = mock_prisma.db.litellm_teamtable
|
||||
|
||||
# Access group has assigned_team_ids but the team's access_group_ids is not synced
|
||||
|
|
@ -1409,18 +1375,14 @@ def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks
|
|||
assert resp.status_code == 204
|
||||
|
||||
# find_unique is called for the out-of-sync team (included via union with assigned_team_ids)
|
||||
mock_team_table.find_unique.assert_awaited_once_with(
|
||||
where={"team_id": "team-out-of-sync"}
|
||||
)
|
||||
mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-out-of-sync"})
|
||||
# No update needed since team's access_group_ids doesn't contain "ag-to-delete"
|
||||
mock_team_table.update.assert_not_awaited()
|
||||
|
||||
|
||||
def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks):
|
||||
"""Delete includes keys from assigned_key_ids even when not found by hasSome query."""
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
|
||||
client_and_mocks
|
||||
)
|
||||
client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks
|
||||
mock_key_table = mock_prisma.db.litellm_verificationtoken
|
||||
|
||||
existing = _make_access_group_record(
|
||||
|
|
@ -1439,9 +1401,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks)
|
|||
resp = client.delete("/v1/access_group/ag-to-delete")
|
||||
assert resp.status_code == 204
|
||||
|
||||
mock_key_table.find_unique.assert_awaited_once_with(
|
||||
where={"token": "token-out-of-sync"}
|
||||
)
|
||||
mock_key_table.find_unique.assert_awaited_once_with(where={"token": "token-out-of-sync"})
|
||||
mock_key_table.update.assert_not_awaited()
|
||||
|
||||
|
||||
|
|
@ -1536,10 +1496,16 @@ def test_list_access_groups_resolves_names_with_one_query_per_table(client_and_m
|
|||
mock_table.find_many = AsyncMock(
|
||||
return_value=[
|
||||
_make_access_group_record(
|
||||
access_group_id="ag-1", access_mcp_server_ids=["mcp-a"], access_agent_ids=["agent-a"], assigned_key_ids=["key-a"]
|
||||
access_group_id="ag-1",
|
||||
access_mcp_server_ids=["mcp-a"],
|
||||
access_agent_ids=["agent-a"],
|
||||
assigned_key_ids=["key-a"],
|
||||
),
|
||||
_make_access_group_record(
|
||||
access_group_id="ag-2", access_mcp_server_ids=["mcp-b"], access_agent_ids=["agent-b"], assigned_key_ids=["key-b"]
|
||||
access_group_id="ag-2",
|
||||
access_mcp_server_ids=["mcp-b"],
|
||||
access_agent_ids=["agent-b"],
|
||||
assigned_key_ids=["key-b"],
|
||||
),
|
||||
]
|
||||
)
|
||||
|
|
@ -1573,7 +1539,10 @@ def test_list_access_groups_skips_lookups_when_nothing_to_resolve(client_and_moc
|
|||
"""Groups with no MCP servers, agents or keys must not trigger an empty IN () query per table."""
|
||||
client, mock_prisma, mock_table, *_ = client_and_mocks
|
||||
mock_table.find_many = AsyncMock(
|
||||
return_value=[_make_access_group_record(access_group_id="ag-1"), _make_access_group_record(access_group_id="ag-2")]
|
||||
return_value=[
|
||||
_make_access_group_record(access_group_id="ag-1"),
|
||||
_make_access_group_record(access_group_id="ag-2"),
|
||||
]
|
||||
)
|
||||
|
||||
resp = client.get("/v1/access_group")
|
||||
|
|
|
|||
|
|
@ -90,6 +90,7 @@ export interface AgentFormValues {
|
|||
guardrails?: string[];
|
||||
entitlement_models?: string[];
|
||||
entitlement_agents?: string[];
|
||||
access_group_ids?: string[];
|
||||
allowed_mcp_servers_and_groups?: McpServerSelection;
|
||||
mcp_tool_permissions?: Record<string, string[]>;
|
||||
defaultInputModes?: string[];
|
||||
|
|
@ -121,6 +122,7 @@ export interface AgentRequestPayload {
|
|||
agent_card_params?: Record<string, unknown>;
|
||||
litellm_params?: Record<string, unknown>;
|
||||
object_permission?: Record<string, unknown>;
|
||||
access_group_ids?: string[];
|
||||
}
|
||||
|
||||
interface AgentFormFieldProps {
|
||||
|
|
|
|||
|
|
@ -21,6 +21,9 @@ vi.mock("./agent_card_discovery", () => ({ default: () => <div data-testid="agen
|
|||
vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ default: () => <div /> }));
|
||||
vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () => <div /> }));
|
||||
vi.mock("@/components/guardrails/GuardrailSelector", () => ({ default: () => <div /> }));
|
||||
vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({
|
||||
useAccessGroups: () => ({ data: [], isLoading: false, isError: false }),
|
||||
}));
|
||||
vi.mock("@/components/common_components/team_dropdown", () => ({ default: () => <div /> }));
|
||||
|
||||
const a2aInfo: AgentCreateInfo = {
|
||||
|
|
|
|||
|
|
@ -44,6 +44,14 @@ vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({
|
|||
default: () => null,
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/AccessGroupSelector", () => ({
|
||||
default: ({ onChange }: { onChange: (value: string[]) => void }) => (
|
||||
<button type="button" data-testid="select-access-group" onClick={() => onChange(["ag-1", "ag-2"])}>
|
||||
Select access group
|
||||
</button>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/team_dropdown", () => ({
|
||||
default: () => null,
|
||||
}));
|
||||
|
|
@ -141,5 +149,27 @@ describe("AddAgentForm logos", () => {
|
|||
await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled());
|
||||
const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0];
|
||||
expect(payload.object_permission).toEqual({ mcp_toolsets: ["ts-1"] });
|
||||
expect(payload).not.toHaveProperty("access_group_ids");
|
||||
});
|
||||
|
||||
it("includes selected access groups in the create payload", async () => {
|
||||
const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never });
|
||||
vi.mocked(networking.createAgentCall).mockReset().mockResolvedValue({
|
||||
agent_id: "agent-1",
|
||||
agent_name: "Test Agent",
|
||||
} as never);
|
||||
vi.mocked(networking.keyListCall).mockResolvedValue({ keys: [] });
|
||||
|
||||
renderForm();
|
||||
await user.click(screen.getByRole("button", { name: "Next →" }));
|
||||
await user.click(screen.getByTestId("select-access-group"));
|
||||
await user.click(screen.getByRole("button", { name: "Next →" }));
|
||||
await user.click(screen.getByRole("button", { name: "Next →" }));
|
||||
await user.click(screen.getByText(/Skip for now/));
|
||||
await user.click(screen.getByRole("button", { name: "Create Agent →" }));
|
||||
|
||||
await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled());
|
||||
const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0];
|
||||
expect(payload.access_group_ids).toEqual(["ag-1", "ag-2"]);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -50,6 +50,7 @@ import {
|
|||
} from "./AgentFormKit";
|
||||
import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector";
|
||||
import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
|
||||
import AccessGroupSelector from "@/components/common_components/AccessGroupSelector";
|
||||
import GuardrailSelector from "@/components/guardrails/GuardrailSelector";
|
||||
import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
|
||||
|
|
@ -113,6 +114,7 @@ const SHARED_INITIAL_VALUES: AgentFormValues = {
|
|||
mcp_tool_permissions: {},
|
||||
entitlement_models: [],
|
||||
entitlement_agents: [],
|
||||
access_group_ids: [],
|
||||
guardrails: [],
|
||||
};
|
||||
|
||||
|
|
@ -374,6 +376,9 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({ visible, onClose, accessTok
|
|||
if (Object.keys(objectPermission).length > 0) {
|
||||
agentData.object_permission = objectPermission;
|
||||
}
|
||||
if (values.access_group_ids?.length) {
|
||||
agentData.access_group_ids = values.access_group_ids;
|
||||
}
|
||||
|
||||
// Wire trace-id flags and budget controls into agent litellm_params (before create call)
|
||||
if (requireTraceIdInbound || requireTraceIdOutbound) {
|
||||
|
|
@ -494,6 +499,22 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({ visible, onClose, accessTok
|
|||
)}
|
||||
</AgentFormField>
|
||||
|
||||
<AgentFormField
|
||||
name="access_group_ids"
|
||||
label={labelWithHint(
|
||||
"Access Groups",
|
||||
"Attach access groups to this agent. Attached groups cap which models, MCP servers, and agents the agent can reach, on top of its key and team permissions. Leave empty to apply no extra cap.",
|
||||
)}
|
||||
>
|
||||
{({ value, onChange }) => (
|
||||
<AccessGroupSelector
|
||||
value={Array.isArray(value) ? (value as string[]) : []}
|
||||
onChange={onChange}
|
||||
placeholder="Select access groups (optional)"
|
||||
/>
|
||||
)}
|
||||
</AgentFormField>
|
||||
|
||||
<Separator className="my-2" />
|
||||
|
||||
<AgentFormField
|
||||
|
|
|
|||
|
|
@ -313,6 +313,10 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => {
|
|||
return agentData;
|
||||
};
|
||||
|
||||
export const parseAccessGroupIdsForForm = (agent: { access_group_ids?: string[] | null }) => ({
|
||||
access_group_ids: agent.access_group_ids ?? [],
|
||||
});
|
||||
|
||||
export const parseMcpPermissionsForForm = (agent: any) => ({
|
||||
allowed_mcp_servers_and_groups: {
|
||||
servers: agent.object_permission?.mcp_servers ?? [],
|
||||
|
|
@ -377,5 +381,6 @@ export const parseAgentForForm = (agent: any) => {
|
|||
// extra_headers: already an array of strings
|
||||
extra_headers: agent.extra_headers ?? [],
|
||||
...parseMcpPermissionsForForm(agent),
|
||||
...parseAccessGroupIdsForForm(agent),
|
||||
};
|
||||
};
|
||||
|
|
|
|||
|
|
@ -25,6 +25,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({
|
|||
|
||||
vi.mock("./agent_card_discovery", () => ({ default: () => <div data-testid="agent-card-discovery" /> }));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({
|
||||
useAccessGroups: () => ({ data: [], isLoading: false, isError: false }),
|
||||
}));
|
||||
|
||||
const A2A_AGENT = {
|
||||
agent_id: "agent-1",
|
||||
agent_name: "my-agent",
|
||||
|
|
@ -176,6 +180,7 @@ describe("AgentInfoView update payload", () => {
|
|||
session_tpm_limit: 333,
|
||||
session_rpm_limit: 444,
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
access_group_ids: [],
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -217,6 +222,7 @@ describe("AgentInfoView update payload", () => {
|
|||
session_tpm_limit: 333,
|
||||
session_rpm_limit: 444,
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
access_group_ids: [],
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -295,6 +301,7 @@ describe("AgentInfoView update payload", () => {
|
|||
model: "langgraph/asst_1",
|
||||
},
|
||||
object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} },
|
||||
access_group_ids: [],
|
||||
});
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -28,6 +28,28 @@ vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({
|
|||
useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }),
|
||||
}));
|
||||
|
||||
vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({
|
||||
useAccessGroups: () => ({
|
||||
data: [{ access_group_id: "ag-1", access_group_name: "support-tools" }],
|
||||
isLoading: false,
|
||||
isError: false,
|
||||
}),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/common_components/AccessGroupSelector", () => ({
|
||||
default: ({ value, onChange }: { value?: string[]; onChange: (value: string[]) => void }) => (
|
||||
<div>
|
||||
<span data-testid="selected-access-groups">{(value ?? []).join(",")}</span>
|
||||
<button type="button" onClick={() => onChange(["ag-1"])}>
|
||||
Attach ag-1
|
||||
</button>
|
||||
<button type="button" onClick={() => onChange([])}>
|
||||
Detach all access groups
|
||||
</button>
|
||||
</div>
|
||||
),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({
|
||||
default: () => <div data-testid="mcp-server-selector" />,
|
||||
}));
|
||||
|
|
@ -76,6 +98,38 @@ describe("AgentInfoView settings", () => {
|
|||
expect(payload.tpm_limit).toBe(42);
|
||||
const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} };
|
||||
expect(payload.object_permission).toEqual(clearedMcpGrants);
|
||||
expect(payload.access_group_ids).toEqual([]);
|
||||
});
|
||||
|
||||
it("sends the newly attached access group in the update payload", async () => {
|
||||
render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
|
||||
|
||||
fireEvent.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Edit Settings" }));
|
||||
fireEvent.click(await screen.findByRole("button", { name: "Attach ag-1" }));
|
||||
expect(screen.getByTestId("selected-access-groups")).toHaveTextContent("ag-1");
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: /Save Changes/ }));
|
||||
|
||||
await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1));
|
||||
const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0];
|
||||
expect(payload.access_group_ids).toEqual(["ag-1"]);
|
||||
});
|
||||
|
||||
it("loads the attached access groups into the editor and sends an empty list once detached", async () => {
|
||||
vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1"] });
|
||||
render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
|
||||
|
||||
fireEvent.click(await screen.findByRole("tab", { name: "Settings" }));
|
||||
fireEvent.click(screen.getByRole("button", { name: "Edit Settings" }));
|
||||
expect(await screen.findByTestId("selected-access-groups")).toHaveTextContent("ag-1");
|
||||
|
||||
fireEvent.click(screen.getByRole("button", { name: "Detach all access groups" }));
|
||||
fireEvent.click(screen.getByRole("button", { name: /Save Changes/ }));
|
||||
|
||||
await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1));
|
||||
const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0];
|
||||
expect(payload.access_group_ids).toEqual([]);
|
||||
});
|
||||
|
||||
it("shows MCP grants with server names on the overview tab", async () => {
|
||||
|
|
@ -88,4 +142,20 @@ describe("AgentInfoView settings", () => {
|
|||
|
||||
expect(await screen.findByText("github (srv-1)")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows attached access groups with their names on the overview tab", async () => {
|
||||
vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1", "ag-unknown"] });
|
||||
|
||||
render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
|
||||
|
||||
expect(await screen.findByText("support-tools (ag-1)")).toBeInTheDocument();
|
||||
expect(screen.getByText("ag-unknown")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows None when the agent has no access groups attached", async () => {
|
||||
render(<AgentInfoView agentId="agent-1" onClose={vi.fn()} accessToken="sk-test" isAdmin={true} />);
|
||||
|
||||
expect(await screen.findByText("Access Groups")).toBeInTheDocument();
|
||||
expect(screen.getByText("None")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ import { Agent } from "@/components/agents/types";
|
|||
import { KeyResponse } from "@/components/key_team_helpers/key_list";
|
||||
import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys";
|
||||
import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers";
|
||||
import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
|
||||
import AccessGroupSelector from "@/components/common_components/AccessGroupSelector";
|
||||
import KeyInfoView from "@/components/templates/key_info_view";
|
||||
import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector";
|
||||
import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions";
|
||||
|
|
@ -26,6 +28,7 @@ import {
|
|||
AGENT_FORM_CONFIG,
|
||||
buildAgentDataFromForm,
|
||||
buildMcpObjectPermission,
|
||||
parseAccessGroupIdsForForm,
|
||||
parseAgentForForm,
|
||||
parseMcpPermissionsForForm,
|
||||
} from "./agent_config";
|
||||
|
|
@ -122,7 +125,11 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
} else {
|
||||
const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) });
|
||||
form.reset({
|
||||
...parseDynamicAgentForForm(data, typeInfo),
|
||||
...parseMcpPermissionsForForm(data),
|
||||
...parseAccessGroupIdsForForm(data),
|
||||
});
|
||||
} else {
|
||||
form.reset(parseAgentForForm(data));
|
||||
}
|
||||
|
|
@ -142,7 +149,11 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
if (agentType !== "a2a") {
|
||||
const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType);
|
||||
if (typeInfo) {
|
||||
form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) });
|
||||
form.reset({
|
||||
...parseDynamicAgentForForm(agent, typeInfo),
|
||||
...parseMcpPermissionsForForm(agent),
|
||||
...parseAccessGroupIdsForForm(agent),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -153,12 +164,18 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" });
|
||||
const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" });
|
||||
const { data: mcpServers = [] } = useMCPServers();
|
||||
const { data: accessGroups = [] } = useAccessGroups();
|
||||
|
||||
const mcpServerLabel = (serverId: string) => {
|
||||
const server = mcpServers.find((s) => s.server_id === serverId);
|
||||
return server?.server_name ? `${server.server_name} (${serverId})` : serverId;
|
||||
};
|
||||
|
||||
const accessGroupLabel = (accessGroupId: string) => {
|
||||
const group = accessGroups.find((g) => g.access_group_id === accessGroupId);
|
||||
return group ? `${group.access_group_name} (${accessGroupId})` : accessGroupId;
|
||||
};
|
||||
|
||||
const discoveryRequest = useMemo(
|
||||
() => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo),
|
||||
[watchedFormValues, selectedAgentTypeInfo, detectedAgentType],
|
||||
|
|
@ -221,6 +238,7 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
await patchAgentCall(accessToken, agentId, {
|
||||
...updateData,
|
||||
object_permission: buildMcpObjectPermission(values),
|
||||
access_group_ids: values.access_group_ids ?? [],
|
||||
});
|
||||
toast.success("Agent updated successfully");
|
||||
setIsEditing(false);
|
||||
|
|
@ -350,6 +368,17 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
<DetailItem label="RPM Limit">{agent.rpm_limit ?? "Unlimited"}</DetailItem>
|
||||
<DetailItem label="Session TPM Limit">{agent.session_tpm_limit ?? "Unlimited"}</DetailItem>
|
||||
<DetailItem label="Session RPM Limit">{agent.session_rpm_limit ?? "Unlimited"}</DetailItem>
|
||||
<DetailItem label="Access Groups">
|
||||
{agent.access_group_ids?.length ? (
|
||||
<div className="space-y-1">
|
||||
{agent.access_group_ids.map((accessGroupId) => (
|
||||
<div key={accessGroupId}>{accessGroupLabel(accessGroupId)}</div>
|
||||
))}
|
||||
</div>
|
||||
) : (
|
||||
"None"
|
||||
)}
|
||||
</DetailItem>
|
||||
<DetailItem label="Created At">{formatDate(agent.created_at)}</DetailItem>
|
||||
<DetailItem label="Updated At">{formatDate(agent.updated_at)}</DetailItem>
|
||||
</DetailList>
|
||||
|
|
@ -489,6 +518,26 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({ agentId, onClose, accessT
|
|||
{rateLimitField("session_rpm_limit", "Session RPM Limit")}
|
||||
</div>
|
||||
|
||||
<Separator className="my-6" />
|
||||
<h3 className="text-lg font-medium mb-4">Access Groups</h3>
|
||||
<FieldGroup>
|
||||
<AgentFormField
|
||||
name="access_group_ids"
|
||||
label={labelWithHint(
|
||||
"Access Groups",
|
||||
"Attached groups cap which models, MCP servers, and agents this agent can reach, on top of its key and team permissions. Leave empty to apply no extra cap.",
|
||||
)}
|
||||
>
|
||||
{({ value, onChange }) => (
|
||||
<AccessGroupSelector
|
||||
value={Array.isArray(value) ? (value as string[]) : []}
|
||||
onChange={onChange}
|
||||
placeholder="Select access groups (optional)"
|
||||
/>
|
||||
)}
|
||||
</AgentFormField>
|
||||
</FieldGroup>
|
||||
|
||||
<Separator className="my-6" />
|
||||
<h3 className="text-lg font-medium mb-4">MCP Servers</h3>
|
||||
<FieldGroup>
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ export interface Agent {
|
|||
[key: string]: any;
|
||||
};
|
||||
object_permission?: AgentObjectPermission;
|
||||
access_group_ids?: string[] | null;
|
||||
keys?: AgentAttachedKey[] | null;
|
||||
spend?: number;
|
||||
tpm_limit?: number | null;
|
||||
|
|
|
|||
|
|
@ -6267,6 +6267,7 @@ export const patchAgentCall = async (
|
|||
rpm_limit?: number | null;
|
||||
session_tpm_limit?: number | null;
|
||||
session_rpm_limit?: number | null;
|
||||
access_group_ids?: string[];
|
||||
},
|
||||
) => {
|
||||
try {
|
||||
|
|
|
|||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -23195,6 +23195,8 @@ export interface components {
|
|||
};
|
||||
/** AgentConfig */
|
||||
AgentConfig: {
|
||||
/** Access Group Ids */
|
||||
access_group_ids?: string[] | null;
|
||||
agent_card_params: components["schemas"]["AgentCard"];
|
||||
/** Agent Name */
|
||||
agent_name: string;
|
||||
|
|
@ -23342,6 +23344,8 @@ export interface components {
|
|||
};
|
||||
/** AgentResponse */
|
||||
AgentResponse: {
|
||||
/** Access Group Ids */
|
||||
access_group_ids?: string[] | null;
|
||||
/** Agent Card Params */
|
||||
agent_card_params: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -34111,6 +34115,8 @@ export interface components {
|
|||
};
|
||||
/** PatchAgentRequest */
|
||||
PatchAgentRequest: {
|
||||
/** Access Group Ids */
|
||||
access_group_ids?: string[] | null;
|
||||
agent_card_params?: components["schemas"]["AgentCard"];
|
||||
/** Agent Name */
|
||||
agent_name?: string;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue