mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Add team member budget columns and fix Next Budget Reset display (#24942)
* Add team member budget columns and fix Next Budget Reset display - Add Spend (Current Period), Total Spend, Team Member Budget, Next Budget Reset, and Rate Limits columns to the team members table - Add total_spend field to LiteLLM_TeamMembership (accumulates across resets; never zeroed by the reset job) - Add budget_reset_at to LiteLLM_BudgetTable Python type and schema - Fix Next Budget Reset showing proxy start time: run ResetBudgetJob once immediately at startup via asyncio.create_task so expired budgets are advanced before the first UI request - Fix update_budget to only recalculate budget_reset_at when duration genuinely changes (avoid clobbering a valid future reset on idempotent startup calls) - Fix _reset_budget_reset_at_date to anchor null-reset budgets to created_at instead of current_time - Add Budget Reset Period selector to Edit Member modal; wire budget_duration through TeamMemberUpdateRequest and _upsert_budget_and_membership so individual member budgets get a budget_reset_at on creation - Wire members with budget_id=null to the team template budget on settings update Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Address PR review: fix test patch target, migration, multi-period reset, wire-up scope - Fix test_upsert_team_member_budget_table: patch target was team_endpoints.prisma_client (module attr) but production code does a lazy import from proxy_server; correct target is litellm.proxy.proxy_server.prisma_client - Split test into two: create path wires up null-budget members (expected), update path does NOT (new assertion for backwards-compat fix) - Move null-budget wire-up to create-only path in upsert_team_member_budget_table to avoid silently constraining members who were intentionally unlimited on every settings save - Add migration 20260401000000_add_total_spend_to_team_membership with ALTER TABLE ... ADD COLUMN IF NOT EXISTS "total_spend" for existing deployments - Fix _reset_budget_reset_at_date else branch: use while loop (matching the None path) so proxies down for multiple periods advance the timestamp all the way to the next future date in one scheduler tick - Add test_reset_budget_reset_at_date_multi_period_catchup to cover this Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * test: restore mock_new_budget.called assertion dropped in previous fix Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: update existing budget row instead of orphaning it in _upsert_budget_and_membership When existing_budget_id is provided (member already has a budget), update the existing LiteLLM_BudgetTable row in-place rather than always creating a new one. The previous code accepted existing_budget_id but never used it, silently orphaning one budget row per team_member_update call that set any budget field. Adds two unit tests covering the update-existing and create-new paths. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: allow clearing budget_duration via UI; sync schema copies Frontend: remove the != null guard in teamMemberUpdateCall so that explicitly selecting 'No reset (unlimited)' sends budget_duration: null to the backend (rather than omitting the field entirely). Backend: add clear_budget_duration flag to _upsert_budget_and_membership. The team_member_update endpoint sets this flag when budget_duration is present in model_fields_set but null, allowing the update path to explicitly null out budget_duration and budget_reset_at on the existing budget row. Also sync litellm/proxy/schema.prisma and litellm-proxy-extras schema to root (missing total_spend column + alignment drift). * Add total_spend to LiteLLM_TeamMembership in all schema.prisma copies The migration already added the column to the DB, but all three schema.prisma files were missing the field, causing Prisma to throw InvalidPrismaInvocation on every team membership spend write. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * Fix early-return guard ignoring clear_budget_duration flag When clear_budget_duration=True with all budget limits None, the guard at the top of _upsert_budget_and_membership was short-circuiting to disconnect the budget, bypassing the intended update. Adding `not clear_budget_duration` to the guard condition lets the function reach the update path. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Roy Vivat
GitHub
parent
096893ea97
commit
19306ff4a4
20 changed files with 1157 additions and 46 deletions
|
|
@ -0,0 +1,4 @@
|
|||
-- Add total_spend column to LiteLLM_TeamMembership
|
||||
-- Tracks lifetime (never-zeroed) spend for a user within a team,
|
||||
-- independent of the current-period spend that resets periodically.
|
||||
ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
|
@ -608,9 +608,10 @@ model LiteLLM_UserNotifications {
|
|||
|
||||
model LiteLLM_TeamMembership {
|
||||
// Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
|
|||
|
|
@ -1943,6 +1943,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
rpm_limit: Optional[int] = None
|
||||
model_max_budget: Optional[dict] = None
|
||||
budget_duration: Optional[str] = None
|
||||
budget_reset_at: Optional[datetime] = None
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
|
|
@ -3622,6 +3623,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
|
|||
team_id: str
|
||||
budget_id: Optional[str] = None
|
||||
spend: Optional[float] = 0.0
|
||||
total_spend: Optional[float] = 0.0
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable]
|
||||
|
||||
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
|
||||
|
|
@ -3750,6 +3752,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
|
|||
rpm_limit: Optional[int] = Field(
|
||||
default=None, description="Requests per minute limit for this team member"
|
||||
)
|
||||
budget_duration: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Budget reset period for this team member (e.g. '24h', '7d', '30d'). Determines how often their individual spend resets.",
|
||||
)
|
||||
|
||||
|
||||
class TeamMemberUpdateResponse(MemberUpdateResponse):
|
||||
|
|
|
|||
|
|
@ -52,13 +52,17 @@ class ResetBudgetJob:
|
|||
self, budgets_to_reset: List[LiteLLM_BudgetTableFull]
|
||||
):
|
||||
"""
|
||||
Resets the budget for all LiteLLM Team Members if their budget has expired
|
||||
Resets the budget for all LiteLLM Team Members if their budget has expired.
|
||||
Only resets the current-period spend; total_spend is never zeroed so it
|
||||
accumulates as a lifetime usage counter.
|
||||
"""
|
||||
budget_ids = [
|
||||
budget.budget_id
|
||||
for budget in budgets_to_reset
|
||||
if budget.budget_id is not None
|
||||
]
|
||||
if not budget_ids:
|
||||
return
|
||||
|
||||
# Reset spend counters for affected team members.
|
||||
# Reset Redis directly so a transient failure doesn't leave stale
|
||||
|
|
@ -658,18 +662,25 @@ class ResetBudgetJob:
|
|||
|
||||
duration_s = duration_in_seconds(duration=budget.budget_duration)
|
||||
|
||||
# Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account
|
||||
if (
|
||||
budget.budget_reset_at is None
|
||||
and budget.created_at + timedelta(seconds=duration_s) > current_time
|
||||
):
|
||||
budget.budget_reset_at = budget.created_at + timedelta(
|
||||
seconds=duration_s
|
||||
)
|
||||
if budget.budget_reset_at is None:
|
||||
# Budget was created without a reset time. Anchor the first
|
||||
# period to created_at so the reset aligns with when the
|
||||
# budget was originally set up, not the current wall-clock
|
||||
# time (which would be proxy-start time when this job runs
|
||||
# at startup and produce a confusing "resets at startup" display).
|
||||
anchor = budget.created_at + timedelta(seconds=duration_s)
|
||||
while anchor <= current_time:
|
||||
anchor += timedelta(seconds=duration_s)
|
||||
budget.budget_reset_at = anchor
|
||||
else:
|
||||
budget.budget_reset_at = current_time + timedelta(
|
||||
seconds=duration_s
|
||||
)
|
||||
# Normal roll-forward: advance from the previous reset time.
|
||||
# Use a while loop so proxies that were down for multiple
|
||||
# periods catch up in one pass rather than firing the
|
||||
# spend-zero operation repeatedly over many scheduler ticks.
|
||||
while budget.budget_reset_at <= current_time:
|
||||
budget.budget_reset_at = budget.budget_reset_at + timedelta(
|
||||
seconds=duration_s
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
"Error resetting budget_reset_at for budget: %s. Item: %s", e, budget
|
||||
|
|
|
|||
|
|
@ -1295,7 +1295,10 @@ class DBSpendUpdateWriter:
|
|||
|
||||
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
|
||||
where={"team_id": team_id, "user_id": user_id},
|
||||
data={"spend": {"increment": response_cost}},
|
||||
data={
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
},
|
||||
)
|
||||
# Transaction succeeded, break out of retry loop
|
||||
break
|
||||
|
|
|
|||
|
|
@ -173,12 +173,38 @@ async def update_budget(
|
|||
except ValueError as e:
|
||||
raise HTTPException(status_code=400, detail={"error": str(e)})
|
||||
|
||||
update_data: dict = {
|
||||
**budget_obj.model_dump(exclude_unset=True),
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
}
|
||||
|
||||
# Recalculate budget_reset_at only when budget_duration is genuinely changing.
|
||||
# We compare against the current DB value so that idempotent calls at proxy
|
||||
# startup (e.g. _update_default_team_member_budget) do not overwrite a
|
||||
# correctly-set future budget_reset_at with "now + duration".
|
||||
# We check budget_obj.budget_duration directly (not model_fields_set) because
|
||||
# internal callers set fields via attribute assignment after construction, which
|
||||
# does not reliably update model_fields_set.
|
||||
if budget_obj.budget_duration is not None and "budget_reset_at" not in update_data:
|
||||
existing = await prisma_client.db.litellm_budgettable.find_unique(
|
||||
where={"budget_id": budget_obj.budget_id}
|
||||
)
|
||||
existing_duration = existing.budget_duration if existing is not None else None
|
||||
duration_changed = existing_duration != budget_obj.budget_duration
|
||||
reset_at_missing = (
|
||||
existing is None or existing.budget_reset_at is None
|
||||
)
|
||||
if duration_changed or reset_at_missing:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
update_data["budget_duration"] = budget_obj.budget_duration
|
||||
update_data["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=budget_obj.budget_duration
|
||||
)
|
||||
|
||||
response = await prisma_client.db.litellm_budgettable.update(
|
||||
where={"budget_id": budget_obj.budget_id},
|
||||
data={
|
||||
**budget_obj.model_dump(exclude_unset=True), # type: ignore
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
}, # type: ignore
|
||||
data=update_data, # type: ignore
|
||||
)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -354,6 +354,8 @@ async def _upsert_budget_and_membership(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
tpm_limit: Optional[int] = None,
|
||||
rpm_limit: Optional[int] = None,
|
||||
budget_duration: Optional[str] = None,
|
||||
clear_budget_duration: bool = False,
|
||||
):
|
||||
"""
|
||||
Helper function to Create/Update or Delete the budget within the team membership
|
||||
|
|
@ -366,11 +368,15 @@ async def _upsert_budget_and_membership(
|
|||
user_api_key_dict: User API Key dictionary containing user information
|
||||
tpm_limit: Tokens per minute limit for the team member
|
||||
rpm_limit: Requests per minute limit for the team member
|
||||
budget_duration: Budget reset period (e.g. '24h', '7d', '30d')
|
||||
clear_budget_duration: When True, explicitly clears budget_duration (and
|
||||
budget_reset_at) on an existing budget. Distinct from budget_duration=None
|
||||
which means "not supplied / don't touch".
|
||||
|
||||
If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership.
|
||||
If max_budget, tpm_limit, rpm_limit, and budget_duration are all None, the user's budget is removed from the team membership.
|
||||
If any of these values exist, a budget is updated or created and linked to the team membership.
|
||||
"""
|
||||
if max_budget is None and tpm_limit is None and rpm_limit is None:
|
||||
if max_budget is None and tpm_limit is None and rpm_limit is None and budget_duration is None and not clear_budget_duration:
|
||||
# disconnect the budget since all limits are None
|
||||
await tx.litellm_teammembership.update(
|
||||
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
|
||||
|
|
@ -389,11 +395,34 @@ async def _upsert_budget_and_membership(
|
|||
create_data["tpm_limit"] = tpm_limit
|
||||
if rpm_limit is not None:
|
||||
create_data["rpm_limit"] = rpm_limit
|
||||
if budget_duration is not None:
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
|
||||
create_data["budget_duration"] = budget_duration
|
||||
create_data["budget_reset_at"] = get_budget_reset_time(
|
||||
budget_duration=budget_duration
|
||||
)
|
||||
|
||||
if existing_budget_id is not None:
|
||||
# Update in-place: patch only the fields that were supplied so we don't
|
||||
# overwrite fields the caller didn't touch (e.g. keep max_budget when
|
||||
# only budget_duration changes). Exclude created_by — that's set once.
|
||||
update_data = {k: v for k, v in create_data.items() if k != "created_by"}
|
||||
if clear_budget_duration:
|
||||
update_data["budget_duration"] = None
|
||||
update_data["budget_reset_at"] = None
|
||||
await tx.litellm_budgettable.update(
|
||||
where={"budget_id": existing_budget_id},
|
||||
data=update_data,
|
||||
)
|
||||
budget_id_to_connect = existing_budget_id
|
||||
else:
|
||||
new_budget = await tx.litellm_budgettable.create(
|
||||
data=create_data,
|
||||
include={"team_membership": True},
|
||||
)
|
||||
budget_id_to_connect = new_budget.budget_id
|
||||
|
||||
new_budget = await tx.litellm_budgettable.create(
|
||||
data=create_data,
|
||||
include={"team_membership": True},
|
||||
)
|
||||
# upsert the team membership with the new/updated budget
|
||||
await tx.litellm_teammembership.upsert(
|
||||
where={
|
||||
|
|
@ -407,12 +436,12 @@ async def _upsert_budget_and_membership(
|
|||
"user_id": user_id,
|
||||
"team_id": team_id,
|
||||
"litellm_budget_table": {
|
||||
"connect": {"budget_id": new_budget.budget_id},
|
||||
"connect": {"budget_id": budget_id_to_connect},
|
||||
},
|
||||
},
|
||||
"update": {
|
||||
"litellm_budget_table": {
|
||||
"connect": {"budget_id": new_budget.budget_id},
|
||||
"connect": {"budget_id": budget_id_to_connect},
|
||||
},
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -231,7 +231,7 @@ class TeamMemberBudgetHandler:
|
|||
updated_kv["metadata"] = {}
|
||||
updated_kv["metadata"]["team_member_budget_id"] = budget_row.budget_id
|
||||
|
||||
else: # budget does not exist
|
||||
else: # budget does not exist — newly creating the template
|
||||
updated_kv = await TeamMemberBudgetHandler.create_team_member_budget_table(
|
||||
data=team_table,
|
||||
new_team_data_json=updated_kv,
|
||||
|
|
@ -242,6 +242,27 @@ class TeamMemberBudgetHandler:
|
|||
team_member_budget_duration=team_member_budget_duration,
|
||||
)
|
||||
|
||||
# Wire up existing members whose budget_id is null to the
|
||||
# newly-created template budget so they immediately inherit the
|
||||
# reset schedule. Only do this on creation (not on every update)
|
||||
# to avoid silently constraining members who were intentionally
|
||||
# left unlimited.
|
||||
if team_table.team_id is not None:
|
||||
new_budget_id: Optional[str] = (
|
||||
updated_kv.get("metadata", {}).get("team_member_budget_id")
|
||||
)
|
||||
if new_budget_id is not None:
|
||||
from litellm.proxy.proxy_server import prisma_client as _prisma_client
|
||||
|
||||
if _prisma_client is not None:
|
||||
await _prisma_client.db.litellm_teammembership.update_many(
|
||||
where={
|
||||
"team_id": team_table.team_id,
|
||||
"budget_id": None,
|
||||
},
|
||||
data={"budget_id": new_budget_id},
|
||||
)
|
||||
|
||||
# Remove team member fields from updated_kv
|
||||
TeamMemberBudgetHandler._clean_team_member_fields(updated_kv)
|
||||
return updated_kv
|
||||
|
|
@ -2427,6 +2448,11 @@ async def team_member_update(
|
|||
break
|
||||
|
||||
### upsert new budget
|
||||
# budget_duration=None means "not supplied" by default; we need to distinguish
|
||||
# that from the user explicitly sending null to clear a previously-set value.
|
||||
clear_budget_duration = (
|
||||
"budget_duration" in data.model_fields_set and data.budget_duration is None
|
||||
)
|
||||
async with prisma_client.db.tx() as tx:
|
||||
await _upsert_budget_and_membership(
|
||||
tx=tx,
|
||||
|
|
@ -2437,6 +2463,8 @@ async def team_member_update(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
tpm_limit=data.tpm_limit,
|
||||
rpm_limit=data.rpm_limit,
|
||||
budget_duration=data.budget_duration,
|
||||
clear_budget_duration=clear_budget_duration,
|
||||
)
|
||||
|
||||
### update team member role
|
||||
|
|
|
|||
|
|
@ -6273,6 +6273,12 @@ class ProxyStartupEvent:
|
|||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
# Run once immediately at startup so expired budget_reset_at values are
|
||||
# advanced before the first UI request. Without this, any budget that
|
||||
# expired exactly at restart time would show a stale past timestamp
|
||||
# until the first scheduled interval fires (~10 minutes later).
|
||||
asyncio.create_task(budget_reset_job.reset_budget())
|
||||
|
||||
### UPDATE SPEND ###
|
||||
scheduler.add_job(
|
||||
update_spend,
|
||||
|
|
|
|||
|
|
@ -608,9 +608,10 @@ model LiteLLM_UserNotifications {
|
|||
|
||||
model LiteLLM_TeamMembership {
|
||||
// Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
|
|||
|
|
@ -608,9 +608,10 @@ model LiteLLM_UserNotifications {
|
|||
|
||||
model LiteLLM_TeamMembership {
|
||||
// Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
|
|||
|
|
@ -18,10 +18,13 @@ from litellm.proxy.utils import ProxyLogging
|
|||
|
||||
# Mock classes for testing
|
||||
class MockLiteLLMTeamMembership:
|
||||
def __init__(self):
|
||||
self.update_many_calls: List[Dict[str, Any]] = []
|
||||
|
||||
async def update_many(
|
||||
self, where: Dict[str, Any], data: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
# Mock the update_many method for litellm_teammembership
|
||||
self.update_many_calls.append({"where": where, "data": data})
|
||||
return {"count": 1}
|
||||
|
||||
|
||||
|
|
@ -41,6 +44,10 @@ class MockDB:
|
|||
self.litellm_teammembership = MockLiteLLMTeamMembership()
|
||||
self.litellm_verificationtoken = MockLiteLLMVerificationToken()
|
||||
|
||||
def reset(self):
|
||||
self.litellm_teammembership = MockLiteLLMTeamMembership()
|
||||
self.litellm_verificationtoken = MockLiteLLMVerificationToken()
|
||||
|
||||
|
||||
class MockPrismaClient:
|
||||
def __init__(self):
|
||||
|
|
@ -444,6 +451,101 @@ def test_reset_budget_for_keys_linked_to_budgets_empty(
|
|||
assert len(calls) == 0
|
||||
|
||||
|
||||
def test_reset_budget_for_team_members_resets_spend(
|
||||
reset_budget_job, mock_prisma_client
|
||||
):
|
||||
"""
|
||||
Test that reset_budget_for_litellm_team_members resets spend to 0
|
||||
for all team memberships linked to the expired budgets.
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
test_budget = type(
|
||||
"LiteLLM_BudgetTableFull",
|
||||
(),
|
||||
{
|
||||
"max_budget": 50.0,
|
||||
"budget_duration": "1mo",
|
||||
"budget_reset_at": now,
|
||||
"budget_id": "team-member-budget-1",
|
||||
"created_at": now - timedelta(days=30),
|
||||
},
|
||||
)
|
||||
|
||||
budgets_to_reset = [test_budget]
|
||||
|
||||
asyncio.run(
|
||||
reset_budget_job.reset_budget_for_litellm_team_members(
|
||||
budgets_to_reset=budgets_to_reset
|
||||
)
|
||||
)
|
||||
|
||||
calls = mock_prisma_client.db.litellm_teammembership.update_many_calls
|
||||
assert len(calls) == 1
|
||||
|
||||
call = calls[0]
|
||||
# spend must be zeroed out for the matching budget IDs
|
||||
assert call["where"]["budget_id"] == {"in": ["team-member-budget-1"]}
|
||||
assert call["data"]["spend"] == 0
|
||||
|
||||
|
||||
def test_reset_budget_for_team_members_preserves_total_spend(
|
||||
reset_budget_job, mock_prisma_client
|
||||
):
|
||||
"""
|
||||
Test that reset_budget_for_litellm_team_members only resets the current-period
|
||||
spend and does NOT reset total_spend. total_spend must accumulate across all
|
||||
budget periods to give a lifetime view of member usage.
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
|
||||
test_budget = type(
|
||||
"LiteLLM_BudgetTableFull",
|
||||
(),
|
||||
{
|
||||
"max_budget": 50.0,
|
||||
"budget_duration": "1mo",
|
||||
"budget_reset_at": now,
|
||||
"budget_id": "team-member-budget-2",
|
||||
"created_at": now - timedelta(days=30),
|
||||
},
|
||||
)
|
||||
|
||||
budgets_to_reset = [test_budget]
|
||||
|
||||
asyncio.run(
|
||||
reset_budget_job.reset_budget_for_litellm_team_members(
|
||||
budgets_to_reset=budgets_to_reset
|
||||
)
|
||||
)
|
||||
|
||||
calls = mock_prisma_client.db.litellm_teammembership.update_many_calls
|
||||
assert len(calls) == 1
|
||||
|
||||
call = calls[0]
|
||||
# total_spend must NOT be present in the reset data — it must never be zeroed
|
||||
assert "total_spend" not in call["data"], (
|
||||
"total_spend should not be reset during a budget period reset; "
|
||||
"it must accumulate across all periods"
|
||||
)
|
||||
assert call["data"]["spend"] == 0
|
||||
|
||||
|
||||
def test_reset_budget_for_team_members_empty_budgets(
|
||||
reset_budget_job, mock_prisma_client
|
||||
):
|
||||
"""
|
||||
Test that reset_budget_for_litellm_team_members does nothing when given
|
||||
an empty list of budgets (e.g. no budget periods have expired).
|
||||
"""
|
||||
asyncio.run(
|
||||
reset_budget_job.reset_budget_for_litellm_team_members(budgets_to_reset=[])
|
||||
)
|
||||
|
||||
calls = mock_prisma_client.db.litellm_teammembership.update_many_calls
|
||||
assert len(calls) == 0
|
||||
|
||||
|
||||
def test_budget_table_reset_also_resets_linked_keys(
|
||||
reset_budget_job, mock_prisma_client
|
||||
):
|
||||
|
|
@ -479,3 +581,185 @@ def test_budget_table_reset_also_resets_linked_keys(
|
|||
)
|
||||
assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]}
|
||||
assert calls[0]["data"]["spend"] == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for _reset_budget_reset_at_date
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_reset_at_date_normal_rollforward():
|
||||
"""
|
||||
When budget_reset_at is set and has expired, roll it forward by one
|
||||
duration interval from the PREVIOUS reset time (not from now).
|
||||
This ensures periods stay aligned regardless of when the job runs.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_BudgetTableFull
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
previous_reset = now - timedelta(days=1) # expired yesterday
|
||||
|
||||
budget = LiteLLM_BudgetTableFull(
|
||||
budget_id="b1",
|
||||
budget_duration="7d",
|
||||
budget_reset_at=previous_reset,
|
||||
created_at=now - timedelta(days=30),
|
||||
)
|
||||
|
||||
result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now)
|
||||
|
||||
expected = previous_reset + timedelta(days=7)
|
||||
assert result.budget_reset_at == expected, (
|
||||
"budget_reset_at should advance by exactly one duration from the previous reset"
|
||||
)
|
||||
assert result.budget_reset_at > now, "new reset must be in the future"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_reset_at_date_none_reset_at_not_yet_expired():
|
||||
"""
|
||||
When budget_reset_at is None but created_at + duration is still in the
|
||||
future, anchor to created_at + duration (the original planned reset time).
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_BudgetTableFull
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
created_at = now - timedelta(days=3) # created 3 days ago
|
||||
# duration is 7d → first reset would be created_at + 7d = 4 days from now
|
||||
|
||||
budget = LiteLLM_BudgetTableFull(
|
||||
budget_id="b2",
|
||||
budget_duration="7d",
|
||||
budget_reset_at=None,
|
||||
created_at=created_at,
|
||||
)
|
||||
|
||||
result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now)
|
||||
|
||||
expected = created_at + timedelta(days=7)
|
||||
assert result.budget_reset_at == expected
|
||||
assert result.budget_reset_at > now
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_reset_at_date_none_reset_at_already_past():
|
||||
"""
|
||||
Regression: when budget_reset_at is None AND created_at + duration is
|
||||
already in the past (old budget that was never given a reset time), the
|
||||
new reset_at must be computed by advancing from created_at in duration
|
||||
steps until it's in the future — NOT set to current_time + duration.
|
||||
|
||||
Setting it to current_time + duration was the bug: it made the reset
|
||||
display show "proxy start time + duration" because this path runs at
|
||||
startup via the ResetBudgetJob scheduler.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_BudgetTableFull
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
# Budget was created 30 days ago with a 7-day duration → 4 full periods
|
||||
# have elapsed; created_at + 7d, +14d, +21d, +28d are all in the past.
|
||||
created_at = now - timedelta(days=30)
|
||||
|
||||
budget = LiteLLM_BudgetTableFull(
|
||||
budget_id="b3",
|
||||
budget_duration="7d",
|
||||
budget_reset_at=None,
|
||||
created_at=created_at,
|
||||
)
|
||||
|
||||
result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now)
|
||||
|
||||
assert result.budget_reset_at is not None
|
||||
assert result.budget_reset_at > now, "reset must be in the future"
|
||||
|
||||
# Must be anchored to created_at, not to `now`.
|
||||
# The correct next reset is created_at + N*7d where N is the smallest
|
||||
# integer that puts it past `now`. For 30 days elapsed: N=5 → +35d.
|
||||
expected = created_at + timedelta(days=35)
|
||||
assert result.budget_reset_at == expected, (
|
||||
f"Expected {expected}, got {result.budget_reset_at}. "
|
||||
"reset_at must be anchored to created_at, not proxy-start time."
|
||||
)
|
||||
|
||||
# Verify it is NOT simply now + duration (the old buggy behaviour).
|
||||
buggy_value = now + timedelta(days=7)
|
||||
assert result.budget_reset_at != pytest.approx(buggy_value, abs=timedelta(minutes=1)), (
|
||||
"reset_at must not be set to current_time + duration (proxy-start-time bug)"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_reset_at_date_multi_period_catchup():
|
||||
"""
|
||||
When budget_reset_at is non-null but multiple periods have elapsed (e.g.
|
||||
the proxy was down for 3 weeks and the period is 7 days), the while loop
|
||||
must advance the timestamp all the way to the next future date in a single
|
||||
call rather than requiring the scheduler to fire 3 times.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_BudgetTableFull
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
# reset_at is 22 days in the past; 7d period → 3 full periods have elapsed
|
||||
previous_reset = now - timedelta(days=22)
|
||||
|
||||
budget = LiteLLM_BudgetTableFull(
|
||||
budget_id="b-multi",
|
||||
budget_duration="7d",
|
||||
budget_reset_at=previous_reset,
|
||||
created_at=now - timedelta(days=30),
|
||||
)
|
||||
|
||||
result = await ResetBudgetJob._reset_budget_reset_at_date(budget, now)
|
||||
|
||||
assert result.budget_reset_at > now, "reset must be in the future after catch-up"
|
||||
# Should be previous_reset + 28d (4 periods) = now + 6d
|
||||
expected = previous_reset + timedelta(days=28)
|
||||
assert result.budget_reset_at == expected, (
|
||||
f"Expected {expected}, got {result.budget_reset_at}. "
|
||||
"Multi-period catch-up must advance through all elapsed periods in one call."
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_budget_at_startup_advances_expired_budget_reset_at(
|
||||
reset_budget_job, mock_prisma_client
|
||||
):
|
||||
"""
|
||||
Regression: on proxy restart, any budget whose budget_reset_at has just
|
||||
expired (e.g. budget created 24 h ago, reset_at = proxy_start_time) must be
|
||||
advanced by the immediate reset_budget() call that now runs at startup.
|
||||
|
||||
Without the startup call, a budget that expired at the exact moment the proxy
|
||||
restarted would show a past timestamp in the UI for up to ~10 minutes (the
|
||||
first scheduled interval), even if the user checked 50 minutes later.
|
||||
"""
|
||||
now = datetime.now(timezone.utc)
|
||||
proxy_start_time = now # simulated proxy start time
|
||||
|
||||
# Budget was created 24 h ago; reset_at is exactly now (proxy_start_time)
|
||||
expired_budget = type(
|
||||
"LiteLLM_BudgetTableFull",
|
||||
(),
|
||||
{
|
||||
"max_budget": 10.0,
|
||||
"budget_duration": "24h",
|
||||
"budget_reset_at": proxy_start_time, # expired right now
|
||||
"budget_id": "template-budget-1",
|
||||
"created_at": proxy_start_time - timedelta(hours=24),
|
||||
},
|
||||
)
|
||||
|
||||
mock_prisma_client.data["budget"] = [expired_budget]
|
||||
|
||||
# Simulate what the startup asyncio.create_task fires
|
||||
await reset_budget_job.reset_budget_for_litellm_budget_table()
|
||||
|
||||
# The budget should have been advanced in DB
|
||||
updated_budgets = mock_prisma_client.updated_data.get("budget", [])
|
||||
assert len(updated_budgets) == 1, "Expected the expired budget to be updated"
|
||||
|
||||
updated = updated_budgets[0]
|
||||
assert updated.budget_reset_at > proxy_start_time, (
|
||||
"After startup reset, budget_reset_at must be in the future, "
|
||||
"not equal to proxy start time"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -316,3 +316,105 @@ async def test_upsert_rpm_only_creates_new_budget(mock_tx, fake_user):
|
|||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# TEST: budget_duration sets budget_reset_at on the new budget
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_with_budget_duration_sets_reset_at(mock_tx, fake_user):
|
||||
"""
|
||||
When budget_duration is provided, the new individual member budget should
|
||||
have both budget_duration and budget_reset_at written so that the
|
||||
Next Budget Reset column shows a real future time instead of '-'.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
fake_reset_time = datetime(2026, 5, 1, 0, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.timezone_utils.get_budget_reset_time",
|
||||
return_value=fake_reset_time,
|
||||
) as mock_get_reset:
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-dur",
|
||||
user_id="user-dur",
|
||||
max_budget=10.0,
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_duration="24h",
|
||||
)
|
||||
|
||||
mock_get_reset.assert_called_once_with(budget_duration="24h")
|
||||
|
||||
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
|
||||
data={
|
||||
"max_budget": 10.0,
|
||||
"budget_duration": "24h",
|
||||
"budget_reset_at": fake_reset_time,
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
},
|
||||
include={"team_membership": True},
|
||||
)
|
||||
|
||||
|
||||
# TEST: budget_duration=None alone does not create a budget (disconnect path)
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_disconnect_when_only_budget_duration_none(mock_tx, fake_user):
|
||||
"""
|
||||
Passing only None values (including budget_duration=None) should still
|
||||
hit the disconnect path, not attempt to create a budget.
|
||||
"""
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-none",
|
||||
user_id="user-none",
|
||||
max_budget=None,
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
tpm_limit=None,
|
||||
rpm_limit=None,
|
||||
budget_duration=None,
|
||||
)
|
||||
|
||||
mock_tx.litellm_teammembership.update.assert_awaited_once()
|
||||
mock_tx.litellm_budgettable.create.assert_not_called()
|
||||
|
||||
|
||||
# TEST: budget_duration alone (no max_budget/limits) creates a budget
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_budget_duration_only_creates_budget(mock_tx, fake_user):
|
||||
"""
|
||||
Setting only budget_duration (no max_budget, tpm, rpm) should still create
|
||||
a budget with the duration so that the member gets a reset schedule.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import patch
|
||||
|
||||
fake_reset_time = datetime(2026, 5, 8, 0, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.timezone_utils.get_budget_reset_time",
|
||||
return_value=fake_reset_time,
|
||||
):
|
||||
await _upsert_budget_and_membership(
|
||||
mock_tx,
|
||||
team_id="team-dur-only",
|
||||
user_id="user-dur-only",
|
||||
max_budget=None,
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=fake_user,
|
||||
budget_duration="7d",
|
||||
)
|
||||
|
||||
mock_tx.litellm_budgettable.create.assert_awaited_once_with(
|
||||
data={
|
||||
"budget_duration": "7d",
|
||||
"budget_reset_at": fake_reset_time,
|
||||
"created_by": fake_user.user_id,
|
||||
"updated_by": fake_user.user_id,
|
||||
},
|
||||
include={"team_membership": True},
|
||||
)
|
||||
mock_tx.litellm_teammembership.update.assert_not_called()
|
||||
|
|
|
|||
|
|
@ -24,6 +24,12 @@ def client_and_mocks(monkeypatch):
|
|||
mock_table = MagicMock()
|
||||
mock_table.create = AsyncMock(side_effect=lambda *, data: data)
|
||||
mock_table.update = AsyncMock(side_effect=lambda *, where, data: {**where, **data})
|
||||
# Default find_unique: return a budget whose duration differs from "7d" so
|
||||
# duration-change tests see a change. Individual tests override this as needed.
|
||||
mock_table.find_unique = AsyncMock(return_value=types.SimpleNamespace(
|
||||
budget_duration=None,
|
||||
budget_reset_at=None,
|
||||
))
|
||||
|
||||
mock_prisma.db = types.SimpleNamespace(
|
||||
litellm_budgettable=mock_table,
|
||||
|
|
@ -265,3 +271,251 @@ async def test_new_budget_invalid_model_max_budget(client_and_mocks, monkeypatch
|
|||
assert resp.status_code in (400, 422), resp.text
|
||||
detail = resp.json()["detail"]
|
||||
assert "model_max_budget" in str(detail) or "dictionary" in str(detail).lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_duration_recalculates_budget_reset_at(client_and_mocks):
|
||||
"""
|
||||
Regression test: changing budget_duration via /budget/update must also update
|
||||
budget_reset_at to reflect the new period.
|
||||
|
||||
Previously, update_budget() stored the new duration but left the old
|
||||
budget_reset_at untouched, so the next reset shown in the UI was still in the
|
||||
past (the old period's reset time).
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
client, _, mock_table = client_and_mocks
|
||||
|
||||
captured_data: dict = {}
|
||||
|
||||
async def capture_update(*, where, data):
|
||||
captured_data.update(data)
|
||||
return {**where, **data}
|
||||
|
||||
mock_table.update = AsyncMock(side_effect=capture_update)
|
||||
|
||||
payload = {
|
||||
"budget_id": "budget_duration_change",
|
||||
"budget_duration": "7d",
|
||||
}
|
||||
resp = client.post("/budget/update", json=payload)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
# budget_reset_at must be present and in the future
|
||||
assert "budget_reset_at" in captured_data, (
|
||||
"budget_reset_at must be recalculated when budget_duration changes"
|
||||
)
|
||||
reset_at = captured_data["budget_reset_at"]
|
||||
if isinstance(reset_at, str):
|
||||
reset_at = datetime.fromisoformat(reset_at)
|
||||
if reset_at.tzinfo is None:
|
||||
reset_at = reset_at.replace(tzinfo=timezone.utc)
|
||||
assert reset_at > datetime.now(timezone.utc), (
|
||||
"budget_reset_at must be in the future after a duration change"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_duration_respects_explicit_budget_reset_at(client_and_mocks):
|
||||
"""
|
||||
If the caller explicitly provides budget_reset_at alongside budget_duration,
|
||||
the supplied value must be used as-is and not overwritten.
|
||||
"""
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
client, _, mock_table = client_and_mocks
|
||||
|
||||
captured_data: dict = {}
|
||||
|
||||
async def capture_update(*, where, data):
|
||||
captured_data.update(data)
|
||||
return {**where, **data}
|
||||
|
||||
mock_table.update = AsyncMock(side_effect=capture_update)
|
||||
|
||||
explicit_reset = (datetime.now(timezone.utc) + timedelta(days=30)).isoformat()
|
||||
payload = {
|
||||
"budget_id": "budget_explicit_reset",
|
||||
"budget_duration": "7d",
|
||||
"budget_reset_at": explicit_reset,
|
||||
}
|
||||
resp = client.post("/budget/update", json=payload)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
stored = captured_data.get("budget_reset_at")
|
||||
if isinstance(stored, str):
|
||||
stored = datetime.fromisoformat(stored)
|
||||
if isinstance(stored, datetime) and stored.tzinfo is None:
|
||||
stored = stored.replace(tzinfo=timezone.utc)
|
||||
|
||||
explicit_dt = datetime.fromisoformat(explicit_reset)
|
||||
# Should be within a second of the explicit value
|
||||
assert abs((stored - explicit_dt).total_seconds()) < 1, (
|
||||
"Explicit budget_reset_at must not be overwritten by the duration recalculation"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_duration_via_internal_call_path(client_and_mocks, monkeypatch):
|
||||
"""
|
||||
Regression test for the internal call path used by upsert_team_member_budget_table.
|
||||
|
||||
When the team member budget duration is changed via the team settings UI, the
|
||||
backend constructs a BudgetNewRequest by first calling BudgetNewRequest(budget_id=...)
|
||||
and then setting budget_duration via attribute assignment AFTER construction:
|
||||
|
||||
budget_request = BudgetNewRequest(budget_id=some_id)
|
||||
budget_request.budget_duration = "7d" # ← attribute assignment, not __init__
|
||||
await update_budget(budget_obj=budget_request, ...)
|
||||
|
||||
In Pydantic v2 without validate_assignment=True, attribute assignment does NOT
|
||||
reliably update model_fields_set. The previous fix checked model_fields_set and
|
||||
therefore silently skipped the budget_reset_at recalculation on this path.
|
||||
|
||||
The new fix checks budget_obj.budget_duration directly and writes budget_reset_at
|
||||
explicitly into the update dict, bypassing model_fields_set entirely.
|
||||
"""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
from datetime import datetime, timezone
|
||||
from litellm.proxy._types import BudgetNewRequest, UserAPIKeyAuth, LitellmUserRoles
|
||||
from litellm.proxy.management_endpoints.budget_management_endpoints import update_budget
|
||||
|
||||
import types as _types
|
||||
|
||||
_, mock_prisma, mock_table = client_and_mocks
|
||||
|
||||
# find_unique: existing budget has no duration (None → "7d" is a real change)
|
||||
mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace(
|
||||
budget_duration=None,
|
||||
budget_reset_at=None,
|
||||
))
|
||||
|
||||
captured_data: dict = {}
|
||||
|
||||
async def capture_update(*, where, data):
|
||||
captured_data.update(data)
|
||||
return {**where, **data}
|
||||
|
||||
mock_table.update = AsyncMock(side_effect=capture_update)
|
||||
|
||||
fake_user = UserAPIKeyAuth(
|
||||
user_id="internal_caller",
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
)
|
||||
|
||||
# Simulate exactly what upsert_team_member_budget_table does:
|
||||
# construct with budget_id only, then set budget_duration via attribute assignment.
|
||||
budget_request = BudgetNewRequest(budget_id="internal-budget-id")
|
||||
budget_request.budget_duration = "7d" # attribute assignment — not via __init__
|
||||
|
||||
await update_budget(budget_obj=budget_request, user_api_key_dict=fake_user)
|
||||
|
||||
assert "budget_reset_at" in captured_data, (
|
||||
"budget_reset_at must be recalculated even when budget_duration was set "
|
||||
"via attribute assignment after BudgetNewRequest construction"
|
||||
)
|
||||
assert "budget_duration" in captured_data, (
|
||||
"budget_duration must be present in the DB update"
|
||||
)
|
||||
reset_at = captured_data["budget_reset_at"]
|
||||
if isinstance(reset_at, str):
|
||||
reset_at = datetime.fromisoformat(reset_at)
|
||||
if reset_at.tzinfo is None:
|
||||
reset_at = reset_at.replace(tzinfo=timezone.utc)
|
||||
assert reset_at > datetime.now(timezone.utc), (
|
||||
"budget_reset_at must be a future datetime"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_same_duration_does_not_overwrite_reset_at(client_and_mocks):
|
||||
"""
|
||||
Regression test: proxy startup calls update_team → upsert_team_member_budget_table
|
||||
→ update_budget with the SAME budget_duration that's already stored in the DB.
|
||||
|
||||
Before this fix, update_budget always recalculated budget_reset_at when
|
||||
budget_duration was set, so every proxy restart would stamp the template
|
||||
budget's reset time to "now + duration", making the UI display proxy-start
|
||||
time as the "Next Budget Reset" for all members.
|
||||
|
||||
When the duration hasn't changed and budget_reset_at is already set to a
|
||||
future value, update_budget must NOT touch budget_reset_at.
|
||||
"""
|
||||
import types as _types
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
client, _, mock_table = client_and_mocks
|
||||
|
||||
future_reset = datetime.now(timezone.utc) + timedelta(days=5)
|
||||
|
||||
# Simulate existing budget with the SAME duration and a valid future reset time
|
||||
mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace(
|
||||
budget_duration="7d",
|
||||
budget_reset_at=future_reset,
|
||||
))
|
||||
|
||||
captured_data: dict = {}
|
||||
|
||||
async def capture_update(*, where, data):
|
||||
captured_data.update(data)
|
||||
return {**where, **data}
|
||||
|
||||
mock_table.update = AsyncMock(side_effect=capture_update)
|
||||
|
||||
payload = {
|
||||
"budget_id": "template-budget-startup",
|
||||
"budget_duration": "7d", # same as what's already in DB
|
||||
"max_budget": 50.0,
|
||||
}
|
||||
resp = client.post("/budget/update", json=payload)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
assert "budget_reset_at" not in captured_data, (
|
||||
"budget_reset_at must NOT be overwritten when budget_duration is unchanged "
|
||||
"and a valid future reset time already exists — proxy-start-time bug"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_budget_changed_duration_recalculates_reset_at(client_and_mocks):
|
||||
"""
|
||||
Complementary to the no-op test: when budget_duration genuinely changes
|
||||
(e.g. user switches from 7d to 1mo in the UI), budget_reset_at must be
|
||||
recalculated to reflect the new period.
|
||||
"""
|
||||
import types as _types
|
||||
from datetime import datetime, timezone, timedelta
|
||||
|
||||
client, _, mock_table = client_and_mocks
|
||||
|
||||
# Existing budget has "7d"; request changes it to "30d"
|
||||
mock_table.find_unique = AsyncMock(return_value=_types.SimpleNamespace(
|
||||
budget_duration="7d",
|
||||
budget_reset_at=datetime.now(timezone.utc) + timedelta(days=3),
|
||||
))
|
||||
|
||||
captured_data: dict = {}
|
||||
|
||||
async def capture_update(*, where, data):
|
||||
captured_data.update(data)
|
||||
return {**where, **data}
|
||||
|
||||
mock_table.update = AsyncMock(side_effect=capture_update)
|
||||
|
||||
payload = {
|
||||
"budget_id": "template-budget-change",
|
||||
"budget_duration": "30d", # different from DB value
|
||||
}
|
||||
resp = client.post("/budget/update", json=payload)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
assert "budget_reset_at" in captured_data, (
|
||||
"budget_reset_at must be recalculated when budget_duration changes"
|
||||
)
|
||||
reset_at = captured_data["budget_reset_at"]
|
||||
if isinstance(reset_at, str):
|
||||
reset_at = datetime.fromisoformat(reset_at)
|
||||
if reset_at.tzinfo is None:
|
||||
reset_at = reset_at.replace(tzinfo=timezone.utc)
|
||||
assert reset_at > datetime.now(timezone.utc)
|
||||
|
|
|
|||
|
|
@ -486,3 +486,75 @@ class TestSetObjectMetadataField:
|
|||
):
|
||||
_set_object_metadata_field(team, "model_rpm_limit", {"x": 1})
|
||||
assert team.metadata == {"model_rpm_limit": {"x": 1}}
|
||||
|
||||
|
||||
class TestUpsertBudgetAndMembershipClearDuration:
|
||||
"""
|
||||
Tests for the clear_budget_duration flag in _upsert_budget_and_membership.
|
||||
Regression for the bug where selecting 'No reset (unlimited)' in the UI
|
||||
could not clear a previously-set budget_duration because null was filtered
|
||||
out before reaching the backend.
|
||||
"""
|
||||
|
||||
def _make_tx(self):
|
||||
tx = MagicMock()
|
||||
tx.litellm_budgettable = MagicMock()
|
||||
tx.litellm_budgettable.update = AsyncMock()
|
||||
tx.litellm_budgettable.create = AsyncMock(
|
||||
return_value=MagicMock(budget_id="new-budget-id")
|
||||
)
|
||||
tx.litellm_teammembership = MagicMock()
|
||||
tx.litellm_teammembership.update = AsyncMock()
|
||||
tx.litellm_teammembership.upsert = AsyncMock()
|
||||
return tx
|
||||
|
||||
def _make_auth(self):
|
||||
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_budget_duration_sets_null_on_existing_budget(self):
|
||||
"""clear_budget_duration=True must null out budget_duration and budget_reset_at."""
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_upsert_budget_and_membership,
|
||||
)
|
||||
|
||||
tx = self._make_tx()
|
||||
await _upsert_budget_and_membership(
|
||||
tx=tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
max_budget=None,
|
||||
existing_budget_id="existing-budget-id",
|
||||
user_api_key_dict=self._make_auth(),
|
||||
budget_duration=None,
|
||||
clear_budget_duration=True,
|
||||
)
|
||||
|
||||
tx.litellm_budgettable.update.assert_awaited_once()
|
||||
update_data = tx.litellm_budgettable.update.call_args.kwargs["data"]
|
||||
assert update_data["budget_duration"] is None
|
||||
assert update_data["budget_reset_at"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_clear_flag_does_not_touch_budget_duration(self):
|
||||
"""Without clear_budget_duration, omitting budget_duration must not overwrite it."""
|
||||
from litellm.proxy.management_endpoints.common_utils import (
|
||||
_upsert_budget_and_membership,
|
||||
)
|
||||
|
||||
tx = self._make_tx()
|
||||
await _upsert_budget_and_membership(
|
||||
tx=tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
max_budget=50.0,
|
||||
existing_budget_id="existing-budget-id",
|
||||
user_api_key_dict=self._make_auth(),
|
||||
budget_duration=None,
|
||||
clear_budget_duration=False,
|
||||
)
|
||||
|
||||
tx.litellm_budgettable.update.assert_awaited_once()
|
||||
update_data = tx.litellm_budgettable.update.call_args.kwargs["data"]
|
||||
assert "budget_duration" not in update_data
|
||||
assert "budget_reset_at" not in update_data
|
||||
|
|
|
|||
|
|
@ -1627,6 +1627,7 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
|
|||
|
||||
team_table = MagicMock(spec=LiteLLM_TeamTable)
|
||||
team_table.metadata = {}
|
||||
team_table.team_id = "test_team_id"
|
||||
team_table.team_alias = "Test Team"
|
||||
team_table.budget_duration = None
|
||||
|
||||
|
|
@ -1639,12 +1640,22 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
|
|||
mock_budget_response = MagicMock()
|
||||
mock_budget_response.budget_id = "new_budget_456"
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
|
||||
new_callable=AsyncMock
|
||||
) as mock_new_budget:
|
||||
mock_new_budget.return_value = mock_budget_response
|
||||
mock_db = MagicMock()
|
||||
mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 0})
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = mock_db
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_budget_response,
|
||||
) as mock_new_budget,
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
mock_prisma,
|
||||
),
|
||||
):
|
||||
result = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
|
|
@ -1661,6 +1672,125 @@ async def test_upsert_team_member_budget_table_no_existing_budget():
|
|||
assert "team_member_budget_duration" not in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_team_member_budget_table_wires_up_null_budget_members_on_create():
|
||||
"""
|
||||
Regression test: when a team budget is being CREATED for the first time,
|
||||
existing team members whose budget_id is null must be assigned the new
|
||||
template budget so that their 'Next Budget Reset' column reflects the reset time.
|
||||
|
||||
Previously, members added before the team's budget duration was configured had
|
||||
budget_id=null. Creating the template budget left those memberships disconnected
|
||||
— they never showed a reset date.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
mock_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
|
||||
)
|
||||
|
||||
# No existing budget_id — triggers the create path
|
||||
team_table = MagicMock(spec=LiteLLM_TeamTable)
|
||||
team_table.metadata = {}
|
||||
team_table.team_id = "team-xyz"
|
||||
team_table.team_alias = "Test Team"
|
||||
team_table.budget_duration = None
|
||||
|
||||
updated_kv: dict = {"team_id": "team-xyz"}
|
||||
|
||||
mock_budget_response = MagicMock()
|
||||
mock_budget_response.budget_id = "new-template-budget-abc"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 2})
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = mock_db
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.new_budget",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_budget_response,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
mock_prisma,
|
||||
),
|
||||
):
|
||||
await TeamMemberBudgetHandler.upsert_team_member_budget_table(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=mock_user,
|
||||
updated_kv=updated_kv,
|
||||
team_member_budget_duration="7d",
|
||||
)
|
||||
|
||||
# update_many must have been called to wire up members with no budget
|
||||
mock_db.litellm_teammembership.update_many.assert_awaited_once()
|
||||
call_kwargs = mock_db.litellm_teammembership.update_many.call_args.kwargs
|
||||
assert call_kwargs["where"]["team_id"] == "team-xyz"
|
||||
assert call_kwargs["where"]["budget_id"] is None
|
||||
assert call_kwargs["data"]["budget_id"] == "new-template-budget-abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_team_member_budget_table_update_does_not_wire_up_null_members():
|
||||
"""
|
||||
When the team already has a template budget and the settings are UPDATED,
|
||||
members with budget_id=null must NOT be silently assigned the budget.
|
||||
Those members were intentionally left unlimited; assigning a budget on
|
||||
every settings save would be a backwards-incompatible behaviour change.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, LiteLLM_TeamTable, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
|
||||
|
||||
mock_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test_user_id"
|
||||
)
|
||||
|
||||
# Existing budget_id — triggers the update path
|
||||
team_table = MagicMock(spec=LiteLLM_TeamTable)
|
||||
team_table.metadata = {"team_member_budget_id": "template-budget-abc"}
|
||||
team_table.team_id = "team-xyz"
|
||||
|
||||
updated_kv: dict = {"team_id": "team-xyz"}
|
||||
|
||||
mock_budget_response = MagicMock()
|
||||
mock_budget_response.budget_id = "template-budget-abc"
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.litellm_teammembership.update_many = AsyncMock(return_value={"count": 0})
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = mock_db
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.budget_management_endpoints.update_budget",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_budget_response,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.prisma_client",
|
||||
mock_prisma,
|
||||
),
|
||||
):
|
||||
await TeamMemberBudgetHandler.upsert_team_member_budget_table(
|
||||
team_table=team_table,
|
||||
user_api_key_dict=mock_user,
|
||||
updated_kv=updated_kv,
|
||||
team_member_budget_duration="7d",
|
||||
)
|
||||
|
||||
# update_many must NOT have been called — no silent budget assignment on update
|
||||
mock_db.litellm_teammembership.update_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_team_with_team_member_budget_duration():
|
||||
"""
|
||||
|
|
@ -6499,6 +6629,84 @@ async def test_create_team_member_budget_table_with_duration():
|
|||
assert result["metadata"]["team_member_budget_id"] == "budget-abc"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_budget_and_membership_updates_existing_budget():
|
||||
"""
|
||||
When existing_budget_id is provided, _upsert_budget_and_membership must UPDATE
|
||||
the existing budget row rather than creating a new one — otherwise each call
|
||||
silently orphans the previous LiteLLM_BudgetTable row.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership
|
||||
|
||||
mock_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"
|
||||
)
|
||||
|
||||
mock_tx = MagicMock()
|
||||
mock_tx.litellm_budgettable.update = AsyncMock()
|
||||
mock_tx.litellm_budgettable.create = AsyncMock()
|
||||
mock_tx.litellm_teammembership.upsert = AsyncMock()
|
||||
|
||||
await _upsert_budget_and_membership(
|
||||
tx=mock_tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
max_budget=50.0,
|
||||
existing_budget_id="budget-existing-123",
|
||||
user_api_key_dict=mock_user,
|
||||
)
|
||||
|
||||
# Must update, not create
|
||||
mock_tx.litellm_budgettable.update.assert_awaited_once()
|
||||
mock_tx.litellm_budgettable.create.assert_not_awaited()
|
||||
|
||||
# Membership must be connected to the existing budget id
|
||||
upsert_call = mock_tx.litellm_teammembership.upsert.call_args
|
||||
connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"]
|
||||
assert connect_id == "budget-existing-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_upsert_budget_and_membership_creates_new_budget_when_none():
|
||||
"""
|
||||
When existing_budget_id is None, _upsert_budget_and_membership must CREATE
|
||||
a new budget row and connect it to the membership.
|
||||
"""
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.management_endpoints.common_utils import _upsert_budget_and_membership
|
||||
|
||||
mock_user = UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"
|
||||
)
|
||||
|
||||
mock_new_budget = MagicMock()
|
||||
mock_new_budget.budget_id = "budget-new-456"
|
||||
|
||||
mock_tx = MagicMock()
|
||||
mock_tx.litellm_budgettable.update = AsyncMock()
|
||||
mock_tx.litellm_budgettable.create = AsyncMock(return_value=mock_new_budget)
|
||||
mock_tx.litellm_teammembership.upsert = AsyncMock()
|
||||
|
||||
await _upsert_budget_and_membership(
|
||||
tx=mock_tx,
|
||||
team_id="team-1",
|
||||
user_id="user-1",
|
||||
max_budget=100.0,
|
||||
existing_budget_id=None,
|
||||
user_api_key_dict=mock_user,
|
||||
)
|
||||
|
||||
# Must create, not update
|
||||
mock_tx.litellm_budgettable.create.assert_awaited_once()
|
||||
mock_tx.litellm_budgettable.update.assert_not_awaited()
|
||||
|
||||
# Membership must be connected to the newly created budget id
|
||||
upsert_call = mock_tx.litellm_teammembership.upsert.call_args
|
||||
connect_id = upsert_call.kwargs["data"]["update"]["litellm_budget_table"]["connect"]["budget_id"]
|
||||
assert connect_id == "budget-new-456"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for _batch_resolve_access_group_resources
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -3759,6 +3759,7 @@ export interface Member {
|
|||
max_budget_in_team?: number | null;
|
||||
tpm_limit?: number | null;
|
||||
rpm_limit?: number | null;
|
||||
budget_duration?: string | null;
|
||||
}
|
||||
|
||||
export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => {
|
||||
|
|
@ -3898,6 +3899,10 @@ export const teamMemberUpdateCall = async (
|
|||
if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) {
|
||||
requestBody.rpm_limit = formValues.rpm_limit;
|
||||
}
|
||||
if (formValues.budget_duration !== undefined) {
|
||||
// Explicitly include null so the backend can clear a previously-set duration.
|
||||
requestBody.budget_duration = formValues.budget_duration;
|
||||
}
|
||||
|
||||
console.log("Final request body:", requestBody);
|
||||
|
||||
|
|
|
|||
|
|
@ -85,8 +85,8 @@ const MemberModal = <T extends BaseMember>({
|
|||
const formData = Object.entries(values).reduce((acc, [key, value]) => {
|
||||
if (typeof value === "string") {
|
||||
const trimmedValue = value.trim();
|
||||
// For empty strings on optional numeric fields, set to null
|
||||
if (trimmedValue === "" && (key === "max_budget_in_team" || key === "tpm_limit" || key === "rpm_limit")) {
|
||||
// For empty strings on optional fields, set to null
|
||||
if (trimmedValue === "" && (key === "max_budget_in_team" || key === "tpm_limit" || key === "rpm_limit" || key === "budget_duration")) {
|
||||
return { ...acc, [key]: null };
|
||||
}
|
||||
return { ...acc, [key]: trimmedValue };
|
||||
|
|
|
|||
|
|
@ -57,6 +57,7 @@ export interface TeamMembership {
|
|||
team_id: string;
|
||||
budget_id: string;
|
||||
spend: number;
|
||||
total_spend: number | null;
|
||||
litellm_budget_table: {
|
||||
budget_id: string;
|
||||
soft_budget: number | null;
|
||||
|
|
@ -66,6 +67,7 @@ export interface TeamMembership {
|
|||
rpm_limit: number | null;
|
||||
model_max_budget: Record<string, number> | null;
|
||||
budget_duration: string | null;
|
||||
budget_reset_at: string | null;
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -381,6 +383,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
max_budget_in_team: values.max_budget_in_team,
|
||||
tpm_limit: values.tpm_limit,
|
||||
rpm_limit: values.rpm_limit,
|
||||
budget_duration: values.budget_duration || null,
|
||||
};
|
||||
MessageManager.destroy(); // Remove all existing toasts
|
||||
|
||||
|
|
@ -1483,6 +1486,24 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
min: 0,
|
||||
placeholder: "Requests per minute limit for this member in this team",
|
||||
},
|
||||
{
|
||||
name: "budget_duration",
|
||||
label: (
|
||||
<span>
|
||||
Budget Reset Period{" "}
|
||||
<Tooltip title="How often this member's individual spend resets. Leave unset to inherit from team settings.">
|
||||
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
|
||||
</Tooltip>
|
||||
</span>
|
||||
),
|
||||
type: "select" as const,
|
||||
options: [
|
||||
{ label: "No reset (unlimited)", value: "" },
|
||||
{ label: "Daily", value: "24h" },
|
||||
{ label: "Weekly", value: "7d" },
|
||||
{ label: "Monthly", value: "30d" },
|
||||
],
|
||||
},
|
||||
],
|
||||
}}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -45,13 +45,27 @@ export default function TeamMemberTab({
|
|||
return "0";
|
||||
};
|
||||
|
||||
// Helper function to get spend for a user
|
||||
// Helper function to get spend for a user (current budget period only)
|
||||
const getUserSpend = (userId: string | null): number | null => {
|
||||
if (!userId) return 0;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return membership?.spend || 0;
|
||||
};
|
||||
|
||||
// Helper function to get total (all-time) spend for a user
|
||||
const getUserTotalSpend = (userId: string | null): number | null => {
|
||||
if (!userId) return 0;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return membership?.total_spend ?? 0;
|
||||
};
|
||||
|
||||
// Helper function to get the next budget reset datetime for a user
|
||||
const getUserBudgetResetAt = (userId: string | null): string | null => {
|
||||
if (!userId) return null;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return membership?.litellm_budget_table?.budget_reset_at ?? null;
|
||||
};
|
||||
|
||||
const getUserBudget = (userId: string | null): string | null => {
|
||||
if (!userId) return null;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
|
|
@ -86,8 +100,8 @@ export default function TeamMemberTab({
|
|||
{
|
||||
title: (
|
||||
<Space direction="horizontal">
|
||||
Team Member Spend (USD)
|
||||
<Tooltip title="This is the amount spent by a user in the team.">
|
||||
Spend (Current Period)
|
||||
<Tooltip title="Amount spent by this member in the current budget period. Resets when the budget period ends.">
|
||||
<InfoCircleOutlined />
|
||||
</Tooltip>
|
||||
</Space>
|
||||
|
|
@ -97,6 +111,20 @@ export default function TeamMemberTab({
|
|||
<Typography.Text>${formatNumberWithCommas(getUserSpend(record.user_id), 4)}</Typography.Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
title: (
|
||||
<Space direction="horizontal">
|
||||
Total Spend (USD)
|
||||
<Tooltip title="Cumulative amount spent by this member across all budget periods.">
|
||||
<InfoCircleOutlined />
|
||||
</Tooltip>
|
||||
</Space>
|
||||
),
|
||||
key: "total_spend",
|
||||
render: (_: unknown, record: Member) => (
|
||||
<Typography.Text>${formatNumberWithCommas(getUserTotalSpend(record.user_id), 4)}</Typography.Text>
|
||||
),
|
||||
},
|
||||
{
|
||||
title: "Team Member Budget (USD)",
|
||||
key: "budget",
|
||||
|
|
@ -109,6 +137,26 @@ export default function TeamMemberTab({
|
|||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
title: (
|
||||
<Space direction="horizontal">
|
||||
Next Budget Reset
|
||||
<Tooltip title="When this member's spend resets for the next budget period.">
|
||||
<InfoCircleOutlined />
|
||||
</Tooltip>
|
||||
</Space>
|
||||
),
|
||||
key: "budget_reset_at",
|
||||
render: (_: unknown, record: Member) => {
|
||||
const resetAt = getUserBudgetResetAt(record.user_id);
|
||||
if (!resetAt) return <Typography.Text>-</Typography.Text>;
|
||||
return (
|
||||
<Typography.Text>
|
||||
{new Date(resetAt).toLocaleString()}
|
||||
</Typography.Text>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
title: (
|
||||
<Space direction="horizontal">
|
||||
|
|
@ -138,6 +186,7 @@ export default function TeamMemberTab({
|
|||
max_budget_in_team: membership?.litellm_budget_table?.max_budget || null,
|
||||
tpm_limit: membership?.litellm_budget_table?.tpm_limit || null,
|
||||
rpm_limit: membership?.litellm_budget_table?.rpm_limit || null,
|
||||
budget_duration: membership?.litellm_budget_table?.budget_duration || null,
|
||||
};
|
||||
setSelectedEditMember(enhancedMember);
|
||||
setIsEditMemberModalVisible(true);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue