Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/agent_endpoints/auth/agent_caller.py: 34%

1"""The human behind an agent's own proxy calls. 

2 

3``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the 

4agent backend. When the agent echoes them back on requests made with its own key, the proxy caps 

5that key at what the invoking user and team may reach. The cap is intersected with, never 

6substituted for, the agent key's own grants and the agent's access group ceiling, so the headers 

7can only narrow access and need no trust. 

8""" 

9 

10from collections.abc import Mapping 

11from typing import Final 

12 

13from litellm._logging import verbose_proxy_logger 

14from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth 

15from litellm.types.agents import ( 

16 AGENT_CALLER_TEAM_ID_HEADER, 

17 AGENT_CALLER_USER_ID_HEADER, 

18 AgentCaller, 

19) 

20 

21 

22def _header(headers: Mapping[str, str], name: str) -> str | None: 

23 value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None) 

24 return value.strip() or None if value is not None else None 

25 

26 

27def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None: 

28 """The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed.""" 

29 if not user_api_key_auth.agent_id: 29 ↛ 31line 29 didn't jump to line 31 because the condition on line 29 was always true

30 return None 

31 user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER) 

32 team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER) 

33 if user_id is None and team_id is None: 

34 return None 

35 return AgentCaller(user_id=user_id, team_id=team_id) 

36 

37 

38def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None: 

39 """A minimal auth context standing for the invoking user and team, so the shared key/team/user 

40 resolvers can be reused unchanged to compute what the caller may reach.""" 

41 caller: Final = user_api_key_auth.agent_caller 

42 if caller is None: 42 ↛ 44line 42 didn't jump to line 44 because the condition on line 42 was always true

43 return None 

44 return UserAPIKeyAuth( 

45 user_id=caller.user_id, 

46 team_id=caller.team_id, 

47 parent_otel_span=user_api_key_auth.parent_otel_span, 

48 ) 

49 

50 

51async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: 

52 """The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team 

53 that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted.""" 

54 from litellm.proxy.auth.auth_checks import get_team_object 

55 from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache 

56 

57 caller: Final = user_api_key_auth.agent_caller 

58 if caller is None or caller.team_id is None: 

59 return None 

60 return await get_team_object( 

61 team_id=caller.team_id, 

62 prisma_client=prisma_client, 

63 user_api_key_cache=user_api_key_cache, 

64 parent_otel_span=user_api_key_auth.parent_otel_span, 

65 proxy_logging_obj=proxy_logging_obj, 

66 ) 

67 

68 

69async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None: 

70 """The invoking user's row, or ``None`` when no user id was echoed or the row does not exist.""" 

71 from litellm.proxy.auth.auth_checks import get_user_object 

72 from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache 

73 

74 caller: Final = user_api_key_auth.agent_caller 

75 if caller is None or caller.user_id is None: 

76 return None 

77 user_object: Final = await get_user_object( 

78 user_id=caller.user_id, 

79 prisma_client=prisma_client, 

80 user_api_key_cache=user_api_key_cache, 

81 user_id_upsert=False, 

82 parent_otel_span=user_api_key_auth.parent_otel_span, 

83 proxy_logging_obj=proxy_logging_obj, 

84 ) 

85 if user_object is None: 

86 verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id) 

87 return user_object