scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(customer-portal): keep existing memberships and scope free tier capacity locks (#51698)

Browse source 
* fix(customer-portal): preserve memberships and isolate capacity locks

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(customer-portal): map model write lock timeouts to retryable failures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
moe-berri 2026-11-14 04:51:54 +08:00 • committed by GitHub
parent 57c6bbfa12
commit 945ee1ca43
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 267 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/management_endpoints/management_v1/customers.py
View file

@ -52,6 +52,11 @@ T: Final = TypeVar("T")
_ERROR_DETAILS: Final[Mapping[ErrorCode, tuple[int, str, str]]] = MappingProxyType(
{
ErrorCode.CAPACITY_BUSY: (
503,
"Customer portal busy",
"Customer portal capacity is being updated. Retry the request.",
),
ErrorCode.FREE_TIER_LIMIT: (
403,
"Free customer portal limit reached",

27
litellm/proxy/management_helpers/customer_hierarchy.py
View file

@ -16,6 +16,7 @@ from litellm.proxy.common_utils.timezone_utils import (
get_budget_reset_time,
)
from litellm.proxy.db.db_span import db_span
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.list_api.common import escape_like
from litellm.proxy.management_helpers.customer_budget_usage import customer_budget_usage
from litellm.repositories.budget_repository import create_budget_row, update_budget_row
@ -46,6 +47,7 @@ if TYPE_CHECKING:
class ErrorCode(str, Enum):
FREE_TIER_LIMIT = "free_tier_limit"
CAPACITY_BUSY = "capacity_busy"
NOT_FOUND = "not_found"
PARENT_NOT_FOUND = "parent_not_found"
ALREADY_EXISTS = "already_exists"
@ -76,6 +78,7 @@ FREE_ORGANIZATION_LIMIT: Final = 5
FREE_TEAMS_PER_ORGANIZATION_LIMIT: Final = 5
FREE_MEMBERS_PER_TEAM_LIMIT: Final = 5
FREE_CAPACITY_ADVISORY_LOCK: Final = 714023581
FREE_CAPACITY_LOCK_TIMEOUT_MS: Final = 1000
class CustomerPortalLimits(BaseModel):
@ -501,14 +504,22 @@ class CustomerHierarchyStore:
self.budgets: TableActions[PrismaRecord] = db.litellm_budgettable # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
self.window_spend: TableActions[LiteLLM_BudgetWindowSpend] = db.litellm_budgetwindowspend # pyright: ignore[reportAttributeAccessIssue] # protocol accepts immutable query mappings
async def lock_free_capacity(self) -> None:
await self.db.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock($1)", FREE_CAPACITY_ADVISORY_LOCK)
async def lock_free_capacity(self, scope: str | None = None) -> None:
await self.db.query_raw("SELECT set_config('lock_timeout', $1, true)", f"{FREE_CAPACITY_LOCK_TIMEOUT_MS}ms")
if scope is None:
await self.db.query_raw("SELECT 1 AS locked FROM pg_advisory_xact_lock($1)", FREE_CAPACITY_ADVISORY_LOCK)
return
await self.db.query_raw(
"SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtextextended($1, $2))",
scope,
FREE_CAPACITY_ADVISORY_LOCK,
)
async def member_capacity_reached(self, team_ids: Sequence[str]) -> bool:
if self.premium_user or not team_ids:
return False
await self.lock_free_capacity()
for team_id in team_ids:
for team_id in sorted(frozenset(team_ids)):
await self.lock_free_capacity(f"members:{team_id}")
if await self.members.count(where={"parent_team_ids": {"has": team_id}}) >= FREE_MEMBERS_PER_TEAM_LIMIT:
return True
return False
@ -876,7 +887,7 @@ class CustomerHierarchyStore:
if parent is None:
return Failure(ErrorCode.PARENT_NOT_FOUND)
if not self.premium_user:
await self.lock_free_capacity()
await self.lock_free_capacity(f"teams:{payload.parent_org_id}")
if await self.teams.count(
where={"parent_org_id": payload.parent_org_id}
) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT and (
@ -1192,7 +1203,7 @@ class CustomerHierarchy:
*,
operation_type: TransactionOperation = TransactionOperation.READ_WRITE,
) -> T | Failure:
from prisma.errors import ForeignKeyViolationError, RecordNotFoundError, UniqueViolationError
from prisma.errors import DataError, ForeignKeyViolationError, RecordNotFoundError, UniqueViolationError
try:
async with self.db.tx() as tx:
@ -1205,6 +1216,10 @@ class CustomerHierarchy:
return Failure(
ErrorCode.HAS_CHILDREN if operation_type is TransactionOperation.DELETE else ErrorCode.PARENT_NOT_FOUND
)
except DataError as error:
if PrismaDBExceptionHandler.is_lock_timeout_error(error):
return Failure(ErrorCode.CAPACITY_BUSY)
raise
async def create_organization(self, payload: OrganizationCreate, actor: str) -> Organization | Failure:
return await self.transaction(lambda store: store.create_organization(payload, actor))

35
litellm/proxy/management_helpers/customer_upsert.py
View file

@ -15,6 +15,7 @@ from litellm.proxy.management_helpers.customer_hierarchy import (
from litellm.types.proxy.management_endpoints.customer_hierarchy import (
Budget,
BudgetTemplate,
CustomerMember,
MemberCreate,
MemberPatch,
MemberUpsert,
@ -159,7 +160,7 @@ async def customer_team(
else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH)
)
if not store.premium_user:
await store.lock_free_capacity()
await store.lock_free_capacity(f"teams:{organization_id}")
if (
await store.teams.find_unique(where={"customer_team_id": default_team_id(organization_id)}) is None
and await store.teams.count(where={"parent_org_id": organization_id}) >= FREE_TEAMS_PER_ORGANIZATION_LIMIT
@ -184,6 +185,23 @@ async def customer_team(
return team if team.parent_org_id == organization_id else Failure(ErrorCode.CUSTOMER_ORG_MISMATCH)
async def member_team_ids(
store: CustomerHierarchyStore,
organization_id: str,
existing: CustomerMember | Failure,
team_id: str | None,
actor: str,
) -> tuple[str, ...] | Failure:
if not isinstance(existing, Failure) and team_id is None:
return existing.parent_team_ids
team: Final = await customer_team(store, organization_id, team_id, actor)
if isinstance(team, Failure):
return team
if isinstance(existing, Failure):
return (team.customer_team_id,)
return (team.customer_team_id, *(tid for tid in existing.parent_team_ids if tid != team.customer_team_id))
async def upsert_member(
store: CustomerHierarchyStore, payload: MemberUpsert, actor: str
) -> MemberUpsertResponse | Failure:
@ -226,23 +244,18 @@ async def upsert_member(
)
if selected is None and override is None:
return Failure(ErrorCode.BUDGET_REQUIRED)
team: Final = await customer_team(store, organization.customer_org_id, payload.parent_team_id, actor)
if isinstance(team, Failure):
return team
team_ids: Final = (
(team.customer_team_id,)
if isinstance(existing, Failure)
else existing.parent_team_ids
if payload.parent_team_id is None
else (team.customer_team_id, *(tid for tid in existing.parent_team_ids if tid != team.customer_team_id))
team_ids: Final = await member_team_ids(
store, organization.customer_org_id, existing, payload.parent_team_id, actor
)
if isinstance(team_ids, Failure):
return team_ids
if isinstance(existing, Failure):
member: Final = await store.create_member(
MemberCreate.model_validate(
{
**patch.model_dump(exclude_unset=True),
"customer_member_id": payload.customer_member_id,
"parent_team_ids": (team.customer_team_id,),
"parent_team_ids": team_ids,
}
),
actor,

63
tests/unit/proxy/management_endpoints/management_v1/test_customers.py
View file

@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Final, cast
import pytest
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from prisma.errors import DataError, RawQueryError
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.auth.customer_hierarchy_auth import customer_member_cache_key, customer_org_alias_cache_key
@ -19,6 +20,7 @@ from litellm.types.proxy.management_endpoints.customer_hierarchy import Customer
from tests.unit.proxy.management_helpers.test_customer_hierarchy import (
NOW,
BoundedChildrenTable,
CapacityErrorDatabase,
Database,
DeleteConflictTable,
Record,
@ -127,6 +129,15 @@ def test_member_upsert_respects_capacity_for_default_team_creation_and_added_mem
"/management/v1/customer_members", json={"customer_member_id": "moving", "parent_team_ids": ["extra-0"]}
)
assert original.status_code == 201, original.text
for patch in ({}, {"alias": "renamed"}):
kept: Final = http.post(
"/management/v1/customer_members/upsert",
json={"customer_org_id": "org", "customer_member_id": "moving", **patch},
)
assert kept.status_code == 200, kept.text
assert kept.json()["created"] is False
assert kept.json()["parent_team_ids"] == ["extra-0"]
assert kept.json()["alias"] == patch.get("alias")
moved: Final = http.post(
"/management/v1/customer_members/upsert",
json={"customer_org_id": "org", "customer_member_id": "moving", "parent_team_id": "team"},
@ -137,11 +148,63 @@ def test_member_upsert_respects_capacity_for_default_team_creation_and_added_mem
assert http.get("/management/v1/customer_teams/org:default").status_code == 404
@pytest.mark.parametrize("error_code", ("55P03", "XX000"))
def test_capacity_contention_is_retryable_without_hiding_other_database_errors(error_code: str) -> None:
http: Final = client(database=CapacityErrorDatabase(error_code, "members:team"), premium_user=False)
family(http)
payload: Final = {"customer_member_id": "blocked", "parent_team_ids": ["team"]}
if error_code != "55P03":
with pytest.raises(RawQueryError, match="advisory lock failed"):
http.post("/management/v1/customer_members", json=payload)
return
response: Final = http.post("/management/v1/customer_members", json=payload)
assert response.status_code == 503, response.text
assert response.json() == {
"type": "urn:litellm:error:customer-hierarchy-capacity-busy",
"title": "Customer portal busy",
"status": 503,
"detail": "Customer portal capacity is being updated. Retry the request.",
}
assert http.get("/management/v1/customer_members/blocked").status_code == 404
def test_customer_portal_limits_require_admin() -> None:
response: Final = client(role=LitellmUserRoles.INTERNAL_USER).get("/management/v1/customer_portal/limits")
assert response.status_code == 403, response.text
@pytest.mark.parametrize(
("error", "retryable"),
(
(DataError({"user_facing_error": {"meta": {"code": "55P03", "message": "lock unavailable"}}}), True),
(DataError({"user_facing_error": {"message": "canceling statement due to lock timeout"}}), True),
(DataError({"user_facing_error": {"message": "other database error", "meta": {"code": "XX000"}}}), False),
),
ids=("sqlstate", "message-only", "other-error"),
)
def test_team_model_write_contention_is_retryable_and_rolls_back(error: DataError, retryable: bool) -> None:
database: Final = Database(team_create_error=error)
http: Final = client(database=database, premium_user=False)
created: Final = http.post("/management/v1/customer_organizations", json={"customer_org_id": "org"})
assert created.status_code == 201, created.text
payload: Final = {"customer_team_id": "blocked", "parent_org_id": "org", "budget": {"max_budget": 10}}
if not retryable:
with pytest.raises(DataError, match="other database error") as raised:
http.post("/management/v1/customer_teams", json=payload)
assert raised.value is error
else:
response: Final = http.post("/management/v1/customer_teams", json=payload)
assert response.status_code == 503, response.text
assert response.json() == {
"type": "urn:litellm:error:customer-hierarchy-capacity-busy",
"title": "Customer portal busy",
"status": 503,
"detail": "Customer portal capacity is being updated. Retry the request.",
}
assert http.get("/management/v1/customer_teams/blocked").status_code == 404
assert not database.litellm_budgettable.rows
def family(http: TestClient) -> dict[str, object]:
org: Final = http.post(
"/management/v1/customer_organizations",

166
tests/unit/proxy/management_helpers/test_customer_hierarchy.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import sqlite3
from collections.abc import AsyncIterator, Mapping, Sequence
@ -5,10 +6,12 @@ from contextlib import asynccontextmanager
from datetime import datetime, timezone
from functools import reduce
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, cast
import pytest
from prisma import Json
from prisma.errors import DataError, RawQueryError
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from litellm.models.budget import BudgetLimit, LiteLLM_BudgetTable
@ -23,8 +26,10 @@ from litellm.proxy.management_helpers.customer_hierarchy import (
)
from litellm.types.proxy.management_endpoints.customer_hierarchy import (
Budget,
CustomerMember,
MemberCreate,
OrganizationCreate,
Team,
TeamCreate,
)
@ -84,9 +89,10 @@ def matches_where(row: Mapping[str, object], where: Mapping[str, object]) -> boo
class Table:
def __init__(self, id_field: str, database: "Database") -> None:
def __init__(self, id_field: str, database: "Database", *, create_error: DataError | None = None) -> None:
self.id_field = id_field
self.database = database
self.create_error = create_error
self.rows: Mapping[str, Record] = {}
async def include(self, row: Record, include: Mapping[str, object] | None) -> Record:
@ -109,6 +115,8 @@ class Table:
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record:
from prisma.errors import UniqueViolationError
if self.create_error is not None:
raise self.create_error
if str(data[self.id_field]) in self.rows:
raise UniqueViolationError({})
row: Final = Record.model_validate(
@ -237,16 +245,63 @@ class DeleteConflictTable(Table):
raise ForeignKeyViolationError({})
class Database:
class AdvisoryLocks:
def __init__(self) -> None:
self.litellm_customerorganizationtable = Table("customer_org_id", self)
self.litellm_customerteamtable = Table("customer_team_id", self)
self.litellm_customermembertable = Table("customer_member_id", self)
self.litellm_budgettable = Table("budget_id", self)
self.litellm_budgetwindowspend = Table("entity_id", self)
self.litellm_verificationtoken = Table("token", self)
self.litellm_projecttable = Table("project_id", self)
self.litellm_customerbudgetincreaserequest = Table("request_id", self)
self.locks: Mapping[tuple[str, ...], asyncio.Lock] = MappingProxyType({})
def lock_for(self, query: str, args: Sequence[object]) -> asyncio.Lock:
key: Final = (query, *(str(arg) for arg in args))
existing: Final = self.locks.get(key)
if existing is not None:
return existing
lock: Final = asyncio.Lock()
self.locks = MappingProxyType({**self.locks, key: lock})
return lock
class PausedCreateTable(Table):
def __init__(self, id_field: str, database: "Database", entered: asyncio.Event, release: asyncio.Event) -> None:
super().__init__(id_field, database)
self.entered = entered
self.release = release
async def create(self, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> Record:
if data.get(self.id_field) == "blocked":
self.entered.set()
await self.release.wait()
return await super().create(data, include)
class Database:
def __init__(
self,
advisory_locks: AdvisoryLocks | None = None,
source: "Database | None" = None,
*,
team_create_error: DataError | None = None,
) -> None:
self.litellm_customerorganizationtable = (
Table("customer_org_id", self) if source is None else source.litellm_customerorganizationtable
)
self.litellm_customerteamtable = (
Table("customer_team_id", self, create_error=team_create_error)
if source is None
else source.litellm_customerteamtable
)
self.litellm_customermembertable = (
Table("customer_member_id", self) if source is None else source.litellm_customermembertable
)
self.litellm_budgettable = Table("budget_id", self) if source is None else source.litellm_budgettable
self.litellm_budgetwindowspend = (
Table("entity_id", self) if source is None else source.litellm_budgetwindowspend
)
self.litellm_verificationtoken = Table("token", self) if source is None else source.litellm_verificationtoken
self.litellm_projecttable = Table("project_id", self) if source is None else source.litellm_projecttable
self.litellm_customerbudgetincreaserequest = (
Table("request_id", self) if source is None else source.litellm_customerbudgetincreaserequest
)
self.advisory_locks = advisory_locks
self.held_locks: tuple[asyncio.Lock, ...] = ()
self.locked_organization_ids: tuple[str, ...] = ()
def has_budget_reference(self, budget_id: str, relation: str) -> bool:
@ -265,9 +320,21 @@ class Database:
table, field = reference
return any(row.model_dump().get(field) == budget_id for row in table.rows.values())
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
if "pg_advisory_xact_lock" in query:
async def acquire_advisory_lock(self, query: str, args: Sequence[object]) -> tuple[Mapping[str, object], ...]:
if self.advisory_locks is None:
return ()
lock: Final = self.advisory_locks.lock_for(query, args)
if lock in self.held_locks:
return ()
await lock.acquire()
self.held_locks = (*self.held_locks, lock)
return ()
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
if "set_config('lock_timeout'" in query:
return ()
if "pg_advisory_xact_lock" in query:
return await self.acquire_advisory_lock(query, args)
if "SELECT EXISTS" in query:
return ({"referenced": False},)
if "FOR UPDATE" in query:
@ -334,6 +401,81 @@ class Database:
for table, rows in zip(tables, snapshots):
table.rows = rows
raise
finally:
for lock in self.held_locks:
lock.release()
self.held_locks = ()
class CapacityErrorDatabase(Database):
def __init__(self, error_code: str, scope: str) -> None:
super().__init__()
self.error_code = error_code
self.scope = scope
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
if "pg_advisory_xact_lock" in query and args and args[0] == self.scope:
raise RawQueryError(
{"user_facing_error": {"meta": {"code": self.error_code, "message": "advisory lock failed"}}}
)
return await super().query_raw(query, *args)
@pytest.mark.asyncio
@pytest.mark.parametrize("kind", ("team", "member"))
@pytest.mark.parametrize("same_parent", (False, True))
async def test_capacity_contention_is_isolated_to_each_parent(kind: str, same_parent: bool) -> None:
locks: Final = AdvisoryLocks()
database: Final = Database(locks)
hierarchy: Final = CustomerHierarchy(
cast("Prisma", database), premium_user=False, clock=lambda: NOW
) # cast-ok: in-memory Prisma edge
for org_id in ("a", "b"):
assert not isinstance(
await hierarchy.create_organization(OrganizationCreate(customer_org_id=org_id), "admin"), Failure
)
if kind == "member":
assert not isinstance(
await hierarchy.create_team(TeamCreate(customer_team_id=org_id, parent_org_id=org_id), "admin"), Failure
)
entered: Final = asyncio.Event()
release: Final = asyncio.Event()
if kind == "team":
database.litellm_customerteamtable = PausedCreateTable("customer_team_id", database, entered, release)
else:
database.litellm_customermembertable = PausedCreateTable("customer_member_id", database, entered, release)
other: Final = CustomerHierarchy(
cast("Prisma", Database(locks, source=database)), premium_user=False, clock=lambda: NOW
) # cast-ok: independent connection to the shared in-memory tables
async def create(store: CustomerHierarchy, parent: str, identity: str) -> Team | CustomerMember | Failure:
if kind == "team":
return await store.create_team(TeamCreate(customer_team_id=identity, parent_org_id=parent), "admin")
return await store.create_member(
MemberCreate(customer_member_id=identity, parent_team_ids=(parent,), budget_override=Budget(max_budget=10)),
"admin",
)
for index in range(4):
assert not isinstance(await create(hierarchy, "a", f"seed-{index}"), Failure)
blocked: Final = asyncio.create_task(create(hierarchy, "a", "blocked"))
await entered.wait()
independent: Final = asyncio.create_task(create(other, "a" if same_parent else "b", "independent"))
try:
if same_parent:
await asyncio.sleep(0)
release.set()
result: Final = await asyncio.wait_for(independent, timeout=10)
if same_parent:
assert result == Failure(ErrorCode.FREE_TIER_LIMIT)
else:
assert not isinstance(result, Failure)
assert result == (
await other.get_team("independent") if kind == "team" else await other.get_member("independent")
)
finally:
release.set()
assert not isinstance(await blocked, Failure)
@pytest.mark.asyncio