Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/ui_session_utils.py: 35%

1"""Helpers to resolve the identity a dashboard UI session token acts as.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Awaitable, Callable 

6from typing import Final 

7 

8from fastapi import HTTPException 

9 

10from litellm._logging import verbose_logger 

11from litellm.constants import UI_SESSION_TOKEN_TEAM_ID 

12from litellm.proxy._types import UserAPIKeyAuth 

13 

14 

15def clone_user_api_key_auth_with_team( 

16 user_api_key_auth: UserAPIKeyAuth, 

17 team_id: str, 

18) -> UserAPIKeyAuth: 

19 """Return a deep copy of the auth context with a different team id.""" 

20 

21 try: 

22 cloned_auth = user_api_key_auth.model_copy() 

23 except AttributeError: 

24 cloned_auth = user_api_key_auth.copy() 

25 cloned_auth.team_id = team_id 

26 return cloned_auth 

27 

28 

29def is_ui_session_credential(user_api_key_auth: UserAPIKeyAuth) -> bool: 

30 """Whether the caller is the dashboard's SSO-minted session token acting as its user, 

31 the only credential shape allowed to widen a request to the owning user's identity.""" 

32 

33 return user_api_key_auth.team_id == UI_SESSION_TOKEN_TEAM_ID and bool(user_api_key_auth.user_id) 

34 

35 

36async def resolve_ui_session_team_ids( 

37 user_api_key_auth: UserAPIKeyAuth, 

38) -> list[str]: 

39 """Resolve the real team ids backing a UI session token.""" 

40 

41 if not is_ui_session_credential(user_api_key_auth): 41 ↛ 44line 41 didn't jump to line 44 because the condition on line 41 was always true

42 return [] 

43 

44 from litellm.proxy.auth.auth_checks import get_user_object 

45 from litellm.proxy.proxy_server import ( 

46 prisma_client, 

47 proxy_logging_obj, 

48 user_api_key_cache, 

49 ) 

50 

51 if prisma_client is None: 

52 verbose_logger.debug("Cannot resolve UI session team ids without DB access") 

53 return [] 

54 

55 try: 

56 user_obj: Final = await get_user_object( 

57 user_id=user_api_key_auth.user_id, 

58 prisma_client=prisma_client, 

59 user_api_key_cache=user_api_key_cache, 

60 user_id_upsert=False, 

61 parent_otel_span=user_api_key_auth.parent_otel_span, 

62 proxy_logging_obj=proxy_logging_obj, 

63 ) 

64 except Exception as exc: # pragma: no cover - defensive logging 

65 verbose_logger.warning( 

66 "Failed to load teams for UI session token user.", 

67 exc, 

68 ) 

69 return [] 

70 

71 if user_obj is None or not user_obj.teams: 

72 return [] 

73 

74 resolved_team_ids: Final[list[str]] = [] 

75 for team_id in user_obj.teams: 

76 if team_id and team_id not in resolved_team_ids: 

77 resolved_team_ids.append(team_id) 

78 return resolved_team_ids 

79 

80 

81async def admitted_user_context(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None: 

82 """THE owner of "resolve this dashboard session's user identity": the same admitted-subject auth a 

83 gateway OAuth session for this user resolves with, carrying the user row's own object permission, 

84 on this request's tracing span. None for any other credential (a caller-passed key is never 

85 widened) and on reload failure, which every caller reads as "no user-level identity available".""" 

86 

87 user_id: Final = user_api_key_auth.user_id 

88 if not is_ui_session_credential(user_api_key_auth) or user_id is None: 88 ↛ 90line 88 didn't jump to line 90 because the condition on line 88 was always true

89 return None 

90 from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( 

91 MCPRequestHandler, 

92 ) 

93 

94 try: 

95 admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) 

96 except HTTPException as e: 

97 verbose_logger.warning("MCP dashboard session: admitted-subject reload failed for %s: %s", user_id, e.detail) 

98 return None 

99 return admitted.model_copy(update={"parent_otel_span": user_api_key_auth.parent_otel_span}) 

100 

101 

102async def acting_user_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth: 

103 """The principal acting-as-user MCP routes resolve permissions with. A non-admin dashboard 

104 session acts as the admitted subject, the same identity a gateway session resolves with, so 

105 server reachability, per-source tool ceilings, rate limits, and billing bind identically on 

106 both surfaces. An admin session keeps its operator view and any caller-passed credential is 

107 returned unchanged, never widened. 

108 

109 Do not combine this with a narrowing that rewrites a single credential's ``object_permission`` 

110 (toolset scope): the admitted subject resolves per grant source and a team source deliberately 

111 carries none of the caller's own grants, so the narrowing would silently evaporate on every 

112 team-granted server. A request carrying such a scope keeps the caller's own credential.""" 

113 

114 if not is_ui_session_credential(user_api_key_auth): 114 ↛ 116line 114 didn't jump to line 116 because the condition on line 114 was always true

115 return user_api_key_auth 

116 from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view 

117 

118 if _user_has_admin_view(user_api_key_auth): 

119 return user_api_key_auth 

120 admitted: Final = await admitted_user_context(user_api_key_auth) 

121 return admitted if admitted is not None else user_api_key_auth 

122 

123 

124async def build_effective_auth_contexts( 

125 user_api_key_auth: UserAPIKeyAuth, 

126) -> list[UserAPIKeyAuth]: 

127 """Every auth context a management or listing surface must resolve a UI session token through: 

128 one per real team backing the session, plus the session user's own admitted identity, so a grant 

129 made directly to the user row is as visible to the dashboard as it is to a gateway session.""" 

130 

131 resolved_team_ids: Final = await resolve_ui_session_team_ids(user_api_key_auth) 

132 team_contexts: Final = ( 

133 [clone_user_api_key_auth_with_team(user_api_key_auth, team_id) for team_id in resolved_team_ids] 

134 if resolved_team_ids 

135 else [user_api_key_auth] 

136 ) 

137 admitted_context: Final = await admitted_user_context(user_api_key_auth) 

138 if admitted_context is None: 138 ↛ 140line 138 didn't jump to line 140 because the condition on line 138 was always true

139 return team_contexts 

140 return [*team_contexts, admitted_context] 

141 

142 

143async def can_access_mcp_server( 

144 user_api_key_auth: UserAPIKeyAuth, 

145 server_id: str, 

146 allowed_servers: Callable[[UserAPIKeyAuth], Awaitable[list[str]]], 

147) -> bool: 

148 """Resolve server access through the same credential contexts as MCP management.""" 

149 for context in await build_effective_auth_contexts(user_api_key_auth): 

150 if server_id in await allowed_servers(context): 

151 return True 

152 return False