mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
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:
parent
d338d3f2d2
commit
82eef2fcca
13 changed files with 565 additions and 11 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
87
litellm/proxy/agent_endpoints/auth/agent_caller.py
Normal file
87
litellm/proxy/agent_endpoints/auth/agent_caller.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""The human behind an agent's own proxy calls.
|
||||
|
||||
``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the
|
||||
agent backend. When the agent echoes them back on requests made with its own key, the proxy caps
|
||||
that key at what the invoking user and team may reach. The cap is intersected with, never
|
||||
substituted for, the agent key's own grants and the agent's access group ceiling, so the headers
|
||||
can only narrow access and need no trust.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth
|
||||
from litellm.types.agents import (
|
||||
AGENT_CALLER_TEAM_ID_HEADER,
|
||||
AGENT_CALLER_USER_ID_HEADER,
|
||||
AgentCaller,
|
||||
)
|
||||
|
||||
|
||||
def _header(headers: Mapping[str, str], name: str) -> str | None:
|
||||
value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None)
|
||||
return value.strip() or None if value is not None else None
|
||||
|
||||
|
||||
def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None:
|
||||
"""The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed."""
|
||||
if not user_api_key_auth.agent_id:
|
||||
return None
|
||||
user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER)
|
||||
team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER)
|
||||
if user_id is None and team_id is None:
|
||||
return None
|
||||
return AgentCaller(user_id=user_id, team_id=team_id)
|
||||
|
||||
|
||||
def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None:
|
||||
"""A minimal auth context standing for the invoking user and team, so the shared key/team/user
|
||||
resolvers can be reused unchanged to compute what the caller may reach."""
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None:
|
||||
return None
|
||||
return UserAPIKeyAuth(
|
||||
user_id=caller.user_id,
|
||||
team_id=caller.team_id,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
)
|
||||
|
||||
|
||||
async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None:
|
||||
"""The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team
|
||||
that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted."""
|
||||
from litellm.proxy.auth.auth_checks import get_team_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None or caller.team_id is None:
|
||||
return None
|
||||
return await get_team_object(
|
||||
team_id=caller.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None:
|
||||
"""The invoking user's row, or ``None`` when no user id was echoed or the row does not exist."""
|
||||
from litellm.proxy.auth.auth_checks import get_user_object
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
caller: Final = user_api_key_auth.agent_caller
|
||||
if caller is None or caller.user_id is None:
|
||||
return None
|
||||
user_object: Final = await get_user_object(
|
||||
user_id=caller.user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
parent_otel_span=user_api_key_auth.parent_otel_span,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
if user_object is None:
|
||||
verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id)
|
||||
return user_object
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -0,0 +1,57 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth, agent_caller_from_headers
|
||||
from litellm.types.agents import AgentCaller
|
||||
|
||||
_AGENT_KEY: Final = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1")
|
||||
|
||||
|
||||
def test_agent_key_echoing_both_ids_acts_for_that_user_and_team() -> None:
|
||||
headers: Final = {"X-LiteLLM-User-Id": " alice ", "x-litellm-team-id": "callers"}
|
||||
|
||||
assert agent_caller_from_headers(headers, _AGENT_KEY) == AgentCaller(user_id="alice", team_id="callers")
|
||||
|
||||
|
||||
def test_agent_key_echoing_only_a_user_id_acts_for_a_teamless_user() -> None:
|
||||
assert agent_caller_from_headers({"x-litellm-user-id": "alice"}, _AGENT_KEY) == AgentCaller(user_id="alice")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("headers", [{}, {"x-litellm-user-id": " ", "x-litellm-team-id": ""}])
|
||||
def test_agent_key_echoing_no_caller_acts_for_itself(headers: dict[str, str]) -> None:
|
||||
assert agent_caller_from_headers(headers, _AGENT_KEY) is None
|
||||
|
||||
|
||||
def test_caller_headers_on_a_key_without_an_agent_are_ignored() -> None:
|
||||
plain_key: Final = UserAPIKeyAuth(api_key="plain-key", user_id="bob")
|
||||
|
||||
assert agent_caller_from_headers({"x-litellm-user-id": "alice", "x-litellm-team-id": "callers"}, plain_key) is None
|
||||
|
||||
|
||||
def test_caller_auth_stands_for_the_invoking_user_not_the_agent() -> None:
|
||||
agent_key: Final = UserAPIKeyAuth(
|
||||
api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1"
|
||||
)
|
||||
agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers")
|
||||
|
||||
caller_auth: Final = agent_caller_auth(agent_key)
|
||||
|
||||
assert caller_auth is not None
|
||||
assert (caller_auth.user_id, caller_auth.team_id, caller_auth.agent_id, caller_auth.api_key) == (
|
||||
"alice",
|
||||
"callers",
|
||||
None,
|
||||
None,
|
||||
)
|
||||
assert agent_caller_auth(_AGENT_KEY) is None
|
||||
|
||||
|
||||
def test_agent_caller_cannot_be_set_from_a_request_payload() -> None:
|
||||
forged: Final = UserAPIKeyAuth.model_validate(
|
||||
{"api_key": "agent-key", "agent_id": "agent-1", "agent_caller": {"user_id": "alice", "team_id": "callers"}}
|
||||
)
|
||||
|
||||
assert forged.agent_caller is None
|
||||
assert "agent_caller" not in forged.model_dump()
|
||||
|
|
@ -9,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"}))
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue