Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/management_helpers/bulk_user_creation.py: 82%

1"""Batched internal user creation behind `POST /management/v1/users/bulk`. 

2 

3The batch is validated with set queries, user rows land in one `create_many`, and every 

4referenced team is written once under its advisory lock instead of once per user. 

5""" 

6 

7import asyncio 

8import json 

9from collections.abc import Awaitable, Callable, Mapping, Sequence 

10from dataclasses import dataclass 

11from datetime import datetime 

12from types import MappingProxyType 

13from typing import TYPE_CHECKING, Final, Literal, TypeAlias, TypeVar 

14 

15from fastapi import HTTPException, Request 

16from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError 

17from typing_extensions import ReadOnly, TypedDict 

18 

19from litellm._logging import verbose_proxy_logger 

20from litellm._uuid import uuid 

21from litellm.integrations.prometheus import PrometheusLogger 

22from litellm.proxy._types import ( 

23 LiteLLM_TeamTable, 

24 LitellmUserRoles, 

25 Member, 

26 NewUserRequestTeam, 

27 OrganizationMemberAddRequest, 

28 OrgMember, 

29 UserAPIKeyAuth, 

30) 

31from litellm.proxy.auth.auth_checks import invalidate_team_member_spend_state 

32from litellm.proxy.auth.litellm_license import LicenseCheck 

33from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time 

34from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler 

35from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks 

36from litellm.proxy.list_api.common import PROBLEM_TYPE_BASE, ManagementProblem 

37from litellm.proxy.management_endpoints.common_utils import ( 

38 _is_user_org_admin_for_team, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses 

39 _is_user_team_admin, # pyright: ignore[reportPrivateUsage] # same team-admin check /user/new uses 

40 validate_budget_duration, 

41) 

42from litellm.proxy.management_endpoints.internal_user_endpoints import ( 

43 _update_internal_new_user_params, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # /user/new defaults; result validated below 

44 check_if_default_team_set, 

45) 

46from litellm.proxy.management_endpoints.key_management_endpoints import ( 

47 _check_permissions_caller_permission, # pyright: ignore[reportPrivateUsage] # same permission check /user/new uses 

48 generate_key_helper_fn, # pyright: ignore[reportUnknownVariableType] # legacy untyped helper; result validated by _KEY_RESPONSE 

49 metadata_json_with_limits, 

50) 

51from litellm.proxy.management_endpoints.organization_endpoints import organization_member_add 

52from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL 

53from litellm.proxy.management_helpers.object_permission_utils import ( 

54 _set_object_permission, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # shared with /user/new; result validated below 

55) 

56from litellm.proxy.management_helpers.utils import ( 

57 _resolve_member_budget_id, # pyright: ignore[reportPrivateUsage] # shared with /team/member_add 

58) 

59from litellm.proxy.utils import PrismaClient 

60from litellm.repositories.prisma_protocols import TableActions 

61from litellm.repositories.team_repository import TeamRepository 

62from litellm.repositories.user_repository import UserRepository 

63from litellm.types.proxy.management_endpoints.internal_user_endpoints import ( 

64 BulkNewUserItem, 

65 BulkNewUserMeta, 

66 BulkNewUserResponse, 

67 UserCreateResult, 

68) 

69from litellm.types.proxy.management_endpoints.management_v1 import ProblemDetail 

70 

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

72 from prisma import Prisma 

73 from prisma import models as prisma_models 

74 

75 from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache 

76 

77BULK_NEW_USER_CONCURRENCY: Final = 10 

78 

79TeamRole: TypeAlias = Literal["user", "admin"] 

80KeyGenerator: TypeAlias = Callable[..., Awaitable[object]] 

81_T: Final = TypeVar("_T") 

82 

83 

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

85class _RowFailure: 

86 index: int 

87 user_id: str | None 

88 user_email: str | None 

89 error: str 

90 

91 

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

93class _PendingUser: 

94 index: int 

95 request: BulkNewUserItem 

96 user_id: str 

97 teams: tuple[NewUserRequestTeam, ...] 

98 

99 

100class _UserRow(BaseModel): 

101 """The `/user/new` body after defaults and object permission were applied.""" 

102 

103 model_config = ConfigDict(extra="ignore") 

104 

105 user_id: str 

106 user_email: str | None = None 

107 user_alias: str | None = None 

108 user_role: str | None = None 

109 team_id: str | None = None 

110 max_budget: float | None = None 

111 spend: float | None = 0.0 

112 models: tuple[str, ...] | None = None 

113 metadata: Mapping[str, object] | None = None 

114 max_parallel_requests: int | None = None 

115 tpm_limit: int | None = None 

116 rpm_limit: int | None = None 

117 budget_duration: str | None = None 

118 allowed_cache_controls: tuple[str, ...] | None = None 

119 sso_user_id: str | None = None 

120 object_permission_id: str | None = None 

121 model_max_budget: Mapping[str, object] | None = None 

122 model_rpm_limit: Mapping[str, object] | None = None 

123 model_tpm_limit: Mapping[str, object] | None = None 

124 mcp_rpm_limit: Mapping[str, int] | None = None 

125 tag_rpm_limit: Mapping[str, int] | None = None 

126 guardrails: tuple[str, ...] | None = None 

127 policies: tuple[str, ...] | None = None 

128 prompts: tuple[str, ...] | None = None 

129 duration: str | None = None 

130 key_alias: str | None = None 

131 aliases: Mapping[str, object] | None = None 

132 config: Mapping[str, object] | None = None 

133 permissions: Mapping[str, object] | None = None 

134 blocked: bool | None = None 

135 agent_id: str | None = None 

136 budget_fallbacks: Mapping[str, tuple[str, ...]] | None = None 

137 budget_limits: tuple[Mapping[str, object], ...] | None = None 

138 organizations: tuple[str, ...] | None = None 

139 

140 

141_USER_ROW: Final = TypeAdapter(_UserRow) 

142 

143 

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

145class _PreparedUser: 

146 pending: _PendingUser 

147 row: _UserRow 

148 

149 

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

151class _TeamAssignment: 

152 user_id: str 

153 user_email: str | None 

154 role: TeamRole 

155 max_budget_in_team: float | None 

156 

157 

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

159class _TeamWrite: 

160 """Outcome of one locked roster write. `failed` maps user ids to the reason they were not added.""" 

161 

162 team_id: str 

163 after: tuple[Member, ...] 

164 added: frozenset[str] 

165 failed: Mapping[str, str] 

166 

167 

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

169class _CreatedUser: 

170 prepared: _PreparedUser 

171 teams: tuple[str, ...] 

172 key: str | None 

173 errors: tuple[str, ...] 

174 

175 

176_ERROR_DETAIL: Final = TypeAdapter(Mapping[str, object]) 

177_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) 

178 

179 

180class _KeyResponse(BaseModel): 

181 token: str 

182 

183 

184_KEY_RESPONSE: Final = TypeAdapter(_KeyResponse) 

185 

186 

187def _error_message(exc: BaseException) -> str: 

188 if not isinstance(exc, HTTPException): 

189 return str(exc) 

190 try: 

191 detail: Final = _ERROR_DETAIL.validate_python(exc.detail) 

192 except ValidationError: 

193 return str(exc.detail) 

194 return str(detail.get("error", detail)) 

195 

196 

197def _requested_teams(item: BulkNewUserItem) -> tuple[NewUserRequestTeam, ...]: 

198 if item.team_id is not None: 

199 return (NewUserRequestTeam(team_id=item.team_id),) 

200 teams: Final = item.teams if item.teams is not None else check_if_default_team_set() 

201 if teams is None: 

202 return () 

203 return tuple(team if isinstance(team, NewUserRequestTeam) else NewUserRequestTeam(team_id=team) for team in teams) 

204 

205 

206def _row_error(item: BulkNewUserItem, user_api_key_dict: UserAPIKeyAuth) -> str | None: 

207 if ( 207 ↛ 211line 207 didn't jump to line 211 because the condition on line 207 was never true

208 item.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) 

209 and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN 

210 ): 

211 return ( 

212 "Only proxy admins can create administrative users (proxy_admin, proxy_admin_viewer). " 

213 f"Attempted to create user with role: {item.user_role}. Your role: {user_api_key_dict.user_role}" 

214 ) 

215 try: 

216 validate_budget_duration(item.budget_duration) 

217 _check_permissions_caller_permission(data=item, user_api_key_dict=user_api_key_dict) 

218 except Exception as exc: # noqa: BLE001 # any validation failure is reported on this row only 

219 return _error_message(exc) 

220 return None 

221 

222 

223def _normalized_email(email: str | None) -> str | None: 

224 return email.strip().lower() if email else None 

225 

226 

227def _partition_rows( 

228 users: Sequence[BulkNewUserItem], user_api_key_dict: UserAPIKeyAuth 

229) -> tuple[tuple[_PendingUser, ...], tuple[_RowFailure, ...]]: 

230 """Assign ids, run the per-row checks and fail later rows that repeat an earlier row's id or email.""" 

231 user_ids: Final = tuple(item.user_id or str(uuid.uuid4()) for item in users) 

232 first_index_by_id: Final = MappingProxyType( 

233 {user_id: index for index, user_id in reversed(tuple(enumerate(user_ids)))} 

234 ) 

235 first_index_by_email: Final = MappingProxyType( 

236 { 

237 email: index 

238 for index, email in reversed(tuple(enumerate(_normalized_email(item.user_email) for item in users))) 

239 if email is not None 

240 } 

241 ) 

242 

243 def classify(index: int, item: BulkNewUserItem) -> _PendingUser | _RowFailure: 

244 user_id: Final = user_ids[index] 

245 email: Final = _normalized_email(item.user_email) 

246 if first_index_by_id[user_id] != index: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true

247 return _RowFailure(index, user_id, item.user_email, f"Duplicate user_id in request: {user_id}") 

248 if email is not None and first_index_by_email[email] != index: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true

249 return _RowFailure(index, user_id, item.user_email, f"Duplicate user_email in request: {item.user_email}") 

250 error: Final = _row_error(item, user_api_key_dict) 

251 if error is not None: 

252 return _RowFailure(index, user_id, item.user_email, error) 

253 return _PendingUser(index, item, user_id, _requested_teams(item)) 

254 

255 outcomes: Final = tuple(classify(index, item) for index, item in enumerate(users)) 

256 return ( 

257 tuple(outcome for outcome in outcomes if isinstance(outcome, _PendingUser)), 

258 tuple(outcome for outcome in outcomes if isinstance(outcome, _RowFailure)), 

259 ) 

260 

261 

262def _user_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_UserTable]": 

263 return UserRepository(prisma_client).table 

264 

265 

266async def _existing_user_conflicts( 

267 prisma_client: PrismaClient, pending: Sequence[_PendingUser] 

268) -> tuple[frozenset[str], frozenset[str]]: 

269 """Return the requested user ids and (lowercased) emails that already exist, using one query each.""" 

270 user_ids: Final = sorted(user.user_id for user in pending) 

271 emails: Final = sorted(frozenset(user.request.user_email for user in pending if user.request.user_email)) 

272 if not user_ids: 

273 return frozenset(), frozenset() 

274 table: Final = _user_table(prisma_client) 

275 id_filter: Final = {"user_id": {"in": user_ids}} # mutable-ok: Prisma query filters are dict-shaped 

276 email_filter: Final = {"user_email": {"in": emails, "mode": "insensitive"}} # mutable-ok: Prisma filter 

277 id_rows: Final = await table.find_many(where=id_filter) 

278 email_rows: Final = await table.find_many(where=email_filter) if emails else () 

279 return ( 

280 frozenset(row.user_id for row in id_rows), 

281 frozenset(lowered for row in email_rows if (lowered := _normalized_email(row.user_email)) is not None), 

282 ) 

283 

284 

285async def _load_teams(prisma_client: PrismaClient, team_ids: frozenset[str]) -> Mapping[str, LiteLLM_TeamTable]: 

286 if not team_ids: 

287 return MappingProxyType({}) 

288 rows: Final = await TeamRepository(prisma_client).table.find_many( 

289 where={"team_id": {"in": sorted(team_ids)}} # mutable-ok: Prisma query filters are dict-shaped 

290 ) 

291 return MappingProxyType({row.team_id: LiteLLM_TeamTable.model_validate(row.model_dump()) for row in rows}) 

292 

293 

294async def _team_permission_error(team: LiteLLM_TeamTable, user_api_key_dict: UserAPIKeyAuth) -> str | None: 

295 if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: 

296 return None 

297 if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team): 

298 return None 

299 if await _is_user_org_admin_for_team(user_api_key_dict=user_api_key_dict, team_obj=team): 

300 return None 

301 return f"Call not allowed. User not proxy admin OR team admin. team_id={team.team_id}" 

302 

303 

304async def _unusable_teams( 

305 prisma_client: PrismaClient, 

306 pending: Sequence[_PendingUser], 

307 user_api_key_dict: UserAPIKeyAuth, 

308) -> tuple[Mapping[str, LiteLLM_TeamTable], Mapping[str, str]]: 

309 """Load every referenced team once and explain, per team id, why rows naming it cannot proceed.""" 

310 team_ids: Final = frozenset(team.team_id for user in pending for team in user.teams) 

311 teams: Final = await _load_teams(prisma_client, team_ids) 

312 permission_errors: Final = await asyncio.gather( 

313 *(_team_permission_error(team, user_api_key_dict) for team in teams.values()) 

314 ) 

315 missing: Final = tuple( 

316 (team_id, f"Team id={team_id} does not exist") for team_id in team_ids if team_id not in teams 

317 ) 

318 denied: Final = tuple( 

319 (team.team_id, error) 

320 for team, error in zip(teams.values(), permission_errors, strict=True) 

321 if error is not None 

322 ) 

323 return teams, MappingProxyType({team_id: error for team_id, error in (*missing, *denied)}) 

324 

325 

326def _db_failure( 

327 user: _PendingUser, 

328 existing_ids: frozenset[str], 

329 existing_emails: frozenset[str], 

330 team_errors: Mapping[str, str], 

331) -> _RowFailure | None: 

332 email: Final = _normalized_email(user.request.user_email) 

333 if user.user_id in existing_ids: 333 ↛ 334line 333 didn't jump to line 334 because the condition on line 333 was never true

334 return _RowFailure(user.index, user.user_id, user.request.user_email, f"User id={user.user_id} already exists") 

335 if email is not None and email in existing_emails: 

336 return _RowFailure( 

337 user.index, user.user_id, user.request.user_email, f"User email={user.request.user_email} already exists" 

338 ) 

339 errors: Final = tuple(team_errors[team.team_id] for team in user.teams if team.team_id in team_errors) 

340 if errors: 

341 return _RowFailure(user.index, user.user_id, user.request.user_email, "; ".join(errors)) 

342 return None 

343 

344 

345async def _prepare_user(user: _PendingUser, prisma_client: PrismaClient) -> _PreparedUser | _RowFailure: 

346 try: 

347 dumped: Final = user.request.model_dump(exclude={"user_id"}) # mutable-ok: pydantic IncEx takes a set 

348 data: Final = {**dumped, "user_id": user.user_id} # mutable-ok: /user/new defaults helper mutates in place 

349 data_json: Final = _JSON_OBJECT.validate_python(_update_internal_new_user_params(data, user.request)) 

350 with_permission: Final = _JSON_OBJECT.validate_python( 

351 await _set_object_permission(data_json=data_json, prisma_client=prisma_client) 

352 ) 

