Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/auth/resolvers/grants.py: 48%

1"""Load a caller's user row, team row, and team membership from the database and validate them together. 

2 

3The virtual-key path reads these off the combined-view SQL join. Every other credential (an IdP JWT, a 

4``lite login`` session token) carries only identifiers, or a snapshot of grants taken when it was minted, so 

5it has to read the live rows on each request. Both of those paths resolve the same rows with the same 

6membership rule, and ``GrantResolver`` is the one place that rule lives. 

7""" 

8 

9from __future__ import annotations 

10 

11from collections.abc import Coroutine, Mapping 

12from dataclasses import dataclass 

13from types import MappingProxyType 

14from typing import TYPE_CHECKING, Final, NoReturn, Protocol, TypeAlias 

15 

16from fastapi import HTTPException, status 

17from pydantic import BaseModel, ValidationError 

18from pydantic.main import IncEx 

19from typing_extensions import assert_never 

20 

21from litellm._logging import verbose_proxy_logger 

22from litellm.proxy._types import ( 

23 LiteLLM_TeamMembership, 

24 LiteLLM_TeamTableCachedObj, 

25 LiteLLM_UserTable, 

26 ProxyErrorTypes, 

27 ProxyException, 

28) 

29from litellm.proxy.auth.auth_checks import ( 

30 TeamNotFoundError, 

31 get_team_membership, 

32 get_team_object, 

33 get_user_object, 

34) 

35from litellm.types.proxy.auth.auth_checks import UserNotFoundError 

36 

37if TYPE_CHECKING: 37 ↛ 38line 37 didn't jump to line 38 because the condition on line 37 was never true

38 from litellm.proxy._types import Span 

39 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache 

40 from litellm.proxy.utils import PrismaClient, ProxyLogging 

41 

42 

43class UserLoader(Protocol): 

44 def __call__( 44 ↛ exitline 44 didn't return from function '__call__' because

45 self, 

46 *, 

47 user_id: str | None, 

48 prisma_client: PrismaClient | None, 

49 user_api_key_cache: UserApiKeyCache, 

50 user_id_upsert: bool, 

51 parent_otel_span: Span | None, 

52 proxy_logging_obj: ProxyLogging | None, 

53 sso_user_id: str | None, 

54 user_email: str | None, 

55 ) -> Coroutine[object, object, LiteLLM_UserTable | None]: ... 

56 

57 

58class TeamLoader(Protocol): 

59 def __call__( 59 ↛ exitline 59 didn't return from function '__call__' because

60 self, 

61 *, 

62 team_id: str, 

63 prisma_client: PrismaClient | None, 

64 user_api_key_cache: UserApiKeyCache, 

65 parent_otel_span: Span | None, 

66 proxy_logging_obj: ProxyLogging | None, 

67 ) -> Coroutine[object, object, LiteLLM_TeamTableCachedObj]: ... 

68 

69 

70class MembershipLoader(Protocol): 

71 def __call__( 71 ↛ exitline 71 didn't return from function '__call__' because

72 self, 

73 *, 

74 user_id: str, 

75 team_id: str, 

76 prisma_client: PrismaClient | None, 

77 user_api_key_cache: UserApiKeyCache, 

78 parent_otel_span: Span | None, 

79 proxy_logging_obj: ProxyLogging | None, 

80 ) -> Coroutine[object, object, LiteLLM_TeamMembership | None]: ... 

81 

82 

83@dataclass(frozen=True, slots=True) 

84class UserLookup: 

85 """The user a credential names, plus the hints ``get_user_object`` may fall back to when the id alone 

86 matches no row.""" 

87 

88 user_id: str | None 

89 user_email: str | None = None 

90 sso_user_id: str | None = None 

91 upsert: bool = False 

92 

93 

94@dataclass(frozen=True, slots=True) 

95class ResolvedGrants: 

96 """The live rows behind a credential. ``effective_user_id`` is the DB row's id when a fuzzy match found a 

97 legacy row under a different id (GH #26789), otherwise the id the credential named.""" 

98 

99 user_object: LiteLLM_UserTable | None 

100 team_object: LiteLLM_TeamTableCachedObj | None 

101 team_membership: LiteLLM_TeamMembership | None 

102 effective_user_id: str | None 

103 

104 

105@dataclass(frozen=True, slots=True) 

106class UserGone: 

107 user_id: str 

108 

109 

110@dataclass(frozen=True, slots=True) 

111class TeamGone: 

112 team_id: str 

113 

114 

115@dataclass(frozen=True, slots=True) 

116class NotAMember: 

117 user_id: str 

118 team_id: str 

119 

120 

121@dataclass(frozen=True, slots=True) 

122class LookupDegraded: 

123 """A row could not be read for a reason that says nothing about the caller: the database is down or a 

124 loader failed. The caller decides whether a grant it already holds may stand in.""" 

125 

126 error: Exception 

127 

128 

129GrantDenial: TypeAlias = UserGone | TeamGone | NotAMember 

130GrantOutcome: TypeAlias = ResolvedGrants | GrantDenial | LookupDegraded 

131 

132 

133_MODELS_COLUMN: Final[Mapping[str, IncEx | bool]] = MappingProxyType({"models": True}) 

134 

135 

136class _UserModelColumn(BaseModel): 

137 """``LiteLLM_UserTable.models`` is a bare ``list``; re-read it with the shape a token's ``models`` takes.""" 

138 

139 models: tuple[str, ...] = () 

140 

141 

142def user_models(user_object: LiteLLM_UserTable) -> tuple[str, ...]: 

143 try: 

144 return _UserModelColumn.model_validate(user_object.model_dump(include=_MODELS_COLUMN)).models 

145 except ValidationError: 

146 return () 

147 

148 

149def canonical_user_id(user_id: str | None, user_object: LiteLLM_UserTable | None) -> str | None: 

150 if user_object is not None and user_object.user_id: 

151 return user_object.user_id 

152 return user_id 

153 

154 

155def raise_public(denial: GrantDenial) -> NoReturn: 

156 match denial: 

157 case UserGone(user_id=user_id): 

158 raise ProxyException( 

159 message=f"Authentication Error, user '{user_id}' no longer exists.", 

160 type=ProxyErrorTypes.auth_error, 

161 param="user_id", 

162 code=status.HTTP_401_UNAUTHORIZED, 

163 ) 

164 case TeamGone(team_id=team_id): 

165 raise TeamNotFoundError(team_id=team_id) 

166 case NotAMember(team_id=team_id): 

167 raise HTTPException( 

168 status_code=status.HTTP_403_FORBIDDEN, 

169 detail=f"Team '{team_id}' is not in your team memberships.", 

170 ) 

171 case _: 

172 assert_never(denial) 

173 

174 

175class GrantResolver: 

176 """Reads the user, membership, and team rows for a credential through injected loaders. 

177 

178 The loaders default to the shared ``auth_checks`` readers. A caller passes its own module's names for them 

179 so the reads stay interceptable where that module's callers already intercept them. ``resolve_identity`` 

180 is the JWT half: user and membership only, since the JWT builder selects the team itself and lets loader 

181 errors surface as they are. ``resolve`` also reads the team row and applies the membership rule, which is 

182 what a credential carrying a grant snapshot needs to refresh it. 

183 """ 

184 

185 def __init__( 

186 self, 

187 prisma_client: PrismaClient | None, 

188 cache: UserApiKeyCache, 

189 *, 

190 parent_otel_span: Span | None = None, 

191 proxy_logging_obj: ProxyLogging | None = None, 

192 load_user: UserLoader = get_user_object, 

193 load_team: TeamLoader = get_team_object, 

194 load_membership: MembershipLoader = get_team_membership, 

195 ) -> None: 

196 self._prisma = prisma_client 

197 self._cache = cache 

198 self._parent_otel_span = parent_otel_span 

199 self._proxy_logging_obj = proxy_logging_obj 

200 self._load_user = load_user 

201 self._load_team = load_team 

202 self._load_membership = load_membership 

203 

204 async def resolve_identity( 

205 self, lookup: UserLookup, team_id: str | None 

206 ) -> tuple[LiteLLM_UserTable | None, LiteLLM_TeamMembership | None, str | None]: 

207 user_object: Final = await self._user(lookup) if lookup.user_id else None 

208 effective_user_id: Final = canonical_user_id(lookup.user_id, user_object) 

209 if effective_user_id != lookup.user_id: 

210 verbose_proxy_logger.debug( 

211 "Auth: rebinding user_id %r -> DB user_id %r (email/sso match)", 

212 lookup.user_id, 

213 effective_user_id, 

214 ) 

215 membership: Final = ( 

216 await self._membership(user_id=effective_user_id, team_id=team_id) 

217 if effective_user_id and team_id 

218 else None 

219 ) 

220 return user_object, membership, effective_user_id 

221 

222 async def resolve(self, lookup: UserLookup, team_id: str | None) -> GrantOutcome: 

223 try: 

224 user_object, membership, effective_user_id = await self.resolve_identity(lookup, team_id) 

225 except UserNotFoundError: 

226 return UserGone(user_id=lookup.user_id or "") 

227 except Exception as error: 

228 return LookupDegraded(error=error) 

229 if team_id is None: 

230 return ResolvedGrants(user_object, None, membership, effective_user_id) 

231 if user_object is not None and team_id not in user_object.teams: 

232 return NotAMember(user_id=user_object.user_id, team_id=team_id) 

233 try: 

234 team_object: Final = await self._load_team( 

235 team_id=team_id, 

236 prisma_client=self._prisma, 

237 user_api_key_cache=self._cache, 

238 parent_otel_span=self._parent_otel_span, 

239 proxy_logging_obj=self._proxy_logging_obj, 

240 ) 

241 except TeamNotFoundError: 

242 return TeamGone(team_id=team_id) 

243 except Exception as error: 

244 return LookupDegraded(error=error) 

245 return ResolvedGrants(user_object, team_object, membership, effective_user_id) 

246 

247 async def _user(self, lookup: UserLookup) -> LiteLLM_UserTable | None: 

248 return await self._load_user( 

249 user_id=lookup.user_id, 

250 prisma_client=self._prisma, 

251 user_api_key_cache=self._cache, 

252 user_id_upsert=lookup.upsert, 

253 parent_otel_span=self._parent_otel_span, 

254 proxy_logging_obj=self._proxy_logging_obj, 

255 user_email=lookup.user_email, 

256 sso_user_id=lookup.sso_user_id, 

257 ) 

258 

259 async def _membership(self, user_id: str, team_id: str) -> LiteLLM_TeamMembership | None: 

260 return await self._load_membership( 

261 user_id=user_id, 

262 team_id=team_id, 

263 prisma_client=self._prisma, 

264 user_api_key_cache=self._cache, 

265 parent_otel_span=self._parent_otel_span, 

266 proxy_logging_obj=self._proxy_logging_obj, 

267 )