From b84f8b6a772d859a6ad762e8429549b86b877ca0 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:24:14 +0000 Subject: [PATCH 01/13] feat(agents): attach access groups to agents and enforce them for models, MCP servers and agent calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + .../mcp_server/auth/user_api_key_auth_mcp.py | 37 ++- litellm/proxy/_lazy_openapi_snapshot.json | 42 +++ litellm/proxy/_types.py | 9 +- .../proxy/agent_endpoints/agent_registry.py | 15 +- .../auth/agent_access_groups.py | 80 ++++++ .../auth/agent_permission_handler.py | 33 ++- litellm/proxy/auth/auth_checks.py | 37 ++- .../access_group_endpoints.py | 72 +++++ litellm/proxy/schema.prisma | 1 + litellm/types/agents.py | 3 + schema.prisma | 1 + .../auth/test_user_api_key_auth_mcp.py | 59 ++++ .../auth/test_agent_access_groups.py | 130 +++++++++ .../auth/test_agent_permission_handler.py | 83 ++++++ .../agent_endpoints/test_agent_registry.py | 133 +++++++++ .../proxy/auth/test_auth_checks.py | 91 ++++++ .../test_access_group_endpoints.py | 271 ++++++++---------- .../agents/_components/AgentFormKit.tsx | 2 + .../add_agent_form.integration.test.tsx | 3 + .../_components/add_agent_form.test.tsx | 30 ++ .../agents/_components/add_agent_form.tsx | 21 ++ .../agents/_components/agent_config.ts | 5 + .../agent_info.integration.test.tsx | 7 + .../agents/_components/agent_info.test.tsx | 70 +++++ .../agents/_components/agent_info.tsx | 53 +++- .../src/components/agents/types.ts | 1 + .../src/components/networking.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 30 files changed, 1138 insertions(+), 161 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql create mode 100644 litellm/proxy/agent_endpoints/auth/agent_access_groups.py create mode 100644 tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py 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 139fb031671..e0b52dd77de 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -71,6 +71,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) 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 b0d57cb6228..bbb3d30864f 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 @@ -1549,10 +1549,18 @@ class MCPRequestHandler: allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent( user_api_key_auth ) - if len(allowed_mcp_servers_for_agent) > 0: + agent_access_group_servers: Final = await MCPRequestHandler._get_agent_access_group_server_ceiling( + user_api_key_auth + ) + if len(allowed_mcp_servers_for_agent) > 0 or agent_access_group_servers is not None: has_lower_level_mcp_restrictions = True - # Intersect: agent can only use servers allowed by BOTH key/team AND agent config - allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent] + # Intersect: agent can only use servers allowed by key/team AND agent config AND agent access groups + allowed_mcp_servers = [ + s + for s in allowed_mcp_servers + if (len(allowed_mcp_servers_for_agent) == 0 or s in allowed_mcp_servers_for_agent) + and (agent_access_group_servers is None or s in agent_access_group_servers) + ] verbose_logger.debug( "Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers ) @@ -3137,6 +3145,29 @@ class MCPRequestHandler: verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] + @staticmethod + async def _get_agent_access_group_server_ceiling( + user_api_key_auth: UserAPIKeyAuth, + ) -> frozenset[str] | None: + """ + Server IDs the agent's attached unified access groups (``LiteLLM_AgentsTable.access_group_ids``) + allow, or None when the agent has none attached. Unlike the object_permission path above, an + attached group set that names no servers is an empty ceiling and denies every server. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + resolve_agent_access_group_ceiling, + ) + + if not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + @staticmethod async def _get_agent_tool_permissions_for_server( server_id: str, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 74f38b3ca6d..cf547dc89d5 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", @@ -3471,6 +3499,20 @@ }, "PatchAgentRequest": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "$ref": "#/components/schemas/AgentCard" }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e8..3a7125acd1f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4127,6 +4127,11 @@ class ProxyErrorTypes(str, enum.Enum): Project does not have access to the model """ + agent_model_access_denied = "agent_model_access_denied" + """ + The agent behind the key does not have access to the model + """ + model_cost_map_missing = "model_cost_map_missing" expired_key = "expired_key" @@ -4201,7 +4206,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org", "project"] + cls, object_type: Literal["key", "user", "team", "org", "project", "agent"] ) -> "ProxyErrorTypes": """ Get the model access error type for object_type @@ -4216,6 +4221,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.org_model_access_denied elif object_type == "project": return cls.project_model_access_denied + elif object_type == "agent": + return cls.agent_model_access_denied @classmethod def get_vector_store_access_error_type_for_object( diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c7b6bca72cf..c8948d8d70e 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 @@ -284,6 +286,12 @@ def _resolved_agent_param_value( return _MISSING_AGENT_PARAM +def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: + if "access_group_ids" not in agent: + return MappingProxyType({}) + return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) + + def _restore_redacted_litellm_params( incoming: Mapping[str, object], existing: Mapping[str, object], @@ -516,6 +524,7 @@ class AgentRegistry: static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None extra_headers_val: Final = agent.get("extra_headers") + access_group_ids_val: Final = agent.get("access_group_ids") create_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -532,6 +541,8 @@ class AgentRegistry: create_data["static_headers"] = static_headers_val if extra_headers_val is not None: create_data["extra_headers"] = extra_headers_val + if access_group_ids_val is not None: + create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val)) if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -601,7 +612,7 @@ class AgentRegistry: existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, object]] = {} + update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if "litellm_params" in agent: @@ -703,6 +714,7 @@ class AgentRegistry: safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) extra_headers_val_u: Final = agent.get("extra_headers") or [] + access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) update_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -710,6 +722,7 @@ class AgentRegistry: "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, + "access_group_ids": access_group_ids_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } 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..4a579de679e --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -0,0 +1,80 @@ +""" +Ceiling that an agent's attached access groups place on requests made with that agent's key. + +Keys and teams use access groups as grants. An agent uses them the way it already uses its +``object_permission``: the union of the attached groups caps what the agent's key can reach, +on top of whatever the key and team allow. A group that cannot be loaded contributes nothing, +so a missing or unreadable group can only narrow the agent, never widen it. +""" + +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Final, TypeAlias + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable +from litellm.types.agents import AgentResponse + +AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]] +AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LiteLLM_AccessGroupTable | None]] + + +@dataclass(frozen=True, slots=True) +class AgentAccessGroupCeiling: + """Everything the agent's attached access groups allow. An empty set denies that resource kind.""" + + access_group_ids: tuple[str, ...] + models: frozenset[str] + mcp_server_ids: frozenset[str] + agent_ids: frozenset[str] + + +async def _load_agent(agent_id: str) -> AgentResponse | None: + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + return await get_agent_with_read_through(agent_id) + + +async def _load_access_group(access_group_id: str) -> LiteLLM_AccessGroupTable | None: + from litellm.proxy.auth.auth_checks import get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as e: + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail + ) + return None + + +async def resolve_agent_access_group_ceiling( + agent_id: str, + load_agent: AgentLoader = _load_agent, + load_access_group: AccessGroupLoader = _load_access_group, +) -> AgentAccessGroupCeiling | None: + """``None`` when the agent has no access groups attached, so nothing is capped.""" + agent: Final = await load_agent(agent_id) + access_group_ids: Final = tuple(agent.access_group_ids or ()) if agent is not None else () + if not access_group_ids: + return None + + loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids)) + groups: Final = tuple(group for group in loaded if group is not None) + return AgentAccessGroupCeiling( + access_group_ids=access_group_ids, + models=frozenset(model for group in groups for model in group.access_model_names), + mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), + agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), + ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index e4dd77e2f82..11d2a68072c 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -66,9 +66,24 @@ class AgentRequestHandler: Resolve the agents the given user/key may reach. ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team - carries any grant. Grants that intersect to nothing stay restricted, so - narrowing a caller can never widen what it reaches. + carries any grant and the agent behind the key has no access groups attached. + Grants that intersect to nothing stay restricted, so narrowing a caller can + never widen what it reaches. """ + key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) + agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth) + if agent_ceiling is None: + return key_team_access + match key_team_access: + case UnrestrictedAgentAccess(): + return RestrictedAgentAccess(agent_ceiling) + case RestrictedAgentAccess(key_team_ids): + return RestrictedAgentAccess(key_team_ids & agent_ceiling) + + @staticmethod + async def _resolve_key_team_agent_access( + user_api_key_auth: UserAPIKeyAuth | None, + ) -> AgentAccess: try: key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) @@ -86,6 +101,20 @@ class AgentRequestHandler: verbose_logger.warning("Failed to get allowed agents: %s", e) return UnrestrictedAgentAccess() + @staticmethod + async def _agent_access_group_ceiling( + user_api_key_auth: UserAPIKeyAuth | None, + ) -> frozenset[str] | None: + """Stable IDs of the agents the calling agent's attached access groups allow; None when none attached.""" + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling + + if user_api_key_auth is None or not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return _to_stable_ids(ceiling.agent_ids) + @staticmethod async def is_agent_allowed( agent_id: str, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3dd2e2d8eb2..5c53ee49717 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1003,6 +1003,9 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) + # 2.4 If the agent behind the key has access groups attached, they cap the models it can call + await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) + ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): @@ -4126,7 +4129,7 @@ def _can_object_call_model( models: list[str], team_model_aliases: dict[str, str] | None = None, team_id: str | None = None, - object_type: Literal["user", "team", "key", "org", "project"] = "user", + object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ @@ -4192,6 +4195,38 @@ def _can_object_call_model( ) +async def _check_agent_access_group_model_access( + model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, +) -> Literal[True]: + """Raises when the key's agent has access groups attached and none of them names the model. + Attached groups that name no model deny every model; ``_can_object_call_model`` would read + an empty allowlist as unrestricted.""" + from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling + + if not model or valid_token is None or not valid_token.agent_id: + return True + ceiling: Final = await resolve_agent_access_group_ceiling(valid_token.agent_id) + if ceiling is None: + return True + if not ceiling.models: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models", + type=ProxyErrorTypes.agent_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return _can_object_call_model( + model=model, + llm_router=llm_router, + models=sorted(ceiling.models), + team_id=valid_token.team_id, + object_type="agent", + ) + + def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: """ Returns True if `model` being accessed is an alias of a team model 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 139fb031671..e0b52dd77de 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -71,6 +71,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) diff --git a/litellm/types/agents.py b/litellm/types/agents.py index dbaaab62d86..12e60352a97 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -189,6 +189,7 @@ class AgentConfig(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] class PatchAgentRequest(TypedDict, total=False): @@ -202,6 +203,7 @@ class PatchAgentRequest(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] # Request/Response models for CRUD endpoints @@ -226,6 +228,7 @@ class AgentResponse(BaseModel): session_rpm_limit: int | None = None static_headers: dict[str, str] | None = None extra_headers: list[str] | None = None + access_group_ids: Sequence[str] | None = None keys: list[AgentKeySummary] | None = None search_score: float | None = None created_at: datetime | None = None diff --git a/schema.prisma b/schema.prisma index 139fb031671..e0b52dd77de 100644 --- a/schema.prisma +++ b/schema.prisma @@ -71,6 +71,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) 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 90ce821d62e..2f8e3d1cb82 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 @@ -4208,6 +4208,65 @@ class TestAgentMCPPermissions: assert sorted(result) == ["server_1", "server_2"] mock_agent.assert_called_once_with(user_api_key_auth) + @pytest.mark.parametrize( + ("group_ceiling", "expected"), + [ + (frozenset({"server_1"}), ["server_1"]), + (frozenset({"server_1", "server_2", "server_3"}), ["server_1", "server_2"]), + (frozenset(), []), + ], + ) + async def test_get_allowed_mcp_servers_agent_access_group_ceiling(self, group_ceiling, expected): + """The agent's attached access groups cap the key/team servers; groups naming no server deny all.""" + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag") + with ( + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]), + patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=group_ceiling), + ): + access = await MCPRequestHandler.get_mcp_server_access(user_api_key_auth=user_api_key_auth) + assert sorted(access.server_ids) == expected + assert access.scope == "scoped" + + async def test_get_allowed_mcp_servers_agent_without_access_groups_is_uncapped(self): + user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag") + with ( + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]), + patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]), + patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=None), + ): + result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=user_api_key_auth) + assert sorted(result) == ["server_1", "server_2"] + + async def test_agent_access_group_server_ceiling_expands_group_servers(self): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + ceiling = AgentAccessGroupCeiling( + access_group_ids=("ag-1",), + models=frozenset(), + mcp_server_ids=frozenset({"server_1"}), + agent_ids=frozenset(), + ) + with ( + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=ceiling), + ), + patch( + "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" + ) as mock_manager, + ): + mock_manager.expand_permission_list.return_value = ["server_1"] + result = await MCPRequestHandler._get_agent_access_group_server_ceiling( + UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag") + ) + assert result == frozenset({"server_1"}) + mock_manager.expand_permission_list.assert_called_once_with(["server_1"]) + + assert await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k")) is None + async def test_get_allowed_mcp_servers_key_team_agent_intersection(self): """Key allows [1, 2], agent allows [2, 3]. Result = [2].""" user_api_key_auth = UserAPIKeyAuth( 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..8c31c9428e7 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -0,0 +1,130 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, +) +from litellm.types.agents import AgentResponse + +_CARD: Final = {"name": "agent", "url": "http://localhost:9999", "version": "1.0.0"} + + +def _agent(access_group_ids: list[str] | None) -> AgentResponse: + return AgentResponse( + agent_id="agent-1", agent_name="agent", agent_card_params=_CARD, access_group_ids=access_group_ids + ) + + +def _group( + group_id: str, + models: tuple[str, ...] = (), + mcp_servers: tuple[str, ...] = (), + agents: tuple[str, ...] = (), +) -> LiteLLM_AccessGroupTable: + return LiteLLM_AccessGroupTable( + access_group_id=group_id, + access_group_name=group_id, + access_model_names=list(models), + access_mcp_server_ids=list(mcp_servers), + access_agent_ids=list(agents), + ) + + +def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]): + async def load_agent(agent_id: str) -> AgentResponse | None: + return agent + + async def load_group(group_id: str) -> LiteLLM_AccessGroupTable | None: + return groups.get(group_id) + + return load_agent, load_group + + +@pytest.mark.asyncio +@pytest.mark.parametrize("access_group_ids", [None, []]) +async def test_agent_without_access_groups_has_no_ceiling(access_group_ids: list[str] | None): + load_agent, load_group = _loaders(_agent(access_group_ids), {"g1": _group("g1", models=("gpt-5",))}) + + assert await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_unknown_agent_has_no_ceiling(): + load_agent, load_group = _loaders(None, {}) + + assert await resolve_agent_access_group_ceiling("missing", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_ceiling_is_the_union_of_every_attached_group(): + load_agent, load_group = _loaders( + _agent(["g1", "g2"]), + { + "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)), + "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)), + }, + ) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "g2"), + models=frozenset({"gpt-5", "claude-sonnet"}), + mcp_server_ids=frozenset({"mcp-a", "mcp-b"}), + agent_ids=frozenset({"agent-b", "agent-c"}), + ) + + +@pytest.mark.asyncio +async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies(): + load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "gone"), + models=frozenset({"gpt-5"}), + mcp_server_ids=frozenset(), + agent_ids=frozenset(), + ) + + +@pytest.mark.asyncio +async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): + load_agent, load_group = _loaders(_agent(["gone"]), {}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling is not None + assert ceiling.models == frozenset() + assert ceiling.mcp_server_ids == frozenset() + assert ceiling.agent_ids == frozenset() + + +@pytest.mark.asyncio +async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + from litellm.proxy.auth import auth_checks + + async def missing_group(**_: object) -> LiteLLM_AccessGroupTable: + raise HTTPException(status_code=404, detail={"error": "Access group doesn't exist in db."}) + + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(auth_checks, "get_access_object", missing_group) + + assert await _load_access_group("gone") is None + + +@pytest.mark.asyncio +async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + + monkeypatch.setattr(proxy_server, "prisma_client", None) + + assert await _load_access_group("ag-1") is None 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..a8a55d332b6 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 @@ -13,6 +13,7 @@ import pytest from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentAccess, AgentRequestHandler, @@ -157,6 +158,88 @@ class TestAgentRequestHandler: is False ), agent_id + @staticmethod + def _ceiling(agent_ids: frozenset[str]) -> AgentAccessGroupCeiling: + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), + models=frozenset(), + mcp_server_ids=frozenset(), + agent_ids=agent_ids, + ) + + async def test_agent_access_groups_cap_an_otherwise_unrestricted_key(self): + """A key with no agent grant of its own may still only reach the agents its + agent's attached access groups name.""" + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + + with ( + patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), + patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta"}))), + ) as mock_ceiling, + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key) is True + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False + mock_ceiling.assert_called_with("caller-agent") + + async def test_agent_access_groups_intersect_with_key_and_team_grants(self): + agent_key: Final = UserAPIKeyAuth( + api_key="test-key", user_id="test-user", team_id="test-team", agent_id="caller-agent" + ) + + with ( + patch.object( + AgentRequestHandler, + "_get_allowed_agents_for_key", + return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta"})), + ), + patch.object( + AgentRequestHandler, + "_get_allowed_agents_for_team", + return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta", "agent-gamma"})), + ), + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta", "agent-gamma"}))), + ), + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + + async def test_agent_access_groups_naming_no_agent_deny_every_agent(self): + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + + with ( + patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), + patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=self._ceiling(frozenset())), + ), + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess(frozenset()) + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False + + async def test_key_without_agent_never_consults_agent_access_groups(self): + plain_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + + with ( + patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), + patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=self._ceiling(frozenset())), + ) as mock_ceiling, + ): + assert await AgentRequestHandler.resolve_agent_access(plain_key) == UnrestrictedAgentAccess() + mock_ceiling.assert_not_called() + async def test_empty_access_group_denies_every_agent(self): """LIT-5143: a key restricted to an access group that resolves to no agents is restricted to nothing, not unrestricted. A failed group lookup still fails open.""" 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..d15a3adadbd 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -990,3 +990,136 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY assert stored_params["is_public"] is True + + +def _agent_row_mock(access_group_ids: list[str]) -> MagicMock: + row: Final = MagicMock() + row.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + "access_group_ids": access_group_ids, + } + row.object_permission = None + return row + + +@pytest.mark.asyncio +async def test_add_agent_to_db_persists_deduplicated_access_group_ids(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock(["ag-1", "ag-2"])) + mock_prisma.db.litellm_agentstable.create = mock_create + + result: Final = await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "access_group_ids": ["ag-1", "ag-2", "ag-1"], + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert tuple(mock_create.call_args.kwargs["data"]["access_group_ids"]) == ("ag-1", "ag-2") + assert result.access_group_ids == ["ag-1", "ag-2"] + + +@pytest.mark.asyncio +async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_default(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params()}, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert "access_group_ids" not in mock_create.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("patch_body", "expected"), + [ + ({"access_group_ids": ["ag-2", "ag-3"]}, ["ag-2", "ag-3"]), + ({"access_group_ids": []}, []), + ({"access_group_ids": None}, []), + ], +) +async def test_patch_agent_in_db_replaces_access_group_ids_when_provided(patch_body: dict, expected: list[str]): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Test Agent", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent=patch_body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old Name", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"agent_name": "New Name"}, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert "access_group_ids" not in mock_update.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body_access_group_ids", "expected"), + [(["ag-9", "ag-9"], ["ag-9"]), (None, []), ("omitted", [])], +) +async def test_update_agent_in_db_always_writes_access_group_ids(body_access_group_ids, expected: list[str]): + """PUT is a full replacement: omitting the field clears any previously attached groups.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"]) + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + body: Final = { + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"model": "bedrock/agentcore/my-agent"}, + **({} if body_access_group_ids == "omitted" else {"access_group_ids": body_access_group_ids}), + } + + await registry.update_agent_in_db( + agent_id="agent-123", agent=body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index f480e096081..ad1742db4b7 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8461,3 +8461,94 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None: def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() -> None: assert request_skips_budget_checks(route="/v1/models", model=None, llm_router=None) is True assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False + + +# Agent access group model ceiling + + +def _agent_model_ceiling(models: frozenset[str]): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + +async def _run_common_checks_for_agent_key(model: str, valid_token: UserAPIKeyAuth): + from fastapi import Request + + from litellm.proxy.auth.auth_checks import common_checks + + return await common_checks( + request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=valid_token, + request=MagicMock(spec=Request), + ) + + +@pytest.mark.asyncio +async def test_common_checks_agent_access_groups_cap_models_even_when_key_allows_them(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + + with patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=_agent_model_ceiling(frozenset({"gpt-5"}))), + ): + assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True + + with pytest.raises(ProxyException) as exc_info: + await _run_common_checks_for_agent_key("claude-sonnet", agent_key) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + + +@pytest.mark.asyncio +async def test_common_checks_agent_access_groups_naming_no_model_deny_every_model(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + + with ( + patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=_agent_model_ceiling(frozenset())), + ), + pytest.raises(ProxyException) as exc_info, + ): + await _run_common_checks_for_agent_key("gpt-5", agent_key) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + + +@pytest.mark.asyncio +async def test_common_checks_agent_without_access_groups_adds_no_model_ceiling(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + + with patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=None), + ) as mock_ceiling: + assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True + assert await _run_common_checks_for_agent_key("claude-sonnet", agent_key) is True + + mock_ceiling.assert_called_with("agent-1") + + +@pytest.mark.asyncio +async def test_common_checks_key_without_agent_never_consults_agent_access_groups(): + plain_key: Final = UserAPIKeyAuth(token="plain-token", models=["gpt-5"]) + + with patch( + "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", + new=AsyncMock(return_value=_agent_model_ceiling(frozenset())), + ) as mock_ceiling: + assert await _run_common_checks_for_agent_key("gpt-5", plain_key) is True + + mock_ceiling.assert_not_called() 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/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..b5301d0e6a0 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,27 @@ describe("AddAgentForm logos", () => { await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; expect(payload.object_permission).toEqual({ mcp_toolsets: ["ts-1"] }); + expect(payload).not.toHaveProperty("access_group_ids"); + }); + + it("includes selected access groups in the create payload", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + vi.mocked(networking.createAgentCall).mockReset().mockResolvedValue({ + agent_id: "agent-1", + agent_name: "Test Agent", + } as never); + vi.mocked(networking.keyListCall).mockResolvedValue({ keys: [] }); + + renderForm(); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByTestId("select-access-group")); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByText(/Skip for now/)); + await user.click(screen.getByRole("button", { name: "Create Agent →" })); + + await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); + const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual(["ag-1", "ag-2"]); }); }); 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 e77c8ba7e41..7714f96804f 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6267,6 +6267,7 @@ export const patchAgentCall = async ( rpm_limit?: number | null; session_tpm_limit?: number | null; session_rpm_limit?: number | null; + access_group_ids?: string[]; }, ) => { try { diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 872875cc535..be193647dca 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23195,6 +23195,8 @@ export interface components { }; /** AgentConfig */ AgentConfig: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name: string; @@ -23342,6 +23344,8 @@ export interface components { }; /** AgentResponse */ AgentResponse: { + /** Access Group Ids */ + access_group_ids?: string[] | null; /** Agent Card Params */ agent_card_params: { [key: string]: unknown; @@ -34111,6 +34115,8 @@ export interface components { }; /** PatchAgentRequest */ PatchAgentRequest: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params?: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name?: string; From 433c804a6cfec9c6674872857f96b435f69f4d96 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:32:39 +0000 Subject: [PATCH 02/13] refactor(agents): mark Callable loader aliases for the type discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/auth/agent_access_groups.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 4a579de679e..f5adc897e11 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -18,8 +18,9 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import LiteLLM_AccessGroupTable from litellm.types.agents import AgentResponse -AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]] -AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LiteLLM_AccessGroupTable | None]] +AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]] # mutable-ok: Callable parameter syntax +LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None +AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax @dataclass(frozen=True, slots=True) @@ -38,7 +39,7 @@ async def _load_agent(agent_id: str) -> AgentResponse | None: return await get_agent_with_read_through(agent_id) -async def _load_access_group(access_group_id: str) -> LiteLLM_AccessGroupTable | None: +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 From 322262db01f9ba69b1a77ab2ec26268c468a3099 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 19:42:56 +0000 Subject: [PATCH 03/13] style(ui): format add_agent_form test with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../agents/_components/add_agent_form.test.tsx | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) 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 b5301d0e6a0..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 @@ -154,10 +154,12 @@ describe("AddAgentForm logos", () => { 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.createAgentCall) + .mockReset() + .mockResolvedValue({ + agent_id: "agent-1", + agent_name: "Test Agent", + } as never); vi.mocked(networking.keyListCall).mockResolvedValue({ keys: [] }); renderForm(); From d743e08432ee738355a144046b57321485401fea Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:04:23 +0000 Subject: [PATCH 04/13] test(agents): inject the access group ceiling resolver instead of patching it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 47 ++++--- .../auth/agent_access_groups.py | 3 + .../auth/agent_permission_handler.py | 15 ++- litellm/proxy/auth/auth_checks.py | 9 +- .../auth/test_user_api_key_auth_mcp.py | 109 ++++++++-------- .../auth/test_agent_permission_handler.py | 118 ++++++++---------- .../proxy/auth/test_auth_checks.py | 93 ++++++-------- 7 files changed, 193 insertions(+), 201 deletions(-) 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 bbb3d30864f..05661584a6b 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,10 @@ 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.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 +188,24 @@ 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: + """Servers left once the agent's object_permission and attached access groups both cap the + key/team result, or None when the agent restricts nothing. An attached group set naming no + server is an empty ceiling, not an absent one, so it denies every server.""" + 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,21 +1568,14 @@ 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), ) - agent_access_group_servers: Final = await MCPRequestHandler._get_agent_access_group_server_ceiling( - user_api_key_auth - ) - if len(allowed_mcp_servers_for_agent) > 0 or agent_access_group_servers is not None: + if agent_capped is not None: has_lower_level_mcp_restrictions = True - # Intersect: agent can only use servers allowed by key/team AND agent config AND agent access groups - allowed_mcp_servers = [ - s - for s in allowed_mcp_servers - if (len(allowed_mcp_servers_for_agent) == 0 or s in allowed_mcp_servers_for_agent) - and (agent_access_group_servers is None or s in agent_access_group_servers) - ] + allowed_mcp_servers = list(agent_capped) verbose_logger.debug( "Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers ) @@ -3148,6 +3163,7 @@ class MCPRequestHandler: @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``) @@ -3157,13 +3173,10 @@ class MCPRequestHandler: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) - from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( - resolve_agent_access_group_ceiling, - ) if not user_api_key_auth.agent_id: return None - ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id) + 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))) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index f5adc897e11..67bb43638e0 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -33,6 +33,9 @@ class AgentAccessGroupCeiling: agent_ids: frozenset[str] +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params + + async def _load_agent(agent_id: str) -> AgentResponse | None: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 11d2a68072c..1759090a29a 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -19,6 +19,10 @@ 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.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -61,6 +65,7 @@ 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. @@ -71,7 +76,7 @@ class AgentRequestHandler: never widen what it reaches. """ key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) - agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth) + agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) if agent_ceiling is None: return key_team_access match key_team_access: @@ -104,13 +109,12 @@ class AgentRequestHandler: @staticmethod async def _agent_access_group_ceiling( user_api_key_auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver, ) -> frozenset[str] | None: """Stable IDs of the agents the calling agent's attached access groups allow; None when none attached.""" - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling - if user_api_key_auth is None or not user_api_key_auth.agent_id: return None - ceiling: Final = await resolve_agent_access_group_ceiling(user_api_key_auth.agent_id) + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) if ceiling is None: return None return _to_stable_ids(ceiling.agent_ids) @@ -119,6 +123,7 @@ class AgentRequestHandler: 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. @@ -132,7 +137,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 5c53ee49717..137cf849389 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -68,6 +68,10 @@ 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.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, @@ -4199,15 +4203,14 @@ 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]: """Raises when the key's agent has access groups attached and none of them names the model. Attached groups that name no model deny every model; ``_can_object_call_model`` would read an empty allowlist as unrestricted.""" - from litellm.proxy.agent_endpoints.auth.agent_access_groups import resolve_agent_access_group_ceiling - if not model or valid_token is None or not valid_token.agent_id: return True - ceiling: Final = await resolve_agent_access_group_ceiling(valid_token.agent_id) + ceiling: Final = await resolve_ceiling(valid_token.agent_id) if ceiling is None: return True if not ceiling.models: 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 2f8e3d1cb82..f8e648fd383 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 ( @@ -4169,6 +4170,27 @@ 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.""" @@ -4208,64 +4230,45 @@ class TestAgentMCPPermissions: assert sorted(result) == ["server_1", "server_2"] mock_agent.assert_called_once_with(user_api_key_auth) - @pytest.mark.parametrize( - ("group_ceiling", "expected"), - [ - (frozenset({"server_1"}), ["server_1"]), - (frozenset({"server_1", "server_2", "server_3"}), ["server_1", "server_2"]), - (frozenset(), []), - ], - ) - async def test_get_allowed_mcp_servers_agent_access_group_ceiling(self, group_ceiling, expected): - """The agent's attached access groups cap the key/team servers; groups naming no server deny all.""" - user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag") - with ( - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]), - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]), - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]), - patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=group_ceiling), - ): - access = await MCPRequestHandler.get_mcp_server_access(user_api_key_auth=user_api_key_auth) - assert sorted(access.server_ids) == expected - assert access.scope == "scoped" - - async def test_get_allowed_mcp_servers_agent_without_access_groups_is_uncapped(self): - user_api_key_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="agent-ag") - with ( - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key", return_value=["server_1", "server_2"]), - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team", return_value=[]), - patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", return_value=[]), - patch.object(MCPRequestHandler, "_get_agent_access_group_server_ceiling", return_value=None), - ): - result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=user_api_key_auth) - assert sorted(result) == ["server_1", "server_2"] - async def test_agent_access_group_server_ceiling_expands_group_servers(self): + from litellm.proxy._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 - ceiling = AgentAccessGroupCeiling( - access_group_ids=("ag-1",), - models=frozenset(), - mcp_server_ids=frozenset({"server_1"}), - agent_ids=frozenset(), - ) - with ( - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=ceiling), - ), - patch( - "litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager" - ) as mock_manager, - ): - mock_manager.expand_permission_list.return_value = ["server_1"] - result = await MCPRequestHandler._get_agent_access_group_server_ceiling( - UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag") + 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(), ) - assert result == frozenset({"server_1"}) - mock_manager.expand_permission_list.assert_called_once_with(["server_1"]) - assert await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k")) is None + 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].""" 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 a8a55d332b6..2a98e6e4feb 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 @@ -11,9 +11,9 @@ 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 +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentAccess, AgentRequestHandler, @@ -159,86 +159,74 @@ class TestAgentRequestHandler: ), agent_id @staticmethod - def _ceiling(agent_ids: frozenset[str]) -> AgentAccessGroupCeiling: - return AgentAccessGroupCeiling( - access_group_ids=("ag-1",), - models=frozenset(), - mcp_server_ids=frozenset(), - agent_ids=agent_ids, + 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"})) - with ( - patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), - patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta"}))), - ) as mock_ceiling, - ): - assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess( - frozenset({"agent-beta"}) - ) - assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key) is True - assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False - mock_ceiling.assert_called_with("caller-agent") - - async def test_agent_access_groups_intersect_with_key_and_team_grants(self): - agent_key: Final = UserAPIKeyAuth( - api_key="test-key", user_id="test-user", team_id="test-team", agent_id="caller-agent" + 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 - with ( - patch.object( - AgentRequestHandler, - "_get_allowed_agents_for_key", - return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta"})), - ), - patch.object( - AgentRequestHandler, - "_get_allowed_agents_for_team", - return_value=RestrictedAgentAccess(frozenset({"agent-alpha", "agent-beta", "agent-gamma"})), - ), - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=self._ceiling(frozenset({"agent-beta", "agent-gamma"}))), - ), - ): - assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess( - frozenset({"agent-beta"}) - ) + async def test_agent_access_groups_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()) - with ( - patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), - patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=self._ceiling(frozenset())), - ), - ): - assert await AgentRequestHandler.resolve_agent_access(agent_key) == RestrictedAgentAccess(frozenset()) - assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key) is False + 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()) - with ( - patch.object(AgentRequestHandler, "_get_allowed_agents_for_key", return_value=UnrestrictedAgentAccess()), - patch.object(AgentRequestHandler, "_get_allowed_agents_for_team", return_value=UnrestrictedAgentAccess()), - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=self._ceiling(frozenset())), - ) as mock_ceiling, - ): - assert await AgentRequestHandler.resolve_agent_access(plain_key) == UnrestrictedAgentAccess() - mock_ceiling.assert_not_called() + 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 diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index ad1742db4b7..b9b9a786d93 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -32,11 +32,13 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver 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, @@ -8466,89 +8468,64 @@ def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() # Agent access group model ceiling -def _agent_model_ceiling(models: frozenset[str]): - from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling +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]] = [] - return AgentAccessGroupCeiling( - access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() - ) + 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() + ) - -async def _run_common_checks_for_agent_key(model: str, valid_token: UserAPIKeyAuth): - from fastapi import Request - - from litellm.proxy.auth.auth_checks import common_checks - - return await common_checks( - request_body={"model": model, "messages": [{"role": "user", "content": "hi"}]}, - team_object=None, - user_object=None, - end_user_object=None, - global_proxy_spend=None, - general_settings={}, - route="/chat/completions", - llm_router=None, - proxy_logging_obj=MagicMock(), - valid_token=valid_token, - request=MagicMock(spec=Request), - ) + return resolve, asked @pytest.mark.asyncio -async def test_common_checks_agent_access_groups_cap_models_even_when_key_allows_them(): +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"})) - with patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=_agent_model_ceiling(frozenset({"gpt-5"}))), - ): - assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True + assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True - with pytest.raises(ProxyException) as exc_info: - await _run_common_checks_for_agent_key("claude-sonnet", agent_key) + 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_common_checks_agent_access_groups_naming_no_model_deny_every_model(): +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 ( - patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=_agent_model_ceiling(frozenset())), - ), - pytest.raises(ProxyException) as exc_info, - ): - await _run_common_checks_for_agent_key("gpt-5", agent_key) + 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_common_checks_agent_without_access_groups_adds_no_model_ceiling(): +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) - with patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=None), - ) as mock_ceiling: - assert await _run_common_checks_for_agent_key("gpt-5", agent_key) is True - assert await _run_common_checks_for_agent_key("claude-sonnet", agent_key) is True - - mock_ceiling.assert_called_with("agent-1") + 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_common_checks_key_without_agent_never_consults_agent_access_groups(): +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()) - with patch( - "litellm.proxy.agent_endpoints.auth.agent_access_groups.resolve_agent_access_group_ceiling", - new=AsyncMock(return_value=_agent_model_ceiling(frozenset())), - ) as mock_ceiling: - assert await _run_common_checks_for_agent_key("gpt-5", plain_key) is True - - mock_ceiling.assert_not_called() + assert await _check_agent_access_group_model_access("gpt-5", plain_key, None, resolve) is True + assert asked == [] From 5b04560997e5838743d852dd6af97749e5979ff9 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:30:37 +0000 Subject: [PATCH 05/13] refactor(agents): keep the model listing cap within the type-discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_access_groups.py | 104 ++++++++++++++---- litellm/proxy/utils.py | 60 +++++++++- 2 files changed, 144 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index 67bb43638e0..e0e5193d0a4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,33 +1,41 @@ -""" -Ceiling that an agent's attached access groups place on requests made with that agent's key. - -Keys and teams use access groups as grants. An agent uses them the way it already uses its -``object_permission``: the union of the attached groups caps what the agent's key can reach, -on top of whatever the key and team allow. A group that cannot be loaded contributes nothing, -so a missing or unreadable group can only narrow the agent, never widen it. -""" - import asyncio -from collections.abc import Awaitable, Callable +from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Final, Protocol, TypeAlias from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_AccessGroupTable -from litellm.types.agents import AgentResponse +from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl -AgentLoader: TypeAlias = Callable[[str], Awaitable[AgentResponse | None]] # mutable-ok: Callable parameter syntax + +class _AgentAccessGroupsRecord(Protocol): + @property + def access_group_ids(self) -> Sequence[str] | None: ... + + +class _AgentIdWhere(TypedDict): + agent_id: ReadOnly[str] + + +AccessGroupIds: TypeAlias = tuple[str, ...] +AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params +AgentRecordFinder: TypeAlias = Callable[[str], Awaitable[_AgentAccessGroupsRecord | None]] # mutable-ok: Callable LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax +_CACHED_IDS: Final = TypeAdapter(list[str]) + @dataclass(frozen=True, slots=True) class AgentAccessGroupCeiling: """Everything the agent's attached access groups allow. An empty set denies that resource kind.""" - access_group_ids: tuple[str, ...] + access_group_ids: AccessGroupIds models: frozenset[str] mcp_server_ids: frozenset[str] agent_ids: frozenset[str] @@ -36,10 +44,69 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params -async def _load_agent(agent_id: str) -> AgentResponse | None: +def agent_access_group_ids_cache_key(agent_id: str) -> str: + return f"agent_access_group_ids:{agent_id}" + + +def _cached_access_group_ids(cached: object) -> AccessGroupIds | None: + if cached is None: + return None + try: + return tuple(_CACHED_IDS.validate_python(cached)) + except ValidationError: + return None + + +async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through - return await get_agent_with_read_through(agent_id) + 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_agent_access_group_ids( + agent_id: str, + cache: DualCache, + find_agent: AgentRecordFinder, + fallback: AccessGroupIdsLoader, +) -> AccessGroupIds: + """The agent row's groups, cached for the management-object TTL and evicted on every agent write.""" + cache_key: Final = agent_access_group_ids_cache_key(agent_id) + cached: Final = _cached_access_group_ids(await cache.async_get_cache(key=cache_key)) + if cached is not None: + return cached + try: + record: Final = await find_agent(agent_id) + except Exception as e: # noqa: BLE001 # prisma raises many error types; the registry snapshot answers instead + verbose_proxy_logger.warning("Failed to read access groups for agent %r, using registry: %s", agent_id, e) + return await fallback(agent_id) + access_group_ids: Final = tuple(record.access_group_ids or ()) if record is not None else () + await cache.async_set_cache(key=cache_key, value=access_group_ids, ttl=get_management_object_ttl(cache)) + return access_group_ids + + +async def _load_agent_access_group_ids(agent_id: str) -> AccessGroupIds: + from litellm.proxy.agent_endpoints.agent_registry import agents_table + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + return await _registry_access_group_ids(agent_id) + db: Final = prisma_client + + async def find_agent(row_agent_id: str) -> _AgentAccessGroupsRecord | None: + return await agents_table(db).find_unique(where=_AgentIdWhere(agent_id=row_agent_id)) + + return await load_agent_access_group_ids(agent_id, user_api_key_cache, find_agent, _registry_access_group_ids) + + +async def evict_agent_access_group_ids(agent_ids: Sequence[str]) -> None: + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast( + cache_keys=tuple(agent_access_group_ids_cache_key(agent_id) for agent_id in agent_ids), + user_api_key_cache=user_api_key_cache, + ) async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: @@ -65,12 +132,11 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: async def resolve_agent_access_group_ceiling( agent_id: str, - load_agent: AgentLoader = _load_agent, + load_access_group_ids: AccessGroupIdsLoader = _load_agent_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.""" - agent: Final = await load_agent(agent_id) - access_group_ids: Final = tuple(agent.access_group_ids or ()) if agent is not None else () + access_group_ids: Final = await load_access_group_ids(agent_id) if not access_group_ids: return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 950ac5e9906..95c221a8f20 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -122,6 +122,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.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -7980,6 +7981,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: Optional["Router"], + 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"], @@ -7992,6 +8038,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. @@ -8095,7 +8142,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: From 3a86567c9d3ffcd12507407215ea14d9d89d2184 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:31:03 +0000 Subject: [PATCH 06/13] fix(agents): evict the cached agent access groups on every agent write and cap the model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/agent_endpoints/agent_registry.py | 3 + litellm/proxy/agent_endpoints/endpoints.py | 4 + .../access_group_endpoints.py | 2 + .../auth/test_agent_access_groups.py | 91 ++++++++++++++++++- .../proxy/utils/helpers/test_model_access.py | 87 ++++++++++++++++-- 5 files changed, 177 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c8948d8d70e..d6b12e830e1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -67,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: ... diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index aa8979a73c6..62783fb412a 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -44,6 +44,7 @@ from litellm.proxy.agent_endpoints.agent_search import ( global_agent_search_index, search_agents, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -696,6 +697,7 @@ async def update_agent( prisma_client=prisma_client, updated_by=updated_by, ) + await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -799,6 +801,7 @@ async def patch_agent( prisma_client=prisma_client, updated_by=updated_by, ) + await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -861,6 +864,7 @@ async def delete_agent( raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.") await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client) + await evict_agent_access_group_ids((agent_id,)) AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index b4923b0a2dc..2694d00b17f 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -16,6 +16,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry +from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -782,6 +783,7 @@ async def delete_access_group( await invalidate_access_group_cache(access_group_id) _detach_access_group_from_agent_registry(detached_agent_ids, access_group_id) + await evict_agent_access_group_ids(detached_agent_ids) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) 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 index 8c31c9428e7..266ce52e3a9 100644 --- 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 @@ -1,11 +1,16 @@ +from collections.abc import Sequence +from dataclasses import dataclass from typing import Final import pytest from fastapi import HTTPException +from litellm.caching.dual_cache import DualCache from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( AgentAccessGroupCeiling, + agent_access_group_ids_cache_key, + load_agent_access_group_ids, resolve_agent_access_group_ceiling, ) from litellm.types.agents import AgentResponse @@ -35,8 +40,8 @@ def _group( def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]): - async def load_agent(agent_id: str) -> AgentResponse | None: - return agent + 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) @@ -105,6 +110,88 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): assert ceiling.agent_ids == frozenset() +@dataclass(frozen=True, slots=True) +class _AgentRow: + access_group_ids: Sequence[str] | None + + +class _FakeAgentTable: + def __init__(self, rows: dict[str, _AgentRow], failing: bool = False) -> None: + self._rows: Final = rows + self._failing: Final = failing + self.reads = 0 + + async def find_agent(self, agent_id: str) -> _AgentRow | None: + self.reads += 1 + if self._failing: + raise RuntimeError("db down") + return self._rows.get(agent_id) + + +async def _registry_snapshot(agent_id: str) -> tuple[str, ...]: + return ("registry-group",) + + +@pytest.mark.asyncio +async def test_agent_row_is_read_once_then_served_from_cache(): + cache: Final = DualCache() + table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1", "g2"])}) + + first: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + second: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + + assert (first, second, table.reads) == (("g1", "g2"), ("g1", "g2"), 1) + + +@pytest.mark.asyncio +async def test_agent_with_no_row_or_no_groups_caches_an_empty_answer(): + cache: Final = DualCache() + table: Final = _FakeAgentTable({"bare": _AgentRow(None)}) + + bare: Final = await load_agent_access_group_ids("bare", cache, table.find_agent, _registry_snapshot) + missing: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) + again: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) + + assert (bare, missing, again, table.reads) == ((), (), (), 2) + + +@pytest.mark.asyncio +async def test_evicted_cache_entry_picks_up_the_patched_row(): + cache: Final = DualCache() + rows: Final = {"agent-1": _AgentRow(["g1"])} + table: Final = _FakeAgentTable(rows) + await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + + rows["agent-1"] = _AgentRow(["g2"]) + stale: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + await cache.async_delete_cache(key=agent_access_group_ids_cache_key("agent-1")) + fresh: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + + assert (stale, fresh) == (("g1",), ("g2",)) + + +@pytest.mark.asyncio +async def test_unreadable_row_falls_back_to_the_registry_without_caching(): + cache: Final = DualCache() + table: Final = _FakeAgentTable({}, failing=True) + + answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + + assert answer == ("registry-group",) + assert await cache.async_get_cache(key=agent_access_group_ids_cache_key("agent-1")) is None + + +@pytest.mark.asyncio +async def test_garbage_in_the_cache_is_treated_as_a_miss(): + cache: Final = DualCache() + await cache.async_set_cache(key=agent_access_group_ids_cache_key("agent-1"), value={"not": "a list"}) + table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1"])}) + + answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + + assert (answer, table.reads) == (("g1",), 1) + + @pytest.mark.asyncio async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server 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"]) From 89330cdac6e3b103421114e1385efe719e417534 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:34:06 +0000 Subject: [PATCH 07/13] refactor(agents): exhaust the agent access match and drop routine comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 3 --- .../auth/agent_permission_handler.py | 14 ++++---------- litellm/proxy/auth/auth_checks.py | 5 +---- 3 files changed, 5 insertions(+), 17 deletions(-) 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 05661584a6b..4bca15190cf 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 @@ -193,9 +193,6 @@ def _agent_capped_servers( agent_servers: Sequence[str], agent_access_group_servers: frozenset[str] | None, ) -> tuple[str, ...] | None: - """Servers left once the agent's object_permission and attached access groups both cap the - key/team result, or None when the agent restricts nothing. An attached group set naming no - server is an empty ceiling, not an absent one, so it denies every server.""" if not agent_servers and agent_access_group_servers is None: return None return tuple( diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 1759090a29a..ea73d4634e3 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -8,7 +8,7 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Final, TypeAlias, assert_never from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts @@ -67,14 +67,7 @@ class AgentRequestHandler: 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. - - ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team - carries any grant and the agent behind the key has no access groups attached. - Grants that intersect to nothing stay restricted, so narrowing a caller can - never widen what it reaches. - """ + """Agents the key may reach: key and team grants intersected with the agent's access group ceiling.""" key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) if agent_ceiling is None: @@ -84,6 +77,8 @@ class AgentRequestHandler: return RestrictedAgentAccess(agent_ceiling) case RestrictedAgentAccess(key_team_ids): return RestrictedAgentAccess(key_team_ids & agent_ceiling) + case _: + assert_never(key_team_access) @staticmethod async def _resolve_key_team_agent_access( @@ -111,7 +106,6 @@ class AgentRequestHandler: user_api_key_auth: UserAPIKeyAuth | None, resolve_ceiling: CeilingResolver, ) -> frozenset[str] | None: - """Stable IDs of the agents the calling agent's attached access groups allow; None when none attached.""" 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) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 137cf849389..df20358bcce 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1007,7 +1007,6 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) - # 2.4 If the agent behind the key has access groups attached, they cap the models it can call await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) ## 2.1 If user can call model (if personal key) @@ -4205,9 +4204,7 @@ async def _check_agent_access_group_model_access( llm_router: Router | None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: - """Raises when the key's agent has access groups attached and none of them names the model. - Attached groups that name no model deny every model; ``_can_object_call_model`` would read - an empty allowlist as unrestricted.""" + """Attached groups naming no model deny every model, unlike the empty allowlist ``_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) From 18c31e6fc6cb4f7a5ad5bb658b7f5d6356664922 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:40:06 +0000 Subject: [PATCH 08/13] fix(agents): import assert_never from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/agent_endpoints/auth/agent_permission_handler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index ea73d4634e3..4e022e48bb4 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -8,7 +8,9 @@ Follows the same pattern as MCP permission handling. import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass -from typing import Final, TypeAlias, assert_never +from typing import Final, TypeAlias + +from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts From 1206fa802b851fb485cf494e234f6a628f86813c Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 20:51:04 +0000 Subject: [PATCH 09/13] refactor(agents): resolve attached access groups from the agent registry instead of the DB on the request path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_access_groups.py | 81 +--------------- litellm/proxy/agent_endpoints/endpoints.py | 4 - .../access_group_endpoints.py | 2 - .../auth/test_agent_access_groups.py | 93 +++---------------- 4 files changed, 14 insertions(+), 166 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py index e0e5193d0a4..49e5407ff88 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -1,35 +1,18 @@ import asyncio -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable from dataclasses import dataclass -from typing import Final, Protocol, TypeAlias +from typing import Final, TypeAlias from fastapi import HTTPException -from pydantic import TypeAdapter, ValidationError -from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.caching.dual_cache import DualCache from litellm.proxy._types import LiteLLM_AccessGroupTable -from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl - - -class _AgentAccessGroupsRecord(Protocol): - @property - def access_group_ids(self) -> Sequence[str] | None: ... - - -class _AgentIdWhere(TypedDict): - agent_id: ReadOnly[str] - AccessGroupIds: TypeAlias = tuple[str, ...] AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params -AgentRecordFinder: TypeAlias = Callable[[str], Awaitable[_AgentAccessGroupsRecord | None]] # mutable-ok: Callable LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax -_CACHED_IDS: Final = TypeAdapter(list[str]) - @dataclass(frozen=True, slots=True) class AgentAccessGroupCeiling: @@ -44,19 +27,6 @@ class AgentAccessGroupCeiling: CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params -def agent_access_group_ids_cache_key(agent_id: str) -> str: - return f"agent_access_group_ids:{agent_id}" - - -def _cached_access_group_ids(cached: object) -> AccessGroupIds | None: - if cached is None: - return None - try: - return tuple(_CACHED_IDS.validate_python(cached)) - except ValidationError: - return None - - async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through @@ -64,51 +34,6 @@ async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: return tuple(agent.access_group_ids or ()) if agent is not None else () -async def load_agent_access_group_ids( - agent_id: str, - cache: DualCache, - find_agent: AgentRecordFinder, - fallback: AccessGroupIdsLoader, -) -> AccessGroupIds: - """The agent row's groups, cached for the management-object TTL and evicted on every agent write.""" - cache_key: Final = agent_access_group_ids_cache_key(agent_id) - cached: Final = _cached_access_group_ids(await cache.async_get_cache(key=cache_key)) - if cached is not None: - return cached - try: - record: Final = await find_agent(agent_id) - except Exception as e: # noqa: BLE001 # prisma raises many error types; the registry snapshot answers instead - verbose_proxy_logger.warning("Failed to read access groups for agent %r, using registry: %s", agent_id, e) - return await fallback(agent_id) - access_group_ids: Final = tuple(record.access_group_ids or ()) if record is not None else () - await cache.async_set_cache(key=cache_key, value=access_group_ids, ttl=get_management_object_ttl(cache)) - return access_group_ids - - -async def _load_agent_access_group_ids(agent_id: str) -> AccessGroupIds: - from litellm.proxy.agent_endpoints.agent_registry import agents_table - from litellm.proxy.proxy_server import prisma_client, user_api_key_cache - - if prisma_client is None: - return await _registry_access_group_ids(agent_id) - db: Final = prisma_client - - async def find_agent(row_agent_id: str) -> _AgentAccessGroupsRecord | None: - return await agents_table(db).find_unique(where=_AgentIdWhere(agent_id=row_agent_id)) - - return await load_agent_access_group_ids(agent_id, user_api_key_cache, find_agent, _registry_access_group_ids) - - -async def evict_agent_access_group_ids(agent_ids: Sequence[str]) -> None: - from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast - from litellm.proxy.proxy_server import user_api_key_cache - - await evict_and_broadcast( - cache_keys=tuple(agent_access_group_ids_cache_key(agent_id) for agent_id in agent_ids), - user_api_key_cache=user_api_key_cache, - ) - - 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 @@ -132,7 +57,7 @@ async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: async def resolve_agent_access_group_ceiling( agent_id: str, - load_access_group_ids: AccessGroupIdsLoader = _load_agent_access_group_ids, + 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.""" diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 62783fb412a..aa8979a73c6 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -44,7 +44,6 @@ from litellm.proxy.agent_endpoints.agent_search import ( global_agent_search_index, search_agents, ) -from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -697,7 +696,6 @@ async def update_agent( prisma_client=prisma_client, updated_by=updated_by, ) - await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -801,7 +799,6 @@ async def patch_agent( prisma_client=prisma_client, updated_by=updated_by, ) - await evict_agent_access_group_ids((agent_id,)) # deregister in memory AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) @@ -864,7 +861,6 @@ async def delete_agent( raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.") await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client) - await evict_agent_access_group_ids((agent_id,)) AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name")) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 2694d00b17f..b4923b0a2dc 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -16,7 +16,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.auth.agent_access_groups import evict_agent_access_group_ids from litellm.proxy.auth.auth_checks import ( _cache_access_object, _cache_key_object, @@ -783,7 +782,6 @@ async def delete_access_group( await invalidate_access_group_cache(access_group_id) _detach_access_group_from_agent_registry(detached_agent_ids, access_group_id) - await evict_agent_access_group_ids(detached_agent_ids) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) 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 index 266ce52e3a9..e744e84d671 100644 --- 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 @@ -1,16 +1,11 @@ -from collections.abc import Sequence -from dataclasses import dataclass from typing import Final import pytest from fastapi import HTTPException -from litellm.caching.dual_cache import DualCache from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( AgentAccessGroupCeiling, - agent_access_group_ids_cache_key, - load_agent_access_group_ids, resolve_agent_access_group_ceiling, ) from litellm.types.agents import AgentResponse @@ -110,86 +105,20 @@ async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): assert ceiling.agent_ids == frozenset() -@dataclass(frozen=True, slots=True) -class _AgentRow: - access_group_ids: Sequence[str] | None - - -class _FakeAgentTable: - def __init__(self, rows: dict[str, _AgentRow], failing: bool = False) -> None: - self._rows: Final = rows - self._failing: Final = failing - self.reads = 0 - - async def find_agent(self, agent_id: str) -> _AgentRow | None: - self.reads += 1 - if self._failing: - raise RuntimeError("db down") - return self._rows.get(agent_id) - - -async def _registry_snapshot(agent_id: str) -> tuple[str, ...]: - return ("registry-group",) - - @pytest.mark.asyncio -async def test_agent_row_is_read_once_then_served_from_cache(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1", "g2"])}) +async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - first: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - second: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) + _, 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 (first, second, table.reads) == (("g1", "g2"), ("g1", "g2"), 1) - - -@pytest.mark.asyncio -async def test_agent_with_no_row_or_no_groups_caches_an_empty_answer(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({"bare": _AgentRow(None)}) - - bare: Final = await load_agent_access_group_ids("bare", cache, table.find_agent, _registry_snapshot) - missing: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) - again: Final = await load_agent_access_group_ids("missing", cache, table.find_agent, _registry_snapshot) - - assert (bare, missing, again, table.reads) == ((), (), (), 2) - - -@pytest.mark.asyncio -async def test_evicted_cache_entry_picks_up_the_patched_row(): - cache: Final = DualCache() - rows: Final = {"agent-1": _AgentRow(["g1"])} - table: Final = _FakeAgentTable(rows) - await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - rows["agent-1"] = _AgentRow(["g2"]) - stale: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - await cache.async_delete_cache(key=agent_access_group_ids_cache_key("agent-1")) - fresh: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert (stale, fresh) == (("g1",), ("g2",)) - - -@pytest.mark.asyncio -async def test_unreadable_row_falls_back_to_the_registry_without_caching(): - cache: Final = DualCache() - table: Final = _FakeAgentTable({}, failing=True) - - answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert answer == ("registry-group",) - assert await cache.async_get_cache(key=agent_access_group_ids_cache_key("agent-1")) is None - - -@pytest.mark.asyncio -async def test_garbage_in_the_cache_is_treated_as_a_miss(): - cache: Final = DualCache() - await cache.async_set_cache(key=agent_access_group_ids_cache_key("agent-1"), value={"not": "a list"}) - table: Final = _FakeAgentTable({"agent-1": _AgentRow(["g1"])}) - - answer: Final = await load_agent_access_group_ids("agent-1", cache, table.find_agent, _registry_snapshot) - - assert (answer, table.reads) == (("g1",), 1) + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset() + ) @pytest.mark.asyncio From 7229fc1952a83387e4126949114e9d253d5caec3 Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 21:00:56 +0000 Subject: [PATCH 10/13] fix(agents): return on every branch of the agent access ceiling so CodeQL sees no fall-through Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../agent_endpoints/auth/agent_permission_handler.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index 4e022e48bb4..fe0a1b13b43 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -10,8 +10,6 @@ from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from typing import Final, TypeAlias -from typing_extensions import assert_never - from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts from litellm.proxy._types import ( @@ -74,13 +72,9 @@ class AgentRequestHandler: agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) if agent_ceiling is None: return key_team_access - match key_team_access: - case UnrestrictedAgentAccess(): - return RestrictedAgentAccess(agent_ceiling) - case RestrictedAgentAccess(key_team_ids): - return RestrictedAgentAccess(key_team_ids & agent_ceiling) - case _: - assert_never(key_team_access) + if isinstance(key_team_access, UnrestrictedAgentAccess): + return RestrictedAgentAccess(agent_ceiling) + return RestrictedAgentAccess(key_team_access.agent_ids & agent_ceiling) @staticmethod async def _resolve_key_team_agent_access( From 25729c521fc944b44f0d518f4b6bc45d87e5c68c Mon Sep 17 00:00:00 2001 From: yassin Date: Thu, 17 Sep 2026 21:22:39 +0000 Subject: [PATCH 11/13] refactor(agents): combine key and team agent grants without a fall-through match Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../auth/agent_permission_handler.py | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index fe0a1b13b43..b7b7638e478 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -48,6 +48,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: @@ -83,19 +99,10 @@ class AgentRequestHandler: 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( From d338d3f2d2f6529de70a273b6f0d622708209c9c Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 09:40:46 +0000 Subject: [PATCH 12/13] style(agents): tidy typing and docstring in access group ceiling helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- litellm/proxy/utils.py | 2 +- .../proxy/agent_endpoints/test_agent_registry.py | 5 ++++- tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 5412a9d9f6a..b5e7ef73d36 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -4332,7 +4332,7 @@ async def _check_agent_access_group_model_access( llm_router: Router | None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> Literal[True]: - """Attached groups naming no model deny every model, unlike the empty allowlist ``_can_object_call_model`` allows.""" + """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) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e736f2fa1c4..a8e7d3232eb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -8188,7 +8188,7 @@ async def _get_access_group_models( async def _agent_access_group_visible_models( user_api_key_dict: "UserAPIKeyAuth", - llm_router: Optional["Router"], + llm_router: "Router | None", include_model_access_groups: bool, return_wildcard_routes: bool, team_id: str | None, 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 d15a3adadbd..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. @@ -1052,7 +1053,9 @@ async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_def ({"access_group_ids": None}, []), ], ) -async def test_patch_agent_in_db_replaces_access_group_ids_when_provided(patch_body: dict, expected: list[str]): +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( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index d8cd578d265..a1179617718 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -8900,9 +8900,6 @@ 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 -# Agent access group model ceiling - - def _agent_model_ceiling_resolver( models: frozenset[str] | None, ) -> tuple[CeilingResolver, list[str]]: From 82eef2fcca5e706914a1e81a3d15094f51436208 Mon Sep 17 00:00:00 2001 From: yassin Date: Mon, 21 Sep 2026 22:32:32 +0000 Subject: [PATCH 13/13] fix(proxy): scope agent permissions to invoking caller An agent key that echoes the x-litellm-user-id / x-litellm-team-id headers forwarded by /a2a is capped at that user's and team's models, MCP servers and agents, on top of its own grants and access group ceiling. The echoed ids only narrow, and nested A2A hops forward the original human caller Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 34 ++++- litellm/proxy/_types.py | 11 ++ .../proxy/agent_endpoints/a2a_endpoints.py | 9 +- .../agent_endpoints/auth/agent_caller.py | 87 +++++++++++++ .../auth/agent_permission_handler.py | 19 ++- litellm/proxy/auth/auth_checks.py | 62 ++++++++- litellm/proxy/auth/user_api_key_auth.py | 4 + litellm/types/agents.py | 16 ++- .../auth/test_user_api_key_auth_mcp.py | 84 +++++++++++++ .../agent_endpoints/auth/test_agent_caller.py | 57 +++++++++ .../auth/test_agent_permission_handler.py | 56 ++++++++- .../agent_endpoints/test_a2a_endpoints.py | 19 +++ .../proxy/auth/test_auth_checks.py | 118 ++++++++++++++++++ 13 files changed, 565 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/agent_endpoints/auth/agent_caller.py create mode 100644 tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py 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 4bca15190cf..f1a0de09162 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 @@ -48,6 +48,7 @@ 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 @@ -1577,14 +1578,21 @@ class MCPRequestHandler: "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 @@ -2927,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. diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c28858b48d9..50ddf52fc3f 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, ) @@ -3247,6 +3248,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) @@ -3278,6 +3288,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_source_team_rpm_limits", None) values.pop("mcp_session_resource_server_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): 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/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 b7b7638e478..9fe74bfee3f 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -23,6 +23,7 @@ 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 @@ -83,14 +84,24 @@ class AgentRequestHandler: user_api_key_auth: UserAPIKeyAuth | None = None, resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> AgentAccess: - """Agents the key may reach: key and team grants intersected with the agent's access group ceiling.""" + """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 key_team_access - if isinstance(key_team_access, UnrestrictedAgentAccess): + return own_access + if isinstance(own_access, UnrestrictedAgentAccess): return RestrictedAgentAccess(agent_ceiling) - return RestrictedAgentAccess(key_team_access.agent_ids & agent_ceiling) + return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) + + @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( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b5e7ef73d36..a477eecba38 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 @@ -72,6 +72,11 @@ 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, @@ -1010,6 +1015,14 @@ async def common_checks( ) 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: @@ -4355,6 +4368,53 @@ async def _check_agent_access_group_model_access( ) +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 de0131772bc..b2a71a27090 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, @@ -3320,6 +3321,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/types/agents.py b/litellm/types/agents.py index 12e60352a97..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 @@ -206,6 +206,20 @@ class PatchAgentRequest(TypedDict, total=False): 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 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 032e0f69a46..91c87cbe146 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 @@ -22,6 +22,7 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.types.agents import AgentCaller @pytest.mark.asyncio @@ -4195,6 +4196,89 @@ def test_agent_capped_servers_without_agent_restrictions_is_uncapped(): 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( 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 2a98e6e4feb..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,7 +9,6 @@ from unittest.mock import AsyncMock, patch import pytest - from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry @@ -21,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: @@ -196,6 +196,60 @@ class TestAgentRequestHandler: 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"})) 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/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index a1179617718..7d5cf7dad9f 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -35,6 +35,7 @@ from litellm.proxy._types import ( 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, @@ -45,11 +46,14 @@ from litellm.proxy.auth.auth_checks import ( _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, @@ -9119,3 +9123,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 == []