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>
This commit is contained in:
yassin 2026-09-21 22:32:32 +00:00
parent d338d3f2d2
commit 82eef2fcca
13 changed files with 565 additions and 11 deletions

View file

@ -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.

View file

@ -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):

View file

@ -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
}

View file

@ -0,0 +1,87 @@
"""The human behind an agent's own proxy calls.
``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the
agent backend. When the agent echoes them back on requests made with its own key, the proxy caps
that key at what the invoking user and team may reach. The cap is intersected with, never
substituted for, the agent key's own grants and the agent's access group ceiling, so the headers
can only narrow access and need no trust.
"""
from collections.abc import Mapping
from typing import Final
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth
from litellm.types.agents import (
AGENT_CALLER_TEAM_ID_HEADER,
AGENT_CALLER_USER_ID_HEADER,
AgentCaller,
)
def _header(headers: Mapping[str, str], name: str) -> str | None:
value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None)
return value.strip() or None if value is not None else None
def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None:
"""The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed."""
if not user_api_key_auth.agent_id:
return None
user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER)
team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER)
if user_id is None and team_id is None:
return None
return AgentCaller(user_id=user_id, team_id=team_id)
def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None:
"""A minimal auth context standing for the invoking user and team, so the shared key/team/user
resolvers can be reused unchanged to compute what the caller may reach."""
caller: Final = user_api_key_auth.agent_caller
if caller is None:
return None
return UserAPIKeyAuth(
user_id=caller.user_id,
team_id=caller.team_id,
parent_otel_span=user_api_key_auth.parent_otel_span,
)
async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
"""The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team
that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted."""
from litellm.proxy.auth.auth_checks import get_team_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
caller: Final = user_api_key_auth.agent_caller
if caller is None or caller.team_id is None:
return None
return await get_team_object(
team_id=caller.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None:
"""The invoking user's row, or ``None`` when no user id was echoed or the row does not exist."""
from litellm.proxy.auth.auth_checks import get_user_object
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
caller: Final = user_api_key_auth.agent_caller
if caller is None or caller.user_id is None:
return None
user_object: Final = await get_user_object(
user_id=caller.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if user_object is None:
verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id)
return user_object

View file

@ -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(

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -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()

View file

@ -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"}))

View file

@ -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):

View file

@ -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 == []