353 return _PreparedUser(user, _USER_ROW.validate_python(with_permission)) 

354 except Exception as exc: # noqa: BLE001 # any preparation failure is reported on this row only 

355 verbose_proxy_logger.warning("/user/bulk_new: could not prepare row %d - %s", user.index, type(exc).__name__) 

356 return _RowFailure(user.index, user.user_id, user.request.user_email, _error_message(exc)) 

357 

358 

359class _UserCreateData(TypedDict): 

360 """One `LiteLLM_UserTable` row as `create_many` takes it; JSON columns are pre-serialized.""" 

361 

362 user_id: ReadOnly[str] 

363 user_email: ReadOnly[str | None] 

364 user_alias: ReadOnly[str | None] 

365 user_role: ReadOnly[str | None] 

366 team_id: ReadOnly[str | None] 

367 max_budget: ReadOnly[float | None] 

368 spend: ReadOnly[float] 

369 models: ReadOnly[tuple[str, ...]] 

370 metadata: ReadOnly[str] 

371 max_parallel_requests: ReadOnly[int | None] 

372 tpm_limit: ReadOnly[int | None] 

373 rpm_limit: ReadOnly[int | None] 

374 budget_duration: ReadOnly[str | None] 

375 budget_reset_at: ReadOnly[datetime | None] 

376 allowed_cache_controls: ReadOnly[tuple[str, ...]] 

377 sso_user_id: ReadOnly[str | None] 

378 object_permission_id: ReadOnly[str | None] 

379 teams: ReadOnly[tuple[str, ...]] 

380 model_max_budget: ReadOnly[str] 

381 

382 

383def _user_create_payload(prepared: _PreparedUser) -> _UserCreateData: 

384 row: Final = prepared.row 

385 metadata_json: Final = metadata_json_with_limits( 

386 row.metadata, 

387 model_rpm_limit=row.model_rpm_limit, 

388 model_tpm_limit=row.model_tpm_limit, 

389 mcp_rpm_limit=row.mcp_rpm_limit, 

390 tag_rpm_limit=row.tag_rpm_limit, 

391 guardrails=row.guardrails, 

392 policies=row.policies, 

393 prompts=row.prompts, 

394 ) 

395 payload: Final[_UserCreateData] = { 

396 "user_id": row.user_id, 

397 "user_email": row.user_email, 

398 "user_alias": row.user_alias, 

399 "user_role": row.user_role, 

400 "team_id": row.team_id, 

401 "max_budget": row.max_budget, 

402 "spend": row.spend or 0.0, 

403 "models": row.models or (), 

404 "metadata": metadata_json, 

405 "max_parallel_requests": row.max_parallel_requests, 

406 "tpm_limit": row.tpm_limit, 

407 "rpm_limit": row.rpm_limit, 

408 "budget_duration": row.budget_duration, 

409 "budget_reset_at": get_budget_reset_time(row.budget_duration) if row.budget_duration else None, 

410 "allowed_cache_controls": row.allowed_cache_controls or (), 

411 "sso_user_id": row.sso_user_id, 

412 "object_permission_id": row.object_permission_id, 

413 "teams": tuple(team.team_id for team in prepared.pending.teams), 

414 "model_max_budget": json.dumps(row.model_max_budget) if row.model_max_budget else "{}", 

415 } 

416 return payload 

417 

418 

419async def _bounded(limit: int, awaitables: Sequence[Awaitable[_T]]) -> tuple[_T | BaseException, ...]: 

420 semaphore: Final = asyncio.Semaphore(limit) 

421 

422 async def run(awaitable: Awaitable[_T]) -> _T: 

423 async with semaphore: 

424 return await awaitable 

425 

426 return tuple(await asyncio.gather(*(run(awaitable) for awaitable in awaitables), return_exceptions=True)) 

427 

428 

429async def _insert_users( 

430 prisma_client: PrismaClient, prepared: Sequence[_PreparedUser] 

431) -> tuple[tuple[_PreparedUser, ...], tuple[_RowFailure, ...]]: 

