feat(agents): attach access groups to agents and enforce them for models, MCP servers and agent calls

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[];

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -73,6 +73,7 @@ model LiteLLM_AgentsTable {
   static_headers       Json?          @default("{}")
   extra_headers        String[]       @default([])
   agent_access_groups String[]       @default([])
+  access_group_ids     String[]       @default([])
   object_permission_id String?
   object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
   spend      Float    @default(0.0)

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -44,6 +44,11 @@
     UserAPIKeyAuth,
     user_api_key_has_admin_view,
 )
+from litellm.proxy.agent_endpoints.auth.agent_access_groups import (
+    CeilingResolver,
+    resolve_agent_access_group_ceiling,
+)
+from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.auth.user_api_key_auth import (
     _get_bearer_token_or_received_api_key,  # pyright: ignore[reportPrivateUsage]  # shared x-litellm-api-key parser lives with user_api_key_auth
@@ -184,6 +189,21 @@ def _has_client_supplied_mcp_auth(
     return bool(mcp_auth_header) or bool(mcp_server_auth_headers)
 
 
+def _agent_capped_servers(
+    allowed_mcp_servers: Sequence[str],
+    agent_servers: Sequence[str],
+    agent_access_group_servers: frozenset[str] | None,
+) -> tuple[str, ...] | None:
+    if not agent_servers and agent_access_group_servers is None:
+        return None
+    return tuple(
+        s
+        for s in allowed_mcp_servers
+        if (not agent_servers or s in agent_servers)
+        and (agent_access_group_servers is None or s in agent_access_group_servers)
+    )
+
+
 def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool:
     """True when this auth is a keyless subject admitted by the gateway session / bridge user
     path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``.
@@ -1546,25 +1566,33 @@ async def get_mcp_server_access(
             # Check agent permissions if agent_id is set on the key
             #########################################################
             if user_api_key_auth and user_api_key_auth.agent_id:
-                allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent(
-                    user_api_key_auth
+                agent_capped: Final = _agent_capped_servers(
+                    allowed_mcp_servers,
+                    await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth),
+                    await MCPRequestHandler._get_agent_access_group_server_ceiling(user_api_key_auth),
                 )
-                if len(allowed_mcp_servers_for_agent) > 0:
+                if agent_capped is not None:
                     has_lower_level_mcp_restrictions = True
-                    # Intersect: agent can only use servers allowed by BOTH key/team AND agent config
-                    allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent]
+                    allowed_mcp_servers = list(agent_capped)
                     verbose_logger.debug(
                         "Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers
                     )
 
+            #########################################################
+            # Cap an agent key at what the user and team that invoked the agent may reach
+            #########################################################
+            caller_capped, caller_restricts = await MCPRequestHandler._apply_agent_caller_ceiling(
+                allowed_mcp_servers, user_api_key_auth
+            )
+
             #########################################################
             # Apply the internal user's own ceiling (the entitlement attached to the human)
             #########################################################
             capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(
-                allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source
+                caller_capped, user_api_key_auth, keyless_source=keyless_source
             )
             allowed_mcp_servers = list(capped)
-            has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts
+            has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or caller_restricts or user_restricts
 
             #########################################################
             # Apply org-level ceiling if org_id is set
@@ -2907,6 +2935,28 @@ async def _apply_user_server_ceiling(
         verbose_logger.debug("Applied user ceiling filter. Final allowed servers: %s", capped)
         return capped, True
 
+    @staticmethod
+    async def _apply_agent_caller_ceiling(
+        allowed_mcp_servers: Sequence[str],
+        user_api_key_auth: UserAPIKeyAuth | None = None,
+    ) -> tuple[tuple[str, ...], bool]:
+        """Narrow an agent key's servers to those the invoking user and team (echoed back by the agent
+        as ``x-litellm-user-id`` / ``x-litellm-team-id``) may reach: the echoed team's grants when it
+        names any, then the echoed user's own entitlement. Raises like the user ceiling when that
+        entitlement is known but unreadable, so the resolver denies rather than widens."""
+        caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None
+        if caller_auth is None:
+            return tuple(allowed_mcp_servers), False
+        team_servers: Final = frozenset(await MCPRequestHandler._get_allowed_mcp_servers_for_team(caller_auth))
+        team_capped: Final = (
+            tuple(server for server in allowed_mcp_servers if server in team_servers)
+            if team_servers
+            else tuple(allowed_mcp_servers)
+        )
+        user_capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(team_capped, caller_auth)
+        verbose_logger.debug("Applied agent caller ceiling. Final allowed servers: %s", user_capped)
+        return user_capped, bool(team_servers) or user_restricts
+
     @staticmethod
     async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool:
         """Whether this human's own entitlement bounds their MCP access at all.
@@ -3137,6 +3187,27 @@ async def _get_allowed_mcp_servers_for_agent(
             verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e)
             return []
 
+    @staticmethod
+    async def _get_agent_access_group_server_ceiling(
+        user_api_key_auth: UserAPIKeyAuth,
+        resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling,
+    ) -> frozenset[str] | None:
+        """
+        Server IDs the agent's attached unified access groups (``LiteLLM_AgentsTable.access_group_ids``)
+        allow, or None when the agent has none attached. Unlike the object_permission path above, an
+        attached group set that names no servers is an empty ceiling and denies every server.
+        """
+        from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+            global_mcp_server_manager,
+        )
+
+        if not user_api_key_auth.agent_id:
+            return None
+        ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id)
+        if ceiling is None:
+            return None
+        return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids)))
+
     @staticmethod
     async def _get_agent_tool_permissions_for_server(
         server_id: str,

litellm/proxy/_lazy_openapi_snapshot.json

@@ -2357,6 +2357,20 @@
         },
         "AgentConfig": {
           "properties": {
+            "access_group_ids": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Access Group Ids"
+            },
             "agent_card_params": {
               "$ref": "#/components/schemas/AgentCard"
             },
@@ -2683,6 +2697,20 @@
         },
         "AgentResponse": {
           "properties": {
+            "access_group_ids": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Access Group Ids"
+            },
             "agent_card_params": {
               "additionalProperties": true,
               "title": "Agent Card Params",
@@ -3506,6 +3534,20 @@
         },
         "PatchAgentRequest": {
           "properties": {
+            "access_group_ids": {
+              "anyOf": [
+                {
+                  "items": {
+                    "type": "string"
+                  },
+                  "type": "array"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "title": "Access Group Ids"
+            },
             "agent_card_params": {
               "$ref": "#/components/schemas/AgentCard"
             },

litellm/proxy/_types.py

@@ -27,6 +27,7 @@
     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 @@ def check_api_key(cls, values):
         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):
@@ -4227,6 +4238,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"
@@ -4301,7 +4317,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
@@ -4316,6 +4332,8 @@ def get_model_access_error_type_for_object(
             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(

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
         }

litellm/proxy/agent_endpoints/agent_registry.py

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict
 
 from pydantic import TypeAdapter, ValidationError
+from typing_extensions import ReadOnly
 
 import litellm
 from litellm.constants import REDACTED_BY_LITELM_STRING
@@ -37,6 +38,7 @@ class AgentRecordDump(TypedDict):
     agent_card_params: dict[str, object]
     static_headers: dict[str, str] | None
     extra_headers: list[str] | None
+    access_group_ids: ReadOnly[Sequence[str] | None]
     object_permission: dict[str, object] | None
     spend: float
     tpm_limit: int | None
@@ -65,6 +67,9 @@ def object_permission_id(self) -> str | None: ...
     @property
     def object_permission(self) -> AgentObjectPermissionRecord | None: ...
 
+    @property
+    def access_group_ids(self) -> Sequence[str] | None: ...
+
     @property
     def spend(self) -> float: ...
 
@@ -284,6 +289,12 @@ def _resolved_agent_param_value(
     return _MISSING_AGENT_PARAM
 
 
+def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]:
+    if "access_group_ids" not in agent:
+        return MappingProxyType({})
+    return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))})
+
+
 def _restore_redacted_litellm_params(
     incoming: Mapping[str, object],
     existing: Mapping[str, object],
@@ -516,6 +527,7 @@ async def add_agent_to_db(
             static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None
 
             extra_headers_val: Final = agent.get("extra_headers")
+            access_group_ids_val: Final = agent.get("access_group_ids")
 
             create_data: Final[dict[str, object]] = {
                 "agent_name": agent_name,
@@ -532,6 +544,8 @@ async def add_agent_to_db(
                 create_data["static_headers"] = static_headers_val
             if extra_headers_val is not None:
                 create_data["extra_headers"] = extra_headers_val
+            if access_group_ids_val is not None:
+                create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val))
             if object_permission_id is not None:
                 create_data["object_permission_id"] = object_permission_id
 
@@ -601,7 +615,7 @@ async def patch_agent_in_db(
             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,13 +717,15 @@ async def update_agent_in_db(
                 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,
                 "litellm_params": litellm_params,
                 "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),
             }