scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

fix(proxy): keep key policy fields and reconcile committed rows in /user/bulk_new

Browse source 
Rows opting into auto_create_key lost blocked, permissions, aliases, config,
agent_id, budget_fallbacks and budget_limits before reaching the key helper.
When create_many commits but the response is lost, re-read which ids landed
and retry only the rest so committed rows report success and get their teams.
Regenerate schema.d.ts and allowlist the endpoint in the Terraform audit.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-14 01:48:44 +00:00
parent ff5b59b173
commit 24a1d772b6
4 changed files with 315 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
litellm/proxy/management_helpers/bulk_user_creation.py
View file

@ -124,6 +124,13 @@ class _UserRow(BaseModel):
prompts: tuple[str, ...] | None = None
duration: str | None = None
key_alias: str | None = None
aliases: Mapping[str, object] | None = None
config: Mapping[str, object] | None = None
permissions: Mapping[str, object] | None = None
blocked: bool | None = None
agent_id: str | None = None
budget_fallbacks: Mapping[str, tuple[str, ...]] | None = None
budget_limits: tuple[Mapping[str, object], ...] | None = None
organizations: tuple[str, ...] | None = None
@ -428,16 +435,26 @@ async def _insert_users(
return tuple(prepared), ()
except Exception as exc: # noqa: BLE001 # fall back to per-row inserts so the failing row can be identified
verbose_proxy_logger.warning("/user/bulk_new: create_many failed, retrying rows individually - %s", exc)
landed_rows: Final = await table.find_many(
where={"user_id": {"in": [payload["user_id"] for payload in payloads]}} # mutable-ok: Prisma filter
)
landed: Final = frozenset(row.user_id for row in landed_rows)
retried: Final = tuple(user for user in prepared if user.row.user_id not in landed)
outcomes: Final = await _bounded(
BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=payload) for payload in payloads)
BULK_NEW_USER_CONCURRENCY, tuple(table.create(data=_user_create_payload(user)) for user in retried)
)
failed: Final = MappingProxyType(
{
user.row.user_id: _RowFailure(
user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome)
)
for user, outcome in zip(retried, outcomes, strict=True)
if isinstance(outcome, BaseException)
}
)
return (
tuple(user for user, outcome in zip(prepared, outcomes, strict=True) if not isinstance(outcome, BaseException)),
tuple(
_RowFailure(user.pending.index, user.pending.user_id, user.row.user_email, _error_message(outcome))
for user, outcome in zip(prepared, outcomes, strict=True)
if isinstance(outcome, BaseException)
),
tuple(user for user in prepared if user.row.user_id not in failed),
tuple(failed.values()),
)
@ -606,9 +623,17 @@ _KEY_FIELDS: Final = MappingProxyType(
for name in (
"user_id",
"team_id",
"agent_id",
"duration",
"key_alias",
"models",
"aliases",
"config",
"permissions",
"blocked",
"spend",
"budget_fallbacks",
"budget_limits",
"metadata",
"max_parallel_requests",
"tpm_limit",

1
terraform/provider/tools/endpointaudit/coverage_allowlist.txt
View file

@ -84,6 +84,7 @@ POST /team/{team_id}/member/{user_id}/reset_spend
POST /team/key/bulk_update
POST /team/permissions_bulk_update
POST /team/{team_id}/disable_logging
POST /user/bulk_new
POST /user/bulk_update
# Alternate method or path for functionality the provider already manages elsewhere

48
tests/test_litellm/proxy/management_helpers/test_bulk_user_creation.py
View file

@ -31,9 +31,10 @@ class _UserRow(BaseModel):
class _UserTable:
"""Enough of the Prisma user table for the bulk path: set lookups, one create_many and per-row fallbacks."""
def __init__(self, fail_ids: frozenset[str] = frozenset()) -> None:
def __init__(self, fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False) -> None:
self.rows: dict[str, _UserRow] = {}
self.fail_ids = fail_ids
self.commit_then_drop = commit_then_drop
self.create_many_calls = 0
async def count(self, where: object = None) -> int:
@ -60,6 +61,8 @@ class _UserTable:
raise RuntimeError("batch insert failed")
for row in rows:
self.rows[row.user_id] = row
if self.commit_then_drop:
raise ConnectionError("connection reset after commit")
return len(rows)
async def update(self, where: dict[str, str], data: dict[str, object]) -> _UserRow:
@ -110,15 +113,22 @@ class _Tx:
class _Db:
def __init__(self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset()) -> None:
self.litellm_usertable = _UserTable(fail_ids)
def __init__(
self, teams: list[LiteLLM_TeamTable], fail_ids: frozenset[str] = frozenset(), commit_then_drop: bool = False
) -> None:
self.litellm_usertable = _UserTable(fail_ids, commit_then_drop)
self.litellm_teamtable = _TeamTable(teams)
self.litellm_teammembership = _MembershipTable()
class _FakePrisma:
def __init__(self, teams: list[LiteLLM_TeamTable] | None = None, fail_ids: frozenset[str] = frozenset()) -> None:
self.db = _Db(teams or [], fail_ids)
def __init__(
self,
teams: list[LiteLLM_TeamTable] | None = None,
fail_ids: frozenset[str] = frozenset(),
commit_then_drop: bool = False,
) -> None:
self.db = _Db(teams or [], fail_ids, commit_then_drop)
self.tx_count = 0
self.locks: list[str] = []
@ -255,6 +265,17 @@ async def test_insert_failure_falls_back_to_per_row_and_reports_only_that_row():
assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"]
@pytest.mark.asyncio
async def test_insert_that_committed_but_lost_its_response_still_counts_as_created():
prisma = _FakePrisma(teams=[_team("t1")], commit_then_drop=True)
response = await _run(prisma, [{"user_id": "u1", "teams": ["t1"]}, {"user_id": "u2"}])
assert [r.success for r in response.results] == [True, True]
assert [r.error for r in response.results] == [None, None]
assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2"}
assert [m.user_id for m in prisma.db.litellm_teamtable.rows["t1"].members_with_roles] == ["u1"]
@pytest.mark.asyncio
async def test_team_write_failure_keeps_user_and_reports_it_on_the_row():
prisma = _FakePrisma(teams=[_team("t1"), _team("t2")])
@ -286,7 +307,17 @@ async def test_keys_are_opt_in_per_row():
prisma,
[
{"user_id": "u1"},
{"user_id": "u2", "auto_create_key": True, "models": ["gpt-4o"], "key_alias": "u2-key"},
{
"user_id": "u2",
"auto_create_key": True,
"models": ["gpt-4o"],
"key_alias": "u2-key",
"blocked": True,
"permissions": {"get_spend_routes": True},
"aliases": {"fast": "gpt-4o"},
"config": {"tier": "gold"},
"budget_fallbacks": {"gpt-4o": ["gpt-4o-mini"]},
},
{"user_id": "u3", "auto_create_key": False},
],
generate_key=generate_key,
@ -296,6 +327,11 @@ async def test_keys_are_opt_in_per_row():
assert len(calls) == 1
assert calls[0]["user_id"] == "u2" and calls[0]["table_name"] == "key"
assert calls[0]["models"] == ("gpt-4o",) and calls[0]["key_alias"] == "u2-key"
assert calls[0]["blocked"] is True
assert calls[0]["permissions"] == {"get_spend_routes": True}
assert calls[0]["aliases"] == {"fast": "gpt-4o"}
assert calls[0]["config"] == {"tier": "gold"}
assert calls[0]["budget_fallbacks"] == {"gpt-4o": ("gpt-4o-mini",)}
assert set(prisma.db.litellm_usertable.rows) == {"u1", "u2", "u3"}

242
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -16478,6 +16478,53 @@ export interface paths {
patch?: never;
trace?: never;
};
"/user/bulk_new": {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
get?: never;
put?: never;
/**
* Bulk New User
* @description Create up to 500 internal users in one request, optionally adding each one to teams.
*
* Every entry in `users` takes the same fields as `/user/new`, with two differences: `auto_create_key`
* defaults to `false` (opt in per user to also get a virtual key back) and `send_invite_email` is not
* supported. Rows are validated together (duplicate ids or emails, unknown teams, roles the caller may not
* grant), inserted in one statement, and each referenced team is written once for all of its new members.
*
* Rows fail independently: a bad row is reported in `results` with `success: false` and an `error`, and the
* other rows still get created. A user that was created but could not be added to one of its teams is
* reported with `success: true`, `teams` listing where they did land, and `error` naming the failed team.
* The whole request is rejected with 403 only if creating the valid rows would exceed the license seat limit.
*
* Usage Example
*
* ```shell
* curl -X POST "http://localhost:4000/user/bulk_new" \
* -H "Content-Type: application/json" \
* -H "Authorization: Bearer sk-1234" \
* -d '{
* "users": [
* {"user_email": "a@example.com", "user_role": "internal_user", "teams": ["team-1"]},
* {"user_email": "b@example.com", "user_role": "internal_user", "auto_create_key": true}
* ]
* }'
* ```
*
* Returns `results` (one entry per input row, in order, with `user_id`, `user_email`, `success`, `teams`,
* `key`, `error`), `total_requested`, `successful_creations` and `failed_creations`.
*/
post: operations["bulk_new_user_user_bulk_new_post"];
delete?: never;
options?: never;
head?: never;
patch?: never;
trace?: never;
};
"/user/bulk_update": {
parameters: {
query?: never;
@ -16781,7 +16828,6 @@ export interface paths {
* - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
* - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
* - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
@ -16887,7 +16933,6 @@ export interface paths {
* - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
* - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
* - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
* - soft_budget: Optional[float] - Get alerts when user crosses given budget, doesn't block requests.
* - model_max_budget: Optional[dict] - Model-specific max budget for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-budgets-to-keys)
* - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}.
* - model_rpm_limit: Optional[float] - Model-specific rpm limit for user. [Docs](https://docs.litellm.ai/docs/proxy/users#add-model-specific-limits-to-keys)
@ -24505,6 +24550,148 @@ export interface components {
/** Budgets */
budgets: string[];
};
/**
* BulkNewUserItem
* @description One row of `/user/bulk_new`: the `/user/new` body, with keys opt-in and invite emails unsupported.
*/
BulkNewUserItem: {
/** Agent Id */
agent_id?: string | null;
/**
* Aliases
* @default {}
*/
aliases: {
[key: string]: unknown;
} | null;
/**
* Allowed Cache Controls
* @default []
*/
allowed_cache_controls: unknown[] | null;
/**
* Auto Create Key
* @default false
*/
auto_create_key: boolean;
/** Blocked */
blocked?: boolean | null;
/** Budget Duration */
budget_duration?: string | null;
/** Budget Fallbacks */
budget_fallbacks?: {
[key: string]: string[];
} | null;
/** Budget Limits */
budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null;
/**
* Config
* @default {}
*/
config: {
[key: string]: unknown;
} | null;
/** Duration */
duration?: string | null;
/** Guardrails */
guardrails?: string[] | null;
/** Key Alias */
key_alias?: string | null;
/** Max Budget */
max_budget?: number | null;
/** Max Parallel Requests */
max_parallel_requests?: number | null;
/** Mcp Rpm Limit */
mcp_rpm_limit?: {
[key: string]: number;
} | null;
/**
* Metadata
* @default {}
*/
metadata: {
[key: string]: unknown;
} | null;
/**
* Model Max Budget
* @default {}
*/
model_max_budget: {
[key: string]: unknown;
} | null;
/** Model Rpm Limit */
model_rpm_limit?: {
[key: string]: unknown;
} | null;
/** Model Tpm Limit */
model_tpm_limit?: {
[key: string]: unknown;
} | null;
/**
* Models
* @default []
*/
models: unknown[] | null;
object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null;
/** Organizations */
organizations?: string[] | null;
/**
* Permissions
* @default {}
*/
permissions: {
[key: string]: unknown;
} | null;
/** Policies */
policies?: string[] | null;
/** Prompts */
prompts?: string[] | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Send Invite Email */
send_invite_email?: boolean | null;
/**
* Spend
* @default 0
*/
spend: number | null;
/** Sso User Id */
sso_user_id?: string | null;
/** Tag Rpm Limit */
tag_rpm_limit?: {
[key: string]: number;
} | null;
/** Team Id */
team_id?: string | null;
/** Teams */
teams?: string[] | components["schemas"]["NewUserRequestTeam"][] | null;
/** Tpm Limit */
tpm_limit?: number | null;
/** User Alias */
user_alias?: string | null;
/** User Email */
user_email?: string | null;
/** User Id */
user_id?: string | null;
/** User Role */
user_role?: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null;
};
/** BulkNewUserRequest */
BulkNewUserRequest: {
/** Users */
users: components["schemas"]["BulkNewUserItem"][];
};
/** BulkNewUserResponse */
BulkNewUserResponse: {
/** Failed Creations */
failed_creations: number;
/** Results */
results: components["schemas"]["UserCreateResult"][];
/** Successful Creations */
successful_creations: number;
/** Total Requested */
total_requested: number;
};
/**
* BulkTeamMemberAddRequest
* @description Request for bulk team member addition
@ -39345,6 +39532,24 @@ export interface components {
*/
severity: "info" | "warning" | "error";
};
/**
* UserCreateResult
* @description Outcome for one row of `/user/bulk_new`. `teams` lists the teams the user was actually added to.
*/
UserCreateResult: {
/** Error */
error?: string | null;
/** Key */
key?: string | null;
/** Success */
success: boolean;
/** Teams */
teams?: string[] | null;
/** User Email */
user_email?: string | null;
/** User Id */
user_id?: string | null;
};
/**
* UserHeaderMapping
* @description Map an incoming HTTP header to a LiteLLM user role.
@ -60631,6 +60836,39 @@ export interface operations {
};
};
};
bulk_new_user_user_bulk_new_post: {
parameters: {
query?: never;
header?: never;
path?: never;
cookie?: never;
};
requestBody: {
content: {
"application/json": components["schemas"]["BulkNewUserRequest"];
};
};
responses: {
/** @description Successful Response */
200: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["BulkNewUserResponse"];
};
};
/** @description Validation Error */
422: {
headers: {
[name: string]: unknown;
};
content: {
"application/json": components["schemas"]["HTTPValidationError"];
};
};
};
};
bulk_user_update_user_bulk_update_post: {
parameters: {
query?: never;