432 """Insert every row in one statement. If that fails, retry rows one at a time so the error lands on its row.""" 

433 if not prepared: 

434 return (), () 

435 table: Final = _user_table(prisma_client) 

436 payloads: Final = tuple(_user_create_payload(user) for user in prepared) 

437 try: 

438 await table.create_many(data=payloads) 

439 return tuple(prepared), () 

440 except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified 

441 verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually", exc_info=True) 

442 outcome_unknown: Final = PrismaDBExceptionHandler.is_database_infrastructure_error(exc) 

443 requested: Final = frozenset(payload["user_id"] for payload in payloads) 

444 landed_rows: Final = await table.find_many(where={"user_id": {"in": list(requested)}}) # mutable-ok: Prisma filter 

445 landed: Final = frozenset(row.user_id for row in landed_rows) 

446 # create_many is one INSERT: after a lost response the full set is ours, any partial set belongs to another request 

447 if outcome_unknown and landed == requested: 447 ↛ 448line 447 didn't jump to line 448 because the condition on line 447 was never true

448 return tuple(prepared), () 

449 taken: Final = tuple(user for user in prepared if user.row.user_id in landed) 

450 retried: Final = tuple(user for user in prepared if user.row.user_id not in landed) 

451 outcomes: Final = await _bounded( 

452 BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried) 

453 ) 

454 failed: Final = MappingProxyType( 

455 { 

456 **{ 

457 user.row.user_id: _RowFailure( 

458 user.pending.index, 

459 user.pending.user_id, 

460 user.row.user_email, 

461 f"User id={user.row.user_id} already exists", 

462 ) 

463 for user in taken 

464 }, 

465 **{ 

466 user.row.user_id: _RowFailure( 

467 user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome) 

468 ) 

469 for user, outcome in zip(retried, outcomes, strict=True) 

470 if isinstance(outcome, BaseException) 

471 }, 

472 } 

473 ) 

474 return ( 

475 tuple(user for user in prepared if user.row.user_id not in failed), 

476 tuple(failed.values()), 

477 ) 

478 

479 

480def _assignments_by_team(created: Sequence[_PreparedUser]) -> Mapping[str, tuple[_TeamAssignment, ...]]: 

481 team_ids: Final = tuple(dict.fromkeys(team.team_id for user in created for team in user.pending.teams)) 

482 return MappingProxyType( 

483 { 

484 team_id: tuple( 

485 _TeamAssignment(user.pending.user_id, user.row.user_email, team.user_role, team.max_budget_in_team) 

486 for user in created 

487 for team in user.pending.teams 

488 if team.team_id == team_id 

489 ) 

490 for team_id in team_ids 

491 } 

492 ) 

493 

494 

495class _MembershipData(TypedDict): 

496 team_id: ReadOnly[str] 

497 user_id: ReadOnly[str] 

498 budget_id: ReadOnly[str | None] 

499 

500 

501class _RosterData(TypedDict): 

502 members_with_roles: ReadOnly[str] 

503 

504 

505class _TeamsData(TypedDict): 

506 teams: ReadOnly[tuple[str, ...]] 

507 

508 

509def _default_member_budget_id(team: LiteLLM_TeamTable) -> str | None: 

510 metadata: Final = ( 

511 _JSON_OBJECT.validate_python( 

512 team.metadata # pyright: ignore[reportUnknownMemberType] # LiteLLM_TeamTable.metadata is a bare dict; validated by the adapter 

513 ) 

514 if team.metadata # pyright: ignore[reportUnknownMemberType] # same bare dict 

515 else None 

516 ) 

517 budget_id: Final = metadata.get("team_member_budget_id") if metadata is not None else None 

518 return budget_id if isinstance(budget_id, str) else None 

519 

520 

521def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]": 

522 return tx.litellm_teamtable # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do 

523 

524 

525def _membership_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamMembership]": 

