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