mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #41634 from BerriAI/litellm_agent_access_groups
feat(agents): attach access groups to agents and enforce them for models, MCP servers and agent calls
This commit is contained in:
Yassin Kortam
GitHub
commit
9cc5b78c33
37 changed files with 1851 additions and 195 deletions
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];
|
||||
|
|
@ -73,6 +73,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)
|
||||
|
|
|
|||
|
|
@ -44,6 +44,11 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
user_api_key_has_admin_view,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
CeilingResolver,
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
from litellm.proxy.auth.user_api_key_auth import (
|
||||
_get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth
|
||||
|
|
@ -184,6 +189,21 @@ def _has_client_supplied_mcp_auth(
|
|||
return bool(mcp_auth_header) or bool(mcp_server_auth_headers)
|
||||
|
||||
|
||||
def _agent_capped_servers(
|
||||
allowed_mcp_servers: Sequence[str],
|
||||
agent_servers: Sequence[str],
|
||||
agent_access_group_servers: frozenset[str] | None,
|
||||
) -> tuple[str, ...] | None:
|
||||
if not agent_servers and agent_access_group_servers is None:
|
||||
return None
|
||||
return tuple(
|
||||
s
|
||||
for s in allowed_mcp_servers
|
||||
if (not agent_servers or s in agent_servers)
|
||||
and (agent_access_group_servers is None or s in agent_access_group_servers)
|
||||
)
|
||||
|
||||
|
||||
def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
|
||||
"""True when this auth is a keyless subject admitted by the gateway session / bridge user
|
||||
path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``.
|
||||
|
|
@ -1546,25 +1566,33 @@ class MCPRequestHandler:
|
|||
# Check agent permissions if agent_id is set on the key
|
||||
#########################################################
|
||||
if user_api_key_auth and user_api_key_auth.agent_id:
|
||||
allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(
|
||||
user_api_key_auth
|
||||
agent_capped: Final = _agent_capped_servers(
|
||||
allowed_mcp_servers,
|
||||
await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth),
|
||||
await MCPRequestHandler._get_agent_access_group_server_ceiling(user_api_key_auth),
|
||||
)
|
||||
if len(allowed_mcp_servers_for_agent) > 0:
|
||||
if agent_capped 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]
|
||||
allowed_mcp_servers = list(agent_capped)
|
||||
verbose_logger.debug(
|
||||
"Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Cap an agent key at what the user and team that invoked the agent may reach
|
||||
#########################################################
|
||||
caller_capped, caller_restricts = await MCPRequestHandler._apply_agent_caller_ceiling(
|
||||
allowed_mcp_servers, user_api_key_auth
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Apply the internal user's own ceiling (the entitlement attached to the human)
|
||||
#########################################################
|
||||
capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(
|
||||
allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source
|
||||
caller_capped, user_api_key_auth, keyless_source=keyless_source
|
||||
)
|
||||
allowed_mcp_servers = list(capped)
|
||||
has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts
|
||||
has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or caller_restricts or user_restricts
|
||||
|
||||
#########################################################
|
||||
# Apply org-level ceiling if org_id is set
|
||||
|
|
@ -2907,6 +2935,28 @@ class MCPRequestHandler:
|
|||
verbose_logger.debug("Applied user ceiling filter. Final allowed servers: %s", capped)
|
||||
return capped, True
|
||||
|
||||
@staticmethod
|
||||
async def _apply_agent_caller_ceiling(
|
||||
allowed_mcp_servers: Sequence[str],
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
) -> tuple[tuple[str, ...], bool]:
|
||||
"""Narrow an agent key's servers to those the invoking user and team (echoed back by the agent
|
||||
as ``x-litellm-user-id`` / ``x-litellm-team-id``) may reach: the echoed team's grants when it
|
||||
names any, then the echoed user's own entitlement. Raises like the user ceiling when that
|
||||
entitlement is known but unreadable, so the resolver denies rather than widens."""
|
||||
caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None
|
||||
if caller_auth is None:
|
||||
return tuple(allowed_mcp_servers), False
|
||||
team_servers: Final = frozenset(await MCPRequestHandler._get_allowed_mcp_servers_for_team(caller_auth))
|
||||
team_capped: Final = (
|
||||
tuple(server for server in allowed_mcp_servers if server in team_servers)
|
||||
if team_servers
|
||||
else tuple(allowed_mcp_servers)
|
||||
)
|
||||
user_capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(team_capped, caller_auth)
|
||||
verbose_logger.debug("Applied agent caller ceiling. Final allowed servers: %s", user_capped)
|
||||
return user_capped, bool(team_servers) or user_restricts
|
||||
|
||||
@staticmethod
|
||||
async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool:
|
||||
"""Whether this human's own entitlement bounds their MCP access at all.
|
||||
|
|
@ -3137,6 +3187,27 @@ 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,
|
||||
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
|
||||
) -> 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,
|
||||
)
|
||||
|
||||
if not user_api_key_auth.agent_id:
|
||||
return None
|
||||
ceiling: Final = await resolve_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",
|
||||
|
|
@ -3506,6 +3534,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"
|
||||
},
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
|||
validate_langfuse_span_scope_value,
|
||||
validate_no_callback_env_reference,
|
||||
)
|
||||
from litellm.types.agents import AgentCaller
|
||||
from litellm.types.integrations.compression_interception import (
|
||||
CompressionSavingsMetadata,
|
||||
)
|
||||
|
|
@ -3274,6 +3275,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
"user id."
|
||||
),
|
||||
)
|
||||
agent_caller: AgentCaller | None = Field(
|
||||
default=None,
|
||||
exclude=True,
|
||||
description=(
|
||||
"Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on "
|
||||
"calls made with its own key. Every check treats it as a ceiling, so a forged value can only "
|
||||
"narrow the agent's access."
|
||||
),
|
||||
)
|
||||
budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True)
|
||||
team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True)
|
||||
|
|
@ -3306,6 +3316,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
values.pop("mcp_session_resource_server_id", None)
|
||||
values.pop("mcp_toolset_id", None)
|
||||
values.pop("via_virtual_key", None)
|
||||
values.pop("agent_caller", None)
|
||||
if values.get("api_key") is not None:
|
||||
values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))})
|
||||
if isinstance(values.get("api_key"), str):
|
||||
|
|
@ -4261,6 +4272,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"
|
||||
|
|
@ -4335,7 +4351,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
|
||||
|
|
@ -4350,6 +4366,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(
|
||||
|
|
|
|||
|
|
@ -146,12 +146,17 @@ def _validate_push_notification_url(url: str) -> None:
|
|||
|
||||
|
||||
def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, str]:
|
||||
"""The human behind this call. An agent key acting for an invoking user forwards that user, not
|
||||
itself, so a chain of agents stays capped at what the original caller may reach."""
|
||||
caller: Final = user_api_key_dict.agent_caller
|
||||
user_id: Final = caller.user_id if caller is not None else user_api_key_dict.user_id
|
||||
team_id: Final = caller.team_id if caller is not None else user_api_key_dict.team_id
|
||||
return MappingProxyType(
|
||||
{
|
||||
name: value
|
||||
for name, value in (
|
||||
("X-LiteLLM-User-Id", user_api_key_dict.user_id),
|
||||
("X-LiteLLM-Team-Id", user_api_key_dict.team_id),
|
||||
("X-LiteLLM-User-Id", user_id),
|
||||
("X-LiteLLM-Team-Id", team_id),
|
||||
)
|
||||
if value
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -65,6 +67,9 @@ class AgentRecord(Protocol):
|
|||
@property
|
||||
def object_permission(self) -> AgentObjectPermissionRecord | None: ...
|
||||
|
||||
@property
|
||||
def access_group_ids(self) -> Sequence[str] | None: ...
|
||||
|
||||
@property
|
||||
def spend(self) -> float: ...
|
||||
|
||||
|
|
@ -284,6 +289,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 +527,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 +544,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 +615,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 +717,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 +725,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),
|
||||
}
|
||||
|
|
|
|||
75
litellm/proxy/agent_endpoints/auth/agent_access_groups.py
Normal file
75
litellm/proxy/agent_endpoints/auth/agent_access_groups.py
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
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
|
||||
|
||||
AccessGroupIds: TypeAlias = tuple[str, ...]
|
||||
AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params
|
||||
LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None
|
||||
AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax
|
||||
|
||||
|
||||
@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: AccessGroupIds
|
||||
models: frozenset[str]
|
||||
mcp_server_ids: frozenset[str]
|
||||
agent_ids: frozenset[str]
|
||||
|
||||
|
||||
CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params
|
||||
|
||||
|
||||
async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds:
|
||||
from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
|
||||
|
||||
agent: Final = await get_agent_with_read_through(agent_id)
|
||||
return tuple(agent.access_group_ids or ()) if agent is not None else ()
|
||||
|
||||
|
||||
async def _load_access_group(access_group_id: str) -> LoadedAccessGroup:
|
||||
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_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids,
|
||||
load_access_group: AccessGroupLoader = _load_access_group,
|
||||
) -> AgentAccessGroupCeiling | None:
|
||||
"""``None`` when the agent has no access groups attached, so nothing is capped."""
|
||||
access_group_ids: Final = await load_access_group_ids(agent_id)
|
||||
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),
|
||||
)
|
||||
87
litellm/proxy/agent_endpoints/auth/agent_caller.py
Normal file
87
litellm/proxy/agent_endpoints/auth/agent_caller.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""The human behind an agent's own proxy calls.
|
||||
|
||||
``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the
|
||||
agent backend. When the agent echoes them back on requests made with its own key, the proxy caps
|
||||
that key at what the invoking user and team may reach. The cap is intersected with, never
|
||||
substituted for, the agent key's own grants and the agent's access group ceiling, so the headers
|
||||
can only narrow access and need no trust.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth
|
||||
from litellm.types.agents import (
|
||||
AGENT_CALLER_TEAM_ID_HEADER,
|
||||
AGENT_CALLER_USER_ID_HEADER,
|
||||
AgentCaller,
|
||||
)
|
||||
|
||||
|
||||
def _header(headers: Mapping[str, str], name: str) -> str | None:
|
||||
value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None)
|
||||
return value.strip() or None if value is not None else None
|
||||
|
||||
|
||||
def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None:
|
||||
"""The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed."""
|
||||
if not user_api_key_auth.agent_id:
|
||||
return None
|
||||
user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER)
|
||||
team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER)
|
||||
if user_id is None and team_id is None:
|
||||
return None
|
||||
return AgentCaller(user_id=user_id, team_id=team_id)
|
||||
|
||||
|
||||
def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None:
|
||||
"""A minimal auth context standing for the invoking user and team, so the shared key/team/user
|
||||
resolvers can be reused unchanged to compute what the caller may reach."""
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None:
|
||||
return None
|
||||
return UserAPIKeyAuth(
|
||||
user_id=caller.user_id,
|
||||
team_id=caller.team_id,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
)
|
||||
|
||||
|
||||
async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
|
||||
"""The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team
|
||||
that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted."""
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None or caller.team_id is None:
|
||||
return None
|
||||
return await get_team_object(
|
||||
team_id=caller.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None:
|
||||
"""The invoking user's row, or ``None`` when no user id was echoed or the row does not exist."""
|
||||
from litellm.proxy.auth.auth_checks import get_user_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None or caller.user_id is None:
|
||||
return None
|
||||
user_object: Final = await get_user_object(
|
||||
user_id=caller.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if user_object is None:
|
||||
verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id)
|
||||
return user_object
|
||||
|
|
@ -19,6 +19,11 @@ from litellm.proxy._types import (
|
|||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
CeilingResolver,
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
|
||||
from litellm.repositories.table_repositories import AgentsRepository
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
|
|
@ -44,6 +49,22 @@ def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]:
|
|||
return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids)
|
||||
|
||||
|
||||
def _restricted_ids(access: AgentAccess) -> frozenset[str] | None:
|
||||
if isinstance(access, UnrestrictedAgentAccess):
|
||||
return None
|
||||
return _to_stable_ids(access.agent_ids)
|
||||
|
||||
|
||||
def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess:
|
||||
key_ids: Final = _restricted_ids(key_access)
|
||||
team_ids: Final = _restricted_ids(team_access)
|
||||
if key_ids is None:
|
||||
return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids)
|
||||
if team_ids is None:
|
||||
return RestrictedAgentAccess(key_ids)
|
||||
return RestrictedAgentAccess(key_ids & team_ids)
|
||||
|
||||
|
||||
class AgentRequestHandler:
|
||||
"""
|
||||
Class to handle agent permission checking, including:
|
||||
|
|
@ -61,35 +82,56 @@ class AgentRequestHandler:
|
|||
@staticmethod
|
||||
async def resolve_agent_access(
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
|
||||
) -> AgentAccess:
|
||||
"""
|
||||
Resolve the agents the given user/key may reach.
|
||||
"""Agents the key may reach: key and team grants, intersected with the agent's access group ceiling
|
||||
and, for an agent key acting on behalf of an invoking user, with that user's team grants."""
|
||||
key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth)
|
||||
caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth)
|
||||
own_access: Final = _intersect_agent_access(key_team_access, caller_access)
|
||||
agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling)
|
||||
if agent_ceiling is None:
|
||||
return own_access
|
||||
if isinstance(own_access, UnrestrictedAgentAccess):
|
||||
return RestrictedAgentAccess(agent_ceiling)
|
||||
return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling)
|
||||
|
||||
``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.
|
||||
"""
|
||||
@staticmethod
|
||||
async def _agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None) -> AgentAccess:
|
||||
caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None
|
||||
if caller_auth is None:
|
||||
return UnrestrictedAgentAccess()
|
||||
return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth)
|
||||
|
||||
@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)
|
||||
|
||||
match (key_access, team_access):
|
||||
case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()):
|
||||
return UnrestrictedAgentAccess()
|
||||
case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)):
|
||||
return RestrictedAgentAccess(_to_stable_ids(team_ids))
|
||||
case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()):
|
||||
return RestrictedAgentAccess(_to_stable_ids(key_ids))
|
||||
case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)):
|
||||
return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids))
|
||||
except Exception as e:
|
||||
verbose_logger.warning("Failed to get allowed agents: %s", e)
|
||||
return UnrestrictedAgentAccess()
|
||||
return _intersect_agent_access(key_access, team_access)
|
||||
|
||||
@staticmethod
|
||||
async def _agent_access_group_ceiling(
|
||||
user_api_key_auth: UserAPIKeyAuth | None,
|
||||
resolve_ceiling: CeilingResolver,
|
||||
) -> frozenset[str] | None:
|
||||
if user_api_key_auth is None or not user_api_key_auth.agent_id:
|
||||
return None
|
||||
ceiling: Final = await resolve_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,
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a specific agent is allowed for the given user/key.
|
||||
|
|
@ -103,7 +145,7 @@ class AgentRequestHandler:
|
|||
"""
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
match await AgentRequestHandler.resolve_agent_access(user_api_key_auth):
|
||||
match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling):
|
||||
case UnrestrictedAgentAccess():
|
||||
return True
|
||||
case RestrictedAgentAccess(allowed_agent_ids):
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ import re
|
|||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
|
@ -68,6 +68,15 @@ from litellm.proxy._types import (
|
|||
SpecialModelNames,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
|
||||
CeilingResolver,
|
||||
resolve_agent_access_group_ceiling,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import (
|
||||
agent_caller_auth,
|
||||
load_agent_caller_team,
|
||||
load_agent_caller_user,
|
||||
)
|
||||
from litellm.proxy.auth.budget_throttle import (
|
||||
budget_throttle_percentage,
|
||||
should_throttle_budget_exceeded,
|
||||
|
|
@ -1006,6 +1015,16 @@ async def common_checks(
|
|||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router)
|
||||
await _check_agent_caller_model_access(
|
||||
model=_model,
|
||||
valid_token=valid_token,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
## 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"):
|
||||
|
|
@ -4251,7 +4270,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]:
|
||||
"""
|
||||
|
|
@ -4317,6 +4336,82 @@ 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,
|
||||
resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
|
||||
) -> Literal[True]:
|
||||
"""Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows."""
|
||||
if not model or valid_token is None or not valid_token.agent_id:
|
||||
return True
|
||||
ceiling: Final = await resolve_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",
|
||||
)
|
||||
|
||||
|
||||
LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None
|
||||
LoadedCallerUser: TypeAlias = LiteLLM_UserTable | None
|
||||
CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerTeam]] # mutable-ok: Callable params
|
||||
CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] # mutable-ok: Callable params
|
||||
|
||||
|
||||
async def _check_agent_caller_model_access(
|
||||
model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str]
|
||||
valid_token: UserAPIKeyAuth | None,
|
||||
llm_router: Router | None,
|
||||
prisma_client: Optional["PrismaClient"],
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
load_team: CallerTeamLoader = load_agent_caller_team,
|
||||
load_user: CallerUserLoader = load_agent_caller_user,
|
||||
) -> None:
|
||||
"""An agent key acting for an invoking user may call only what that user's own key could: the
|
||||
invoking team's models (and per-member scope) when a team was echoed, else the user's models."""
|
||||
if not model or valid_token is None:
|
||||
return
|
||||
caller_auth: Final = agent_caller_auth(valid_token)
|
||||
if caller_auth is None:
|
||||
return
|
||||
caller_team: Final = await load_team(valid_token)
|
||||
if caller_team is not None:
|
||||
await can_team_access_model(
|
||||
model=model,
|
||||
team_object=caller_team,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
await _check_team_member_model_access(
|
||||
model=model,
|
||||
team_object=caller_team,
|
||||
valid_token=caller_auth,
|
||||
llm_router=llm_router,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return
|
||||
caller_user: Final = await load_user(valid_token)
|
||||
if caller_user is None:
|
||||
return
|
||||
await can_user_call_model(model=model, llm_router=llm_router, user_object=caller_user)
|
||||
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ from litellm.integrations.otel.runtime import phase_span, seed_request_identity
|
|||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
ExperimentalUIJWTToken,
|
||||
TeamNotFoundError,
|
||||
|
|
@ -3330,6 +3331,9 @@ async def user_api_key_auth(
|
|||
raise body_parse_exception
|
||||
raise
|
||||
user_api_key_auth_obj.budget_reservation = None
|
||||
user_api_key_auth_obj.agent_caller = agent_caller_from_headers(
|
||||
_safe_get_request_headers(request), user_api_key_auth_obj
|
||||
)
|
||||
_seed_request_destinations(user_api_key_auth_obj, request)
|
||||
|
||||
# A body that never parsed is authenticated (so the trace carries identity
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -73,6 +73,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)
|
||||
|
|
|
|||
|
|
@ -148,6 +148,7 @@ from litellm.proxy._types import (
|
|||
Member,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
|
||||
|
|
@ -8261,6 +8262,51 @@ async def _get_access_group_models(
|
|||
return tuple(dict.fromkeys((*team_group_models, *key_group_models)))
|
||||
|
||||
|
||||
async def _agent_access_group_visible_models(
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
llm_router: "Router | None",
|
||||
include_model_access_groups: bool,
|
||||
return_wildcard_routes: bool,
|
||||
team_id: str | None,
|
||||
resolve_agent_ceiling: CeilingResolver,
|
||||
) -> frozenset[str] | None:
|
||||
"""Models an agent key may still list once its attached access groups cap it, ``None`` when
|
||||
nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on."""
|
||||
from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models
|
||||
|
||||
if not user_api_key_dict.agent_id:
|
||||
return None
|
||||
ceiling: Final = await resolve_agent_ceiling(user_api_key_dict.agent_id)
|
||||
if ceiling is None:
|
||||
return None
|
||||
if llm_router is None:
|
||||
return ceiling.models
|
||||
proxy_model_list: Final = llm_router.get_model_names()
|
||||
model_access_groups: Final = llm_router.get_model_access_groups()
|
||||
granted: Final = get_team_models(
|
||||
team_models=sorted(ceiling.models),
|
||||
proxy_model_list=proxy_model_list,
|
||||
model_access_groups=model_access_groups,
|
||||
include_model_access_groups=include_model_access_groups,
|
||||
)
|
||||
if not granted:
|
||||
return frozenset()
|
||||
return frozenset(
|
||||
get_complete_model_list(
|
||||
key_models=granted,
|
||||
team_models=(),
|
||||
proxy_model_list=proxy_model_list,
|
||||
user_model=None,
|
||||
infer_model_from_keys=False,
|
||||
return_wildcard_routes=return_wildcard_routes,
|
||||
llm_router=llm_router,
|
||||
model_access_groups=model_access_groups,
|
||||
include_model_access_groups=include_model_access_groups,
|
||||
team_id=team_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def get_available_models_for_user(
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
llm_router: Optional["Router"],
|
||||
|
|
@ -8273,6 +8319,7 @@ async def get_available_models_for_user(
|
|||
only_model_access_groups: bool = False,
|
||||
return_wildcard_routes: bool = False,
|
||||
user_api_key_cache: Optional["UserApiKeyCache"] = None,
|
||||
resolve_agent_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Get the list of models available to a user based on their API key and team permissions.
|
||||
|
|
@ -8376,7 +8423,18 @@ async def get_available_models_for_user(
|
|||
team_id=effective_team_id,
|
||||
)
|
||||
|
||||
return all_models
|
||||
agent_visible: Final = await _agent_access_group_visible_models(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
include_model_access_groups=include_model_access_groups,
|
||||
return_wildcard_routes=return_wildcard_routes,
|
||||
team_id=effective_team_id,
|
||||
resolve_agent_ceiling=resolve_agent_ceiling,
|
||||
)
|
||||
if agent_visible is None:
|
||||
return all_models
|
||||
capped: Final = [m for m in all_models if m in agent_visible] # mutable-ok: callers expect the list all_models is
|
||||
return capped
|
||||
|
||||
|
||||
def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None:
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, PrivateAttr, StrictInt
|
||||
from pydantic import BaseModel, ConfigDict, PrivateAttr, StrictInt
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
|
|
@ -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,21 @@ 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]
|
||||
|
||||
|
||||
AGENT_CALLER_USER_ID_HEADER: Final = "x-litellm-user-id"
|
||||
AGENT_CALLER_TEAM_ID_HEADER: Final = "x-litellm-team-id"
|
||||
|
||||
|
||||
class AgentCaller(BaseModel):
|
||||
"""The user and team that invoked an agent, echoed back by the agent on its own proxy calls.
|
||||
Only ever narrows what the agent's key may do."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
user_id: str | None = None
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
# Request/Response models for CRUD endpoints
|
||||
|
|
@ -226,6 +242,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
|
||||
|
|
|
|||
|
|
@ -73,6 +73,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)
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from starlette.datastructures import Headers
|
|||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
UnloadableEntitlementError,
|
||||
_agent_capped_servers,
|
||||
_is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.proxy._types import (
|
|||
SpecialMCPServerNames,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.types.agents import AgentCaller
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4169,10 +4171,114 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission():
|
|||
global_mcp_server_manager.registry.pop("direct-server", None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("agent_servers", "group_ceiling", "expected"),
|
||||
[
|
||||
([], frozenset({"server_1"}), ("server_1",)),
|
||||
([], frozenset({"server_1", "server_2", "server_3"}), ("server_1", "server_2")),
|
||||
([], frozenset(), ()),
|
||||
(["server_2"], frozenset({"server_1", "server_2"}), ("server_2",)),
|
||||
(["server_1"], frozenset({"server_2"}), ()),
|
||||
(["server_1"], None, ("server_1",)),
|
||||
],
|
||||
)
|
||||
def test_agent_capped_servers_intersects_agent_config_and_access_groups(agent_servers, group_ceiling, expected):
|
||||
"""The agent's attached access groups cap the key/team servers alongside its own
|
||||
object_permission; groups naming no server deny all."""
|
||||
assert _agent_capped_servers(["server_1", "server_2"], agent_servers, group_ceiling) == expected
|
||||
|
||||
|
||||
def test_agent_capped_servers_without_agent_restrictions_is_uncapped():
|
||||
assert _agent_capped_servers(["server_1", "server_2"], [], None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestAgentMCPPermissions:
|
||||
"""Test agent-level MCP server and tool permission intersection."""
|
||||
|
||||
@staticmethod
|
||||
def _agent_key_acting_for(user_id: str, team_id: str | None) -> UserAPIKeyAuth:
|
||||
agent_key = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1")
|
||||
agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id)
|
||||
return agent_key
|
||||
|
||||
@staticmethod
|
||||
def _team_servers(grants: dict[str, list[str]]) -> AsyncMock:
|
||||
async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str]:
|
||||
assert user_api_key_auth is not None
|
||||
return grants.get(user_api_key_auth.team_id or "", [])
|
||||
|
||||
return AsyncMock(side_effect=by_team)
|
||||
|
||||
@staticmethod
|
||||
def _user_servers(grants: dict[str, list[str] | None]) -> AsyncMock:
|
||||
async def by_user(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str] | None:
|
||||
assert user_api_key_auth is not None
|
||||
return grants.get(user_api_key_auth.user_id or "", [])
|
||||
|
||||
return AsyncMock(side_effect=by_user)
|
||||
|
||||
async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_servers(self):
|
||||
"""LIT-8014: the agent's own key reaches server_1 and server_2, but the human who invoked it
|
||||
belongs to a team granted only server_2, so on their behalf the agent reaches only server_2."""
|
||||
agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers")
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"])
|
||||
),
|
||||
patch.object( # test-quality-ok: same seam, keyed by which team is being asked about
|
||||
MCPRequestHandler,
|
||||
"_get_allowed_mcp_servers_for_team",
|
||||
self._team_servers({"callers": ["server_2", "server_3"]}),
|
||||
),
|
||||
patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[])
|
||||
),
|
||||
patch.object( # test-quality-ok: neither the agent's owner nor the caller has a personal grant
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({})
|
||||
),
|
||||
):
|
||||
assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_2"]
|
||||
|
||||
async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_servers(self):
|
||||
agent_key = self._agent_key_acting_for(user_id="alice", team_id=None)
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"])
|
||||
),
|
||||
patch.object( # test-quality-ok: same seam
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({})
|
||||
),
|
||||
patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[])
|
||||
),
|
||||
patch.object( # test-quality-ok: same seam, keyed by which user is being asked about
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": ["server_1"]})
|
||||
),
|
||||
):
|
||||
assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_1"]
|
||||
|
||||
async def test_agent_key_acting_for_a_caller_whose_entitlement_is_unreadable_reaches_nothing(self):
|
||||
agent_key = self._agent_key_acting_for(user_id="alice", team_id=None)
|
||||
|
||||
with (
|
||||
patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1"])
|
||||
),
|
||||
patch.object( # test-quality-ok: same seam
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({})
|
||||
),
|
||||
patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[])
|
||||
),
|
||||
patch.object( # test-quality-ok: None is the resolver's own "entitlement unresolvable" signal
|
||||
MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": None})
|
||||
),
|
||||
):
|
||||
assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == []
|
||||
|
||||
async def test_get_allowed_mcp_servers_agent_intersection(self):
|
||||
"""Key/team allow [server_1, server_2]; agent allows [server_1]. Result = [server_1]."""
|
||||
user_api_key_auth = UserAPIKeyAuth(
|
||||
|
|
@ -4208,6 +4314,46 @@ class TestAgentMCPPermissions:
|
|||
assert sorted(result) == ["server_1", "server_2"]
|
||||
mock_agent.assert_called_once_with(user_api_key_auth)
|
||||
|
||||
async def test_agent_access_group_server_ceiling_expands_group_servers(self):
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
|
||||
from litellm.types.mcp import MCPTransport
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
asked: list[str] = []
|
||||
|
||||
async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
|
||||
asked.append(agent_id)
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",),
|
||||
models=frozenset(),
|
||||
mcp_server_ids=frozenset({"aliased-server"}),
|
||||
agent_ids=frozenset(),
|
||||
)
|
||||
|
||||
global_mcp_server_manager.registry["ag-server-id"] = MCPServer(
|
||||
server_id="ag-server-id",
|
||||
name="ag-server",
|
||||
server_name="ag-server",
|
||||
alias="aliased-server",
|
||||
url="https://ag-server.example.com",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
try:
|
||||
result = await MCPRequestHandler._get_agent_access_group_server_ceiling(
|
||||
UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag"), resolve
|
||||
)
|
||||
finally:
|
||||
global_mcp_server_manager.registry.pop("ag-server-id", None)
|
||||
|
||||
assert result == frozenset({"ag-server-id"})
|
||||
assert asked == ["agent-ag"]
|
||||
assert (
|
||||
await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k"), resolve)
|
||||
is None
|
||||
)
|
||||
assert asked == ["agent-ag"]
|
||||
|
||||
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,146 @@
|
|||
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) -> tuple[str, ...]:
|
||||
return tuple(agent.access_group_ids or ()) if agent is not None else ()
|
||||
|
||||
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_agent_loader_reads_the_attached_groups_from_the_registry():
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
_, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))})
|
||||
global_agent_registry.register_agent(_agent(["g1"]))
|
||||
try:
|
||||
ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group)
|
||||
finally:
|
||||
global_agent_registry.deregister_agent("agent")
|
||||
|
||||
assert ceiling == AgentAccessGroupCeiling(
|
||||
access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), 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
|
||||
|
|
@ -0,0 +1,57 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth, agent_caller_from_headers
|
||||
from litellm.types.agents import AgentCaller
|
||||
|
||||
_AGENT_KEY: Final = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1")
|
||||
|
||||
|
||||
def test_agent_key_echoing_both_ids_acts_for_that_user_and_team() -> None:
|
||||
headers: Final = {"X-LiteLLM-User-Id": " alice ", "x-litellm-team-id": "callers"}
|
||||
|
||||
assert agent_caller_from_headers(headers, _AGENT_KEY) == AgentCaller(user_id="alice", team_id="callers")
|
||||
|
||||
|
||||
def test_agent_key_echoing_only_a_user_id_acts_for_a_teamless_user() -> None:
|
||||
assert agent_caller_from_headers({"x-litellm-user-id": "alice"}, _AGENT_KEY) == AgentCaller(user_id="alice")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers", [{}, {"x-litellm-user-id": " ", "x-litellm-team-id": ""}])
|
||||
def test_agent_key_echoing_no_caller_acts_for_itself(headers: dict[str, str]) -> None:
|
||||
assert agent_caller_from_headers(headers, _AGENT_KEY) is None
|
||||
|
||||
|
||||
def test_caller_headers_on_a_key_without_an_agent_are_ignored() -> None:
|
||||
plain_key: Final = UserAPIKeyAuth(api_key="plain-key", user_id="bob")
|
||||
|
||||
assert agent_caller_from_headers({"x-litellm-user-id": "alice", "x-litellm-team-id": "callers"}, plain_key) is None
|
||||
|
||||
|
||||
def test_caller_auth_stands_for_the_invoking_user_not_the_agent() -> None:
|
||||
agent_key: Final = UserAPIKeyAuth(
|
||||
api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1"
|
||||
)
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
|
||||
caller_auth: Final = agent_caller_auth(agent_key)
|
||||
|
||||
assert caller_auth is not None
|
||||
assert (caller_auth.user_id, caller_auth.team_id, caller_auth.agent_id, caller_auth.api_key) == (
|
||||
"alice",
|
||||
"callers",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
assert agent_caller_auth(_AGENT_KEY) is None
|
||||
|
||||
|
||||
def test_agent_caller_cannot_be_set_from_a_request_payload() -> None:
|
||||
forged: Final = UserAPIKeyAuth.model_validate(
|
||||
{"api_key": "agent-key", "agent_id": "agent-1", "agent_caller": {"user_id": "alice", "team_id": "callers"}}
|
||||
)
|
||||
|
||||
assert forged.agent_caller is None
|
||||
assert "agent_caller" not in forged.model_dump()
|
||||
|
|
@ -9,10 +9,10 @@ from unittest.mock import AsyncMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
|
||||
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
||||
AgentAccess,
|
||||
AgentRequestHandler,
|
||||
|
|
@ -20,6 +20,7 @@ from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
|
|||
UnrestrictedAgentAccess,
|
||||
accessible_agents,
|
||||
)
|
||||
from litellm.types.agents import AgentCaller
|
||||
|
||||
|
||||
def _registry_with(*agent_names: str) -> AgentRegistry:
|
||||
|
|
@ -157,6 +158,130 @@ class TestAgentRequestHandler:
|
|||
is False
|
||||
), agent_id
|
||||
|
||||
@staticmethod
|
||||
def _ceiling_resolver(agent_ids: frozenset[str] | None) -> tuple[CeilingResolver, list[str]]:
|
||||
"""A resolver that records the agent ids it was asked about and answers with a fixed
|
||||
ceiling, or None when the agent has no access groups attached."""
|
||||
asked: Final[list[str]] = []
|
||||
|
||||
async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
|
||||
asked.append(agent_id)
|
||||
if agent_ids is None:
|
||||
return None
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",), models=frozenset(), mcp_server_ids=frozenset(), agent_ids=agent_ids
|
||||
)
|
||||
|
||||
return resolve, asked
|
||||
|
||||
@staticmethod
|
||||
def _key_granting(agent_ids: list[str], agent_id: str | None) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="test-user",
|
||||
agent_id=agent_id,
|
||||
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="obj-1", agents=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")
|
||||
resolve, asked = self._ceiling_resolver(frozenset({"agent-beta"}))
|
||||
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset({"agent-beta"})
|
||||
)
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key, resolve) is True
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False
|
||||
assert asked == ["caller-agent"] * 3
|
||||
|
||||
@staticmethod
|
||||
def _team_grants(grants: dict[str, AgentAccess]) -> AsyncMock:
|
||||
async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> AgentAccess:
|
||||
assert user_api_key_auth is not None
|
||||
return grants.get(user_api_key_auth.team_id or "", UnrestrictedAgentAccess())
|
||||
|
||||
return AsyncMock(side_effect=by_team)
|
||||
|
||||
async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_agents(self):
|
||||
"""LIT-8014: the agent's key and access groups reach alpha and beta, but the human who
|
||||
invoked it belongs to a team granted only beta, so on their behalf the agent reaches only beta."""
|
||||
agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent")
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
resolve, _ = self._ceiling_resolver(frozenset({"agent-alpha", "agent-beta", "agent-gamma"}))
|
||||
|
||||
with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam
|
||||
AgentRequestHandler,
|
||||
"_get_allowed_agents_for_team",
|
||||
self._team_grants({"callers": RestrictedAgentAccess(frozenset({"agent-beta", "agent-gamma"}))}),
|
||||
) as mock_team:
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset({"agent-beta"})
|
||||
)
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False
|
||||
|
||||
assert {call.args[0].team_id for call in mock_team.call_args_list} == {None, "callers"}
|
||||
|
||||
async def test_agent_key_acting_for_a_user_whose_team_grants_no_agent_reaches_none(self):
|
||||
agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent")
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
resolve, _ = self._ceiling_resolver(None)
|
||||
|
||||
with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam
|
||||
AgentRequestHandler,
|
||||
"_get_allowed_agents_for_team",
|
||||
self._team_grants({"callers": RestrictedAgentAccess(frozenset())}),
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset()
|
||||
)
|
||||
|
||||
async def test_agent_key_acting_for_an_ungranted_caller_keeps_its_own_agents(self):
|
||||
agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent")
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
resolve, _ = self._ceiling_resolver(None)
|
||||
|
||||
with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam
|
||||
AgentRequestHandler, "_get_allowed_agents_for_team", self._team_grants({})
|
||||
):
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset({"agent-alpha"})
|
||||
)
|
||||
|
||||
|
||||
async def test_agent_access_groups_intersect_with_key_grants(self):
|
||||
agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent")
|
||||
resolve, _ = self._ceiling_resolver(frozenset({"agent-beta", "agent-gamma"}))
|
||||
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset({"agent-beta"})
|
||||
)
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-gamma", agent_key, resolve) is False
|
||||
|
||||
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")
|
||||
resolve, _ = self._ceiling_resolver(frozenset())
|
||||
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(frozenset())
|
||||
assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False
|
||||
|
||||
async def test_agent_without_access_groups_keeps_key_grants(self):
|
||||
agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent")
|
||||
resolve, asked = self._ceiling_resolver(None)
|
||||
|
||||
assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(
|
||||
frozenset({"agent-alpha"})
|
||||
)
|
||||
assert asked == ["caller-agent"]
|
||||
|
||||
async def test_key_without_agent_never_consults_agent_access_groups(self):
|
||||
plain_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
|
||||
resolve, asked = self._ceiling_resolver(frozenset())
|
||||
|
||||
assert await AgentRequestHandler.resolve_agent_access(plain_key, resolve) == UnrestrictedAgentAccess()
|
||||
assert asked == []
|
||||
|
||||
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."""
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.agents import AgentCaller
|
||||
|
||||
AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]]
|
||||
|
||||
|
|
@ -511,6 +512,24 @@ async def test_message_methods_forward_caller_identity_headers(method: str):
|
|||
assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_agent_calling_another_agent_forwards_the_human_who_invoked_it(method: str):
|
||||
"""LIT-8014: an agent acting for alice calls a second agent through the proxy. That hop must
|
||||
carry alice, not the first agent's owner, so the chain stays capped at what alice may reach."""
|
||||
mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS)
|
||||
agent_key = UserAPIKeyAuth(api_key="sk-agent", user_id="agent-owner", team_id="agent-team", agent_id="agent-1")
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
|
||||
captured = await _invoke_message_method(method, mock_request, agent_key)
|
||||
|
||||
forwarded_headers = captured.agent_extra_headers or {}
|
||||
assert (forwarded_headers.get("X-LiteLLM-User-Id"), forwarded_headers.get("X-LiteLLM-Team-Id")) == (
|
||||
"alice",
|
||||
"callers",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
|
||||
async def test_message_methods_send_the_entra_bearer_for_azure_agents(method: str):
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.proxy.agent_endpoints.agent_registry import (
|
|||
_restore_redacted_litellm_params,
|
||||
redact_sensitive_agent_litellm_params,
|
||||
)
|
||||
from litellm.types.agents import PatchAgentRequest
|
||||
|
||||
# Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression
|
||||
# fixtures) -- never a real key shape, and must never appear in any response.
|
||||
|
|
@ -990,3 +991,138 @@ 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: PatchAgentRequest, 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)
|
||||
|
|
|
|||
|
|
@ -34,20 +34,26 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
WebhookEvent,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver
|
||||
from litellm.types.agents import AgentCaller
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
ExperimentalUIJWTToken,
|
||||
_cache_management_object,
|
||||
_can_object_call_model,
|
||||
_can_object_call_vector_stores,
|
||||
_check_agent_access_group_model_access,
|
||||
_check_end_user_budget,
|
||||
_check_team_member_budget,
|
||||
_fetch_key_object_from_db_with_reconnect,
|
||||
_get_fuzzy_user_object,
|
||||
CallerTeamLoader,
|
||||
CallerUserLoader,
|
||||
_get_team_db_check,
|
||||
_log_budget_lookup_failure,
|
||||
_tag_max_budget_check,
|
||||
_team_max_budget_check,
|
||||
_virtual_key_max_budget_alert_check,
|
||||
_check_agent_caller_model_access,
|
||||
_virtual_key_max_budget_check,
|
||||
_virtual_key_soft_budget_check,
|
||||
get_key_object,
|
||||
|
|
@ -8930,6 +8936,69 @@ def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models()
|
|||
assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False
|
||||
|
||||
|
||||
def _agent_model_ceiling_resolver(
|
||||
models: frozenset[str] | None,
|
||||
) -> tuple[CeilingResolver, list[str]]:
|
||||
"""Resolver that records the agent ids it was asked about and answers with a fixed model
|
||||
ceiling, or None when the agent has no access groups attached."""
|
||||
asked: Final[list[str]] = []
|
||||
|
||||
async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
|
||||
asked.append(agent_id)
|
||||
if models is None:
|
||||
return None
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset()
|
||||
)
|
||||
|
||||
return resolve, asked
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_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"])
|
||||
resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5"}))
|
||||
|
||||
assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied
|
||||
assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN)
|
||||
assert asked == ["agent-1", "agent-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_access_groups_naming_no_model_deny_every_model():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[])
|
||||
resolve, _ = _agent_model_ceiling_resolver(frozenset())
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_without_access_groups_adds_no_model_ceiling():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"])
|
||||
resolve, asked = _agent_model_ceiling_resolver(None)
|
||||
|
||||
assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True
|
||||
assert await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve) is True
|
||||
assert asked == ["agent-1", "agent-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_without_agent_never_consults_agent_access_groups():
|
||||
plain_key: Final = UserAPIKeyAuth(token="plain-token", models=["gpt-5"])
|
||||
resolve, asked = _agent_model_ceiling_resolver(frozenset())
|
||||
|
||||
assert await _check_agent_access_group_model_access("gpt-5", plain_key, None, resolve) is True
|
||||
assert asked == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_check_temp_budget_increase_extends_cap():
|
||||
"""Spend above max_budget but below max_budget + active temp increase
|
||||
|
|
@ -9086,3 +9155,117 @@ async def test_team_member_budget_check_adds_temp_increase_to_live_team_default(
|
|||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
)
|
||||
assert exc_info.value.max_budget == expected_cap
|
||||
|
||||
|
||||
def _agent_key_acting_for(user_id: str | None, team_id: str | None) -> UserAPIKeyAuth:
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"])
|
||||
agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id)
|
||||
return agent_key
|
||||
|
||||
|
||||
def _caller_loaders(
|
||||
team: LiteLLM_TeamTable | None,
|
||||
user: LiteLLM_UserTable | None,
|
||||
) -> tuple[CallerTeamLoader, CallerUserLoader, list[str]]:
|
||||
"""Loaders that hand back fixed caller rows and record the agent_caller they were asked about."""
|
||||
asked: Final[list[str]] = []
|
||||
|
||||
async def load_team(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
|
||||
asked.append(f"team:{valid_token.agent_caller.team_id if valid_token.agent_caller else None}")
|
||||
return team
|
||||
|
||||
async def load_user(valid_token: UserAPIKeyAuth) -> LiteLLM_UserTable | None:
|
||||
asked.append(f"user:{valid_token.agent_caller.user_id if valid_token.agent_caller else None}")
|
||||
return user
|
||||
|
||||
return load_team, load_user, asked
|
||||
|
||||
|
||||
async def _cache_with_membership(user_id: str, team_id: str, allowed_models: list[str] | None) -> UserApiKeyCache:
|
||||
from litellm.proxy._types import LiteLLM_TeamMembership
|
||||
from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key
|
||||
|
||||
cache: Final = UserApiKeyCache()
|
||||
await cache.async_set_cache(
|
||||
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
|
||||
value=LiteLLM_TeamMembership(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(allowed_models=allowed_models) if allowed_models else None,
|
||||
),
|
||||
model_type=LiteLLM_TeamMembership,
|
||||
)
|
||||
return cache
|
||||
|
||||
|
||||
async def _check_caller_models(
|
||||
agent_key: UserAPIKeyAuth,
|
||||
model: str,
|
||||
load_team: CallerTeamLoader,
|
||||
load_user: CallerUserLoader,
|
||||
cache: UserApiKeyCache | None = None,
|
||||
) -> None:
|
||||
await _check_agent_caller_model_access(
|
||||
model=model,
|
||||
valid_token=agent_key,
|
||||
llm_router=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=cache or UserApiKeyCache(),
|
||||
proxy_logging_obj=MagicMock(),
|
||||
load_team=load_team,
|
||||
load_user=load_user,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_acting_for_a_team_is_capped_at_that_teams_models():
|
||||
"""LIT-8014: the invoking team may only call gpt-5, so the agent's own claude grant does not help."""
|
||||
agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a")
|
||||
load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=["gpt-5"]), None)
|
||||
cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=None)
|
||||
|
||||
await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache)
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied
|
||||
assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN)
|
||||
assert asked == ["team:team-a", "team:team-a"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_acting_for_a_team_member_is_capped_at_the_members_scope():
|
||||
agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a")
|
||||
load_team, load_user, _ = _caller_loaders(
|
||||
LiteLLM_TeamTable(team_id="team-a", models=["gpt-5", "claude-sonnet"]), None
|
||||
)
|
||||
cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=["gpt-5"])
|
||||
|
||||
await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache)
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache)
|
||||
|
||||
assert "User=alice, Team=team-a" in exc_info.value.internal_message
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_models():
|
||||
agent_key: Final = _agent_key_acting_for(user_id="alice", team_id=None)
|
||||
load_team, load_user, asked = _caller_loaders(None, LiteLLM_UserTable(user_id="alice", models=["gpt-5"]))
|
||||
|
||||
await _check_caller_models(agent_key, "gpt-5", load_team, load_user)
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user)
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.user_model_access_denied
|
||||
assert asked == ["team:None", "user:alice", "team:None", "user:alice"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_without_an_echoed_caller_keeps_its_own_models():
|
||||
agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"])
|
||||
load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=[]), None)
|
||||
|
||||
await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user)
|
||||
|
||||
assert asked == []
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -110,8 +110,7 @@ def test_create_model_info_response_happy_path_no_metadata():
|
|||
"owned_by": result["owned_by"],
|
||||
"created_is_int": isinstance(result["created"], int),
|
||||
"metadata_absent": "metadata" not in result,
|
||||
"max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int)
|
||||
and result["max_input_tokens"] > 0,
|
||||
"max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) and result["max_input_tokens"] > 0,
|
||||
"max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int)
|
||||
and result["max_output_tokens"] > 0,
|
||||
}
|
||||
|
|
@ -205,9 +204,7 @@ def test_validate_model_access_happy_path_single_model_in_list():
|
|||
|
||||
def test_validate_model_access_happy_path_batch_all_accessible():
|
||||
summary = {
|
||||
"result": validate_model_access(
|
||||
"gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]
|
||||
),
|
||||
"result": validate_model_access("gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]),
|
||||
"input": "gpt-4o,claude-haiku",
|
||||
"available": ["gpt-4o", "claude-haiku", "gemini"],
|
||||
}
|
||||
|
|
@ -389,9 +386,7 @@ async def test_get_available_models_for_user_error_path_complete_list_raises(
|
|||
def _boom(**_kwargs):
|
||||
raise RuntimeError("downstream failure")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.model_checks.get_complete_model_list", _boom
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.auth.model_checks.get_complete_model_list", _boom)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key="sk-test-key",
|
||||
user_id="user-1",
|
||||
|
|
@ -481,6 +476,7 @@ async def test_get_available_models_for_user_without_access_groups_grants_nothin
|
|||
)
|
||||
assert result == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_available_models_for_user_resolves_key_access_group_models(
|
||||
monkeypatch,
|
||||
|
|
@ -521,3 +517,78 @@ async def test_get_available_models_for_user_resolves_key_access_group_models(
|
|||
user_api_key_cache=MagicMock(),
|
||||
)
|
||||
assert result == ["model-b"]
|
||||
|
||||
|
||||
def _agent_ceiling(models: frozenset[str] | None):
|
||||
from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling
|
||||
|
||||
async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None:
|
||||
if models is None:
|
||||
return None
|
||||
return AgentAccessGroupCeiling(
|
||||
access_group_ids=("ag-agent",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset()
|
||||
)
|
||||
|
||||
return resolve
|
||||
|
||||
|
||||
def _agent_key(models: list[str]) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(api_key="sk-agent-key", user_id="user-1", agent_id="agent-1", models=models)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_listing_is_capped_to_its_access_groups():
|
||||
result = await get_available_models_for_user(
|
||||
user_api_key_dict=_agent_key(["model-a", "model-b", "model-c"]),
|
||||
llm_router=_router_with_models(["model-a", "model-b", "model-c"]),
|
||||
general_settings={},
|
||||
user_model=None,
|
||||
resolve_agent_ceiling=_agent_ceiling(frozenset({"model-b", "model-d"})),
|
||||
)
|
||||
assert result == ["model-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_key_listing_is_empty_when_its_groups_grant_no_model():
|
||||
result = await get_available_models_for_user(
|
||||
user_api_key_dict=_agent_key(["model-a"]),
|
||||
llm_router=_router_with_models(["model-a"]),
|
||||
general_settings={},
|
||||
user_model=None,
|
||||
resolve_agent_ceiling=_agent_ceiling(frozenset()),
|
||||
)
|
||||
assert result == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_ceiling_expands_a_model_access_group_name_for_listing():
|
||||
router = _router_with_models(["model-a", "model-b"])
|
||||
router.get_model_access_groups.return_value = {"fast-models": ["model-b"]}
|
||||
result = await get_available_models_for_user(
|
||||
user_api_key_dict=_agent_key(["model-a", "model-b"]),
|
||||
llm_router=router,
|
||||
general_settings={},
|
||||
user_model=None,
|
||||
resolve_agent_ceiling=_agent_ceiling(frozenset({"fast-models"})),
|
||||
)
|
||||
assert result == ["model-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups():
|
||||
router = _router_with_models(["model-a", "model-b"])
|
||||
plain_key = await get_available_models_for_user(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-plain", user_id="user-1", models=["model-a", "model-b"]),
|
||||
llm_router=router,
|
||||
general_settings={},
|
||||
user_model=None,
|
||||
resolve_agent_ceiling=_agent_ceiling(frozenset({"model-a"})),
|
||||
)
|
||||
agent_without_groups = await get_available_models_for_user(
|
||||
user_api_key_dict=_agent_key(["model-a", "model-b"]),
|
||||
llm_router=router,
|
||||
general_settings={},
|
||||
user_model=None,
|
||||
resolve_agent_ceiling=_agent_ceiling(None),
|
||||
)
|
||||
assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"])
|
||||
|
|
|
|||
|
|
@ -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,29 @@ 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;
|
||||
|
|
|
|||
|
|
@ -6211,6 +6211,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
|
|
@ -23640,6 +23640,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;
|
||||
|
|
@ -23787,6 +23789,8 @@ export interface components {
|
|||
};
|
||||
/** AgentResponse */
|
||||
AgentResponse: {
|
||||
/** Access Group Ids */
|
||||
access_group_ids?: string[] | null;
|
||||
/** Agent Card Params */
|
||||
agent_card_params: {
|
||||
[key: string]: unknown;
|
||||
|
|
@ -34983,6 +34987,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