526 return tx.litellm_teammembership # pyright: ignore[reportReturnType] # TableActions widens the generated inputs to Mapping, as the repositories do 

527 

528 

529async def _write_team_roster( 

530 prisma_client: PrismaClient, 

531 team: LiteLLM_TeamTable, 

532 members: Sequence[_TeamAssignment], 

533 user_api_key_dict: UserAPIKeyAuth, 

534 litellm_proxy_admin_name: str, 

535) -> _TeamWrite: 

536 """Add every new member to one team under its advisory lock: one roster rewrite and one membership insert.""" 

537 try: 

538 async with prisma_client.tx() as tx: 

539 await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team.team_id) 

540 roster: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, team.team_id) 

541 if roster is None: 

542 raise ValueError(f"Team id={team.team_id} does not exist") 

543 already_present: Final = frozenset(member.user_id for member in roster if member.user_id) 

544 new_members: Final = tuple(member for member in members if member.user_id not in already_present) 

545 budget_ids: Final = tuple( 

546 [ 

547 await _resolve_member_budget_id( 

548 prisma_client=prisma_client, 

549 user_api_key_dict=user_api_key_dict, 

550 litellm_proxy_admin_name=litellm_proxy_admin_name, 

551 max_budget_in_team=member.max_budget_in_team, 

552 allowed_models=team.default_team_member_models or None, 

553 budget_duration=None, 

554 default_team_budget_id=_default_member_budget_id(team), 

555 tx=tx, # pyright: ignore[reportArgumentType] # MemberWriteTx lags the generated Prisma signatures, same as /team/member_add 

556 ) 

557 for member in new_members 

558 ] 

559 ) 

560 await _membership_tx_db(tx).create_many( 

561 data=tuple( 

562 _MembershipData(team_id=team.team_id, user_id=member.user_id, budget_id=budget_id) 

563 for member, budget_id in zip(new_members, budget_ids, strict=True) 

564 ), 

565 skip_duplicates=True, 

566 ) 

567 after: Final = ( 

568 *roster, 

569 *(Member(user_id=m.user_id, user_email=m.user_email, role=m.role) for m in new_members), 

570 ) 

571 await _team_tx_db(tx).update( 

572 where={"team_id": team.team_id}, # mutable-ok: Prisma query filters are dict-shaped 

573 data=_RosterData(members_with_roles=json.dumps(tuple(member.model_dump() for member in after))), 

574 ) 

575 return _TeamWrite( 

576 team_id=team.team_id, 

577 after=after, 

578 added=frozenset(member.user_id for member in members), 

579 failed=MappingProxyType({}), 

580 ) 

581 except Exception as exc: # noqa: BLE001 # the team write failure is reported on each affected row 

582 verbose_proxy_logger.exception("/user/bulk_new: failed to add %d members to a team", len(members)) 

583 message: Final = f"Failed to add user to team {team.team_id}: {_error_message(exc)}" 

584 return _TeamWrite( 

585 team_id=team.team_id, 

586 after=(), 

587 added=frozenset(), 

588 failed=MappingProxyType({member.user_id: message for member in members}), 

589 ) 

590 

591 

592async def _detach_failed_teams( 

593 prisma_client: PrismaClient, created: Sequence[_PreparedUser], writes: Mapping[str, _TeamWrite] 

594) -> None: 

595 """Users are inserted with `teams` already set; drop the teams whose roster write did not take them.""" 

596 table: Final = _user_table(prisma_client) 

597 updates: Final = tuple( 

598 table.update( 

599 where={"user_id": user.row.user_id}, # mutable-ok: Prisma query filters are dict-shaped 

600 data=_TeamsData(teams=landed), 

601 ) 

602 for user in created 

603 if (landed := _row_teams(user, writes)[0]) != tuple(team.team_id for team in user.pending.teams) 

604 ) 

605 for outcome in await _bounded(BULK_NEW_USER_CONCURRENCY, updates): 605 ↛ 606line 605 didn't jump to line 606 because the loop on line 605 never started

