diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql new file mode 100644 index 00000000000..d594b0056df --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 31fc76f7a89..85996430bc5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b0640e4f0dd..2e5cad29b0a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 6f9a2d8c96d..1293a05e4a6 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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" }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 03caca399c3..3f2d2f57fd4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..2a189a76545 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -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 } diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c7b6bca72cf..d6b12e830e1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -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), } diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py new file mode 100644 index 00000000000..49e5407ff88 --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -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), + ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_caller.py b/litellm/proxy/agent_endpoints/auth/agent_caller.py new file mode 100644 index 00000000000..47d43e8f71b --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_caller.py @@ -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 diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index e4dd77e2f82..9fe74bfee3f 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -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): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c2279fb2fe1..b3efde7f05f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4371ce4fda8..08f64e610bd 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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 diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index a6cc5140b15..b4923b0a2dc 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -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 ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 31fc76f7a89..85996430bc5 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index de5f545f109..bfd412db1a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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: diff --git a/litellm/types/agents.py b/litellm/types/agents.py index dbaaab62d86..7f8d8c6af66 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -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 diff --git a/schema.prisma b/schema.prisma index 31fc76f7a89..85996430bc5 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 087c5a03498..5c885a3168d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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( diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py new file mode 100644 index 00000000000..e744e84d671 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -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 diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py new file mode 100644 index 00000000000..b08964503c8 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py @@ -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() diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 383b72e5c58..a87716375e8 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -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.""" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 441e9640ef9..b9a260f5b14 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -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): diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 231626c7eb5..b036e0dac4d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -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) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 44c1d6a3c6b..0cab6535f3d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 == [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index d687f8d1c8d..13e57408bd0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -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") diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 5fb4392eec6..7f77f938323 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -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"]) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index 15b85001cdc..8e100d0c3ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -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; defaultInputModes?: string[]; @@ -121,6 +122,7 @@ export interface AgentRequestPayload { agent_card_params?: Record; litellm_params?: Record; object_permission?: Record; + access_group_ids?: string[]; } interface AgentFormFieldProps { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index ebc97891744..457ee656415 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -21,6 +21,9 @@ vi.mock("./agent_card_discovery", () => ({ default: () =>
({ default: () =>
})); vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () =>
})); vi.mock("@/components/guardrails/GuardrailSelector", () => ({ default: () =>
})); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), +})); vi.mock("@/components/common_components/team_dropdown", () => ({ default: () =>
})); const a2aInfo: AgentCreateInfo = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx index ccac244f019..fbf5cf8c1fb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx @@ -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 }) => ( + + ), +})); + 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"]); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index e71fed40209..5bd6ea9b83a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -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 = ({ 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 = ({ visible, onClose, accessTok )} + + {({ value, onChange }) => ( + + )} + + { 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), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index 357f924cbe7..37e00766a75 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -25,6 +25,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ vi.mock("./agent_card_discovery", () => ({ default: () =>
})); +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: [], }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 29d97a20afe..7e6c7c0e05c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -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 }) => ( +
+ {(value ?? []).join(",")} + + +
+ ), +})); + vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ default: () =>
, })); @@ -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(); + + 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(); + + 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(); + + 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(); + + expect(await screen.findByText("Access Groups")).toBeInTheDocument(); + expect(screen.getByText("None")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index 6cb99e9692f..adac456f232 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -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 = ({ 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 = ({ 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 = ({ 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 = ({ 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 = ({ agentId, onClose, accessT {agent.rpm_limit ?? "Unlimited"} {agent.session_tpm_limit ?? "Unlimited"} {agent.session_rpm_limit ?? "Unlimited"} + + {agent.access_group_ids?.length ? ( +
+ {agent.access_group_ids.map((accessGroupId) => ( +
{accessGroupLabel(accessGroupId)}
+ ))} +
+ ) : ( + "None" + )} + {formatDate(agent.created_at)} {formatDate(agent.updated_at)} @@ -489,6 +518,26 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

Access Groups

+ + + {({ value, onChange }) => ( + + )} + + +

MCP Servers

diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 24ff0c0e12c..6adb3fa9dda 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -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; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index f4011a653ab..82399792674 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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 { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 58b5cedc2cf..839bea0c5d6 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;