606 if isinstance(outcome, BaseException): 

607 verbose_proxy_logger.warning( 

608 "/user/bulk_new: could not detach failed teams from user - %s", type(outcome).__name__ 

609 ) 

610 

611 

612async def _publish_team_writes(writes: Sequence[_TeamWrite], user_api_key_cache: "UserApiKeyCache") -> None: 

613 prometheus_logger: Final = PrometheusLogger.get_instance() 

614 for write in writes: 614 ↛ 615line 614 didn't jump to line 615 because the loop on line 614 never started

615 if prometheus_logger is None or not write.added: 

616 continue 

617 try: 

618 prometheus_logger.set_team_members_metric( 

619 LiteLLM_TeamTable( 

620 team_id=write.team_id, 

621 members_with_roles=write.after, # pyright: ignore[reportArgumentType] # pydantic coerces the tuple into the declared list 

622 ) 

623 ) 

624 except Exception: # noqa: BLE001 # metrics are best-effort and must not fail the request 

625 verbose_proxy_logger.debug("Prometheus: failed to emit team members metric", exc_info=True) 

626 evictions: Final = await _bounded( 

627 BULK_NEW_USER_CONCURRENCY, 

628 tuple( 

629 invalidate_team_member_spend_state( 

630 user_id=user_id, team_id=write.team_id, user_api_key_cache=user_api_key_cache 

631 ) 

632 for write in writes 

633 for user_id in write.added 

634 ), 

635 ) 

636 for eviction in evictions: 636 ↛ 637line 636 didn't jump to line 637 because the loop on line 636 never started

637 if isinstance(eviction, BaseException): 

638 verbose_proxy_logger.warning("/user/bulk_new: cache eviction failed - %s", type(eviction).__name__) 

639 

640 

641_KEY_FIELDS: Final = MappingProxyType( 

642 { 

643 name: True 

644 for name in ( 

645 "user_id", 

646 "team_id", 

647 "agent_id", 

648 "duration", 

649 "key_alias", 

650 "models", 

651 "aliases", 

652 "config", 

653 "permissions", 

654 "blocked", 

655 "spend", 

656 "budget_fallbacks", 

657 "budget_limits", 

658 "metadata", 

659 "max_parallel_requests", 

660 "tpm_limit", 

661 "rpm_limit", 

662 "allowed_cache_controls", 

663 "model_max_budget", 

664 "model_rpm_limit", 

665 "model_tpm_limit", 

666 "mcp_rpm_limit", 

667 "tag_rpm_limit", 

668 "guardrails", 

669 "policies", 

670 "prompts", 

671 "object_permission_id", 

672 ) 

673 } 

674) 

675 

676 

677async def _generate_key(prepared: _PreparedUser, generate_key: KeyGenerator) -> str: 

678 response: Final = _KEY_RESPONSE.validate_python( 

679 await generate_key( 

680 request_type="key", table_name="key", **prepared.row.model_dump(include=_KEY_FIELDS, exclude_none=True) 

681 ) 

682 ) 

683 return response.token 

684 

685 

686async def _add_to_organizations( 

687 prepared: _PreparedUser, organizations: Sequence[str], user_api_key_dict: UserAPIKeyAuth 

688) -> None: 

689 for organization_id in organizations: 689 ↛ exitline 689 didn't return from function '_add_to_organizations' because the loop on line 689 didn't complete

690 await organization_member_add( 

691 data=OrganizationMemberAddRequest( 

692 organization_id=organization_id, 

693 member=OrgMember(user_id=prepared.row.user_id, role=LitellmUserRoles.INTERNAL_USER), 

694 ), 

695 http_request=Request(scope={"type": "http", "path": "/user/bulk_new"}), # mutable-ok: ASGI scopes are dicts 

696 user_api_key_dict=user_api_key_dict, 

697 ) 

698 

699 

700async def _run_per_user( 

701 created: Sequence[_PreparedUser],