mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): reconcile team member total_spend with current cycle spend (#48260)
* fix(proxy): reconcile team member total_spend with current cycle spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): read /team/list as a list and drain the scripted provider Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(repositories): compose the membership reset writer and drop SQL text asserts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: type the cascade repair test and flatten memberships without a stacked comprehension Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: keep the reset job test file's existing typing untouched Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: joshua <joshua@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
84754cca0a
commit
24a7bf6af7
16 changed files with 655 additions and 37 deletions
|
|
@ -65,6 +65,7 @@ from litellm.repositories.table_repositories import (
|
|||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.unit_of_work import (
|
||||
LinkedSpendResetWrites,
|
||||
SpendResetWrites,
|
||||
budget_cascade_unit_of_work,
|
||||
enduser_spend_reset_unit_of_work,
|
||||
spend_reset_unit_of_work,
|
||||
|
|
@ -235,7 +236,7 @@ def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
|
|||
|
||||
|
||||
def _queue_budget_linked_resets(
|
||||
writes: LinkedSpendResetWrites,
|
||||
writes: SpendResetWrites,
|
||||
cascade: "_BudgetCascade",
|
||||
extra: Mapping[str, object] = MappingProxyType({}),
|
||||
) -> None:
|
||||
|
|
@ -908,6 +909,7 @@ class ResetBudgetJob:
|
|||
db_span("reset_budget_cascade", "LiteLLM_BudgetTable"),
|
||||
budget_cascade_unit_of_work(self._new_batch) as uow,
|
||||
):
|
||||
uow.team_memberships.queue_total_spend_repair(cascade.budget_ids)
|
||||
_queue_budget_linked_resets(uow.team_memberships, cascade)
|
||||
_queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE)
|
||||
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
|
||||
|
|
|
|||
|
|
@ -472,6 +472,8 @@ _TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS
|
|||
# recreates the member. The rows travel as one JSON document, not as a numeric array: Prisma
|
||||
# types a raw array parameter from the first batch a connection sees, so after an all-$0 batch
|
||||
# (integers) every later fractional batch on that connection failed with "improper binary format".
|
||||
# total_spend is floored at spend before adding: rows written before total_spend existed (migration
|
||||
# 20260421135425, default 0) carry their history only in spend, so a plain increment keeps them behind forever.
|
||||
_TEAM_MEMBER_SPEND_SQL: Final = """
|
||||
INSERT INTO "LiteLLM_TeamMembership" (user_id, team_id, spend, total_spend)
|
||||
SELECT member.user_id, member.team_id, member.cost, member.cost
|
||||
|
|
@ -487,7 +489,7 @@ WHERE EXISTS (
|
|||
)
|
||||
ON CONFLICT (user_id, team_id) DO UPDATE
|
||||
SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend,
|
||||
total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend
|
||||
total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend) + EXCLUDED.total_spend
|
||||
"""
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1013,7 +1013,7 @@ async def get_all_team_memberships(
|
|||
|
||||
returned_tm: Final[list[LiteLLM_TeamMembership]] = []
|
||||
for tm in team_memberships:
|
||||
returned_tm.append(LiteLLM_TeamMembership.model_validate(tm.model_dump()))
|
||||
returned_tm.append(with_total_spend_floor(LiteLLM_TeamMembership.model_validate(tm.model_dump())))
|
||||
|
||||
return returned_tm
|
||||
|
||||
|
|
@ -4202,6 +4202,23 @@ def _check_not_resetting_own_spend(user_id: str, user_api_key_dict: UserAPIKeyAu
|
|||
_raise_reset_spend_error(status.HTTP_403_FORBIDDEN, "Cannot reset your own spend. Ask a proxy admin.")
|
||||
|
||||
|
||||
# One statement so the validated spend cannot go stale between the read and the write: the WHERE re-checks
|
||||
# reset_to <= spend, and total_spend is floored at spend before spend drops so the history survives the reset.
|
||||
_RESET_TEAM_MEMBER_SPEND_SQL: Final = """
|
||||
UPDATE "LiteLLM_TeamMembership"
|
||||
SET total_spend = GREATEST(total_spend, spend), spend = $3::float8
|
||||
WHERE user_id = $1 AND team_id = $2 AND spend >= $3::float8
|
||||
"""
|
||||
|
||||
|
||||
def with_total_spend_floor(membership: LiteLLM_TeamMembership) -> LiteLLM_TeamMembership:
|
||||
"""Rows written before total_spend existed (migration 20260421135425) hold their history only in spend."""
|
||||
spend: Final = membership.spend or 0.0
|
||||
if (membership.total_spend or 0.0) >= spend:
|
||||
return membership
|
||||
return membership.model_copy(update={"total_spend": spend})
|
||||
|
||||
|
||||
def _raise_reset_spend_error(status_code: int, message: str) -> NoReturn:
|
||||
detail: Final = {"error": message}
|
||||
raise HTTPException(status_code=status_code, detail=detail)
|
||||
|
|
@ -4286,10 +4303,12 @@ async def reset_team_member_spend_fn(
|
|||
current_spend: Final = membership.spend or 0.0
|
||||
reset_to: Final = _validate_team_member_reset_spend_value(data.reset_to, membership)
|
||||
|
||||
await _team_membership_db(prisma_client).update(
|
||||
where=membership_where,
|
||||
data={"spend": reset_to},
|
||||
)
|
||||
reset_rows: Final = await _execute_raw(prisma_client)(_RESET_TEAM_MEMBER_SPEND_SQL, user_id, team_id, reset_to)
|
||||
if reset_rows == 0:
|
||||
_raise_reset_spend_error(
|
||||
status.HTTP_409_CONFLICT,
|
||||
f"Member spend changed since it was read; reset_to ({reset_to}) must be <= current spend. Retry.",
|
||||
)
|
||||
|
||||
await invalidate_team_member_spend_state(
|
||||
user_id=user_id,
|
||||
|
|
@ -5548,16 +5567,17 @@ async def team_member_me(
|
|||
litellm_budget_table=None,
|
||||
)
|
||||
|
||||
floored: Final = with_total_spend_floor(membership)
|
||||
return TeamMemberInfoResponse(
|
||||
user_id=caller_user_id,
|
||||
team_id=team_id,
|
||||
team_alias=team_table.team_alias,
|
||||
role=member_role,
|
||||
user_email=user_email,
|
||||
spend=membership.spend,
|
||||
total_spend=membership.total_spend,
|
||||
budget_id=membership.budget_id,
|
||||
litellm_budget_table=membership.litellm_budget_table,
|
||||
spend=floored.spend,
|
||||
total_spend=floored.total_spend,
|
||||
budget_id=floored.budget_id,
|
||||
litellm_budget_table=floored.litellm_budget_table,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -171,4 +171,6 @@ class PrismaBatch(Protocol):
|
|||
@property
|
||||
def litellm_projecttable(self) -> BatchTable: ...
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None: ...
|
||||
|
||||
async def commit(self) -> None: ...
|
||||
|
|
|
|||
|
|
@ -15,11 +15,11 @@ instead of full-model writes, which trip ``prisma.errors.DataError`` on rows
|
|||
carrying fields the update input type rejects (see #27730).
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Final
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
|
||||
from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch
|
||||
|
||||
|
|
@ -62,6 +62,13 @@ class TeamSpendResetWrites:
|
|||
)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class SpendResetWrites(Protocol):
|
||||
def queue_spend_zero(self, where: Mapping[str, object]) -> None: ...
|
||||
|
||||
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LinkedSpendResetWrites:
|
||||
table: BatchTable
|
||||
|
|
@ -78,6 +85,30 @@ class LinkedSpendResetWrites:
|
|||
)
|
||||
|
||||
|
||||
# Membership rows written before total_spend existed (migration 20260421135425, default 0) hold their
|
||||
# history only in spend. Zeroing spend would erase it, so the floor lands in the same batch, ahead of the reset.
|
||||
_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL: Final = """
|
||||
UPDATE "LiteLLM_TeamMembership"
|
||||
SET total_spend = spend
|
||||
WHERE budget_id = ANY($1::text[]) AND spend > total_spend
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TeamMembershipSpendResetWrites:
|
||||
resets: LinkedSpendResetWrites
|
||||
batch: PrismaBatch
|
||||
|
||||
def queue_spend_zero(self, where: Mapping[str, object]) -> None:
|
||||
self.resets.queue_spend_zero(where)
|
||||
|
||||
def queue_spend_decrement(self, where: Mapping[str, object], amount: float) -> None:
|
||||
self.resets.queue_spend_decrement(where, amount)
|
||||
|
||||
def queue_total_spend_repair(self, budget_ids: Sequence[str]) -> None:
|
||||
self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BudgetWindowWrites:
|
||||
table: BatchTable
|
||||
|
|
@ -106,7 +137,7 @@ class BudgetCascadeUnitOfWork:
|
|||
``enduser_spend_reset_unit_of_work`` because one tier can gate millions.
|
||||
"""
|
||||
|
||||
team_memberships: LinkedSpendResetWrites
|
||||
team_memberships: TeamMembershipSpendResetWrites
|
||||
keys: LinkedSpendResetWrites
|
||||
organizations: LinkedSpendResetWrites
|
||||
tags: LinkedSpendResetWrites
|
||||
|
|
@ -132,7 +163,9 @@ async def budget_cascade_unit_of_work(
|
|||
) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]:
|
||||
batch = new_batch()
|
||||
yield BudgetCascadeUnitOfWork(
|
||||
team_memberships=LinkedSpendResetWrites(table=batch.litellm_teammembership),
|
||||
team_memberships=TeamMembershipSpendResetWrites(
|
||||
resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch
|
||||
),
|
||||
keys=LinkedSpendResetWrites(table=batch.litellm_verificationtoken),
|
||||
organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
|
||||
tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,273 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from tests.integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value
|
||||
from tests.integration._support.database import read_rows, write_rows
|
||||
from tests.integration._support.provider import PROVIDER_URL, SharedProvider
|
||||
from tests.integration._support.wire import Reply
|
||||
|
||||
_CALL_COST: Final = 0.02
|
||||
JSON_VALUE: Final = TypeAdapter(JsonValue)
|
||||
_LEGACY_SPEND: Final = 612.4
|
||||
_LEGACY_TOTAL: Final = 148.1
|
||||
_HEALTHY_SPEND: Final = 12.25
|
||||
_HEALTHY_TOTAL: Final = 97.8
|
||||
|
||||
|
||||
def _completion() -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "scripted"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _priced_model(scenario: Scenario) -> str:
|
||||
return scenario.model(api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
|
||||
|
||||
def _member_with_budget(gateway: Gateway, scenario: Scenario, team_id: str, *, max_budget: float) -> str:
|
||||
user_id: Final = scenario.user(user_role="internal_user")
|
||||
gateway.post(
|
||||
"/team/member_add",
|
||||
{"team_id": team_id, "member": {"role": "user", "user_id": user_id}, "max_budget_in_team": max_budget},
|
||||
)
|
||||
return user_id
|
||||
|
||||
|
||||
def _seed_membership(user_id: str, team_id: str, *, spend: float, total_spend: float) -> None:
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_TeamMembership" SET spend = %s::float8, total_spend = %s::float8 '
|
||||
"WHERE user_id = %s AND team_id = %s",
|
||||
(str(spend), str(total_spend), user_id, team_id),
|
||||
)
|
||||
|
||||
|
||||
def _membership(user_id: str, team_id: str) -> tuple[float, float]:
|
||||
rows: Final = read_rows(
|
||||
'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s',
|
||||
(user_id, team_id),
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
return float(str(rows[0]["spend"])), float(str(rows[0]["total_spend"]))
|
||||
|
||||
|
||||
def _membership_reaches(user_id: str, team_id: str, *, total_spend: float) -> tuple[float, float]:
|
||||
return eventually(
|
||||
lambda: _membership(user_id, team_id),
|
||||
lambda row: row[1] == pytest.approx(total_spend),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _served(gateway: Gateway, provider: SharedProvider, model: str, key: str) -> None:
|
||||
provider.expect(_completion())
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"reconcile {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert len(provider.received()) == 1
|
||||
|
||||
|
||||
def _reset_spend(gateway: Gateway, team_id: str, user_id: str, reset_to: float) -> httpx.Response:
|
||||
return gateway.request("POST", f"/team/{team_id}/member/{user_id}/reset_spend", {"reset_to": reset_to})
|
||||
|
||||
|
||||
def _membership_figures(row: JsonValue) -> tuple[str, tuple[float, float]]:
|
||||
membership: Final = object_value(row)
|
||||
return string_value(membership["user_id"]), (
|
||||
float(str(membership["spend"])),
|
||||
float(str(membership["total_spend"])),
|
||||
)
|
||||
|
||||
|
||||
def _reported_memberships(gateway: Gateway, path: str, params: dict[str, str]) -> dict[str, tuple[float, float]]:
|
||||
"""/team/info answers with one team object, /team/list with a list of them."""
|
||||
response: Final = gateway.request("GET", path, params=params)
|
||||
assert response.status_code == 200, f"GET {path}: {response.status_code} {response.text}"
|
||||
body: Final = JSON_VALUE.validate_json(response.content)
|
||||
teams: Final = [object_value(body)] if isinstance(body, dict) else [object_value(team) for team in list_value(body)]
|
||||
rows: Final = [
|
||||
row for team in teams for row in list_value(team["team_memberships"])
|
||||
] # comprehension-ok: flatten memberships across teams
|
||||
return dict(_membership_figures(row) for row in rows)
|
||||
|
||||
|
||||
def test_a_flush_lifts_a_legacy_total_to_the_cycle_spend_and_leaves_a_healthy_row_on_its_baseline(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
|
||||
legacy_key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
healthy_key: Final = scenario.key(user_id=healthy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, legacy_key)
|
||||
_served(gateway, provider, model, healthy_key)
|
||||
|
||||
legacy_row: Final = _membership_reaches(
|
||||
legacy, team_id, total_spend=max(_LEGACY_SPEND, _LEGACY_TOTAL) + _CALL_COST
|
||||
)
|
||||
assert legacy_row[0] == pytest.approx(_LEGACY_SPEND + _CALL_COST)
|
||||
healthy_row: Final = _membership_reaches(healthy, team_id, total_spend=_HEALTHY_TOTAL + _CALL_COST)
|
||||
assert healthy_row[0] == pytest.approx(_HEALTHY_SPEND + _CALL_COST)
|
||||
|
||||
|
||||
def test_team_info_and_team_list_report_a_legacy_total_no_lower_than_the_cycle_spend(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
healthy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_seed_membership(healthy, team_id, spend=_HEALTHY_SPEND, total_spend=_HEALTHY_TOTAL)
|
||||
|
||||
info: Final = _reported_memberships(gateway, "/team/info", {"team_id": team_id})
|
||||
listed: Final = _reported_memberships(gateway, "/team/list", {"team_id": team_id})
|
||||
|
||||
for reported in (info, listed):
|
||||
assert reported[legacy] == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_SPEND))
|
||||
assert reported[healthy] == (pytest.approx(_HEALTHY_SPEND), pytest.approx(_HEALTHY_TOTAL))
|
||||
assert _membership(legacy, team_id) == (pytest.approx(_LEGACY_SPEND), pytest.approx(_LEGACY_TOTAL))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reset_to", [0.0, 100.0], ids=["to_zero", "to_nonzero"])
|
||||
def test_an_admin_reset_keeps_the_legacy_cycle_spend_as_the_lifetime_floor(gateway: Gateway, reset_to: float) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
|
||||
response: Final = _reset_spend(gateway, team_id, legacy, reset_to)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert _membership(legacy, team_id) == (pytest.approx(reset_to), pytest.approx(_LEGACY_SPEND))
|
||||
|
||||
|
||||
def test_concurrent_admin_resets_leave_one_winner_and_a_lifetime_floor_at_the_old_cycle_spend(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
"""Two admins reset the same member at once. Each UPDATE re-checks reset_to <= spend, so
|
||||
the higher reset either lands first or is refused with 409 once the lower one has landed.
|
||||
Whatever the order, the row ends at the lower value with the old cycle spend preserved."""
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
higher, lower = pool.map(lambda value: _reset_spend(gateway, team_id, legacy, value), (600.0, 500.0))
|
||||
|
||||
assert lower.status_code == 200, lower.text
|
||||
assert higher.status_code in (200, 409), higher.text
|
||||
assert _membership(legacy, team_id) == (pytest.approx(500.0), pytest.approx(_LEGACY_SPEND))
|
||||
|
||||
|
||||
def test_an_admin_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, key)
|
||||
reset: Final = _reset_spend(gateway, team_id, legacy, 0.0)
|
||||
assert reset.status_code == 200, reset.text
|
||||
|
||||
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
|
||||
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
|
||||
assert total >= spend
|
||||
|
||||
|
||||
def _expire_member_budget(user_id: str, team_id: str, *, budget_duration: str) -> None:
|
||||
write_rows(
|
||||
"UPDATE \"LiteLLM_BudgetTable\" SET budget_duration = %s, budget_reset_at = NOW() - INTERVAL '1 hour' "
|
||||
'WHERE budget_id = (SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s)',
|
||||
(budget_duration, user_id, team_id),
|
||||
)
|
||||
|
||||
|
||||
async def _run_scheduled_reset() -> None:
|
||||
client: Final = PrismaClient(
|
||||
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache())
|
||||
)
|
||||
await client.connect()
|
||||
try:
|
||||
await ResetBudgetJob(
|
||||
ProxyLogging(user_api_key_cache=DualCache()), client
|
||||
).reset_budget_for_litellm_budget_table()
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spend, max_budget, carried",
|
||||
[(8.0, 10.0, 0.0), (10.0, 10.0, 0.0), (15.0, 10.0, 5.0)],
|
||||
ids=["below_cap", "at_cap", "above_cap_carries_overage"],
|
||||
)
|
||||
def test_a_scheduled_reset_lifts_the_legacy_total_before_zeroing_or_carrying_the_cycle_spend(
|
||||
gateway: Gateway, monkeypatch: pytest.MonkeyPatch, spend: float, max_budget: float, carried: float
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "budget_rollover", True)
|
||||
with gateway.scenario() as scenario:
|
||||
team_id: Final = scenario.team()
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=max_budget)
|
||||
_seed_membership(legacy, team_id, spend=spend, total_spend=3.0)
|
||||
_expire_member_budget(legacy, team_id, budget_duration="1d")
|
||||
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
|
||||
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
assert _membership(legacy, team_id) == (pytest.approx(carried), pytest.approx(spend))
|
||||
|
||||
|
||||
def test_a_scheduled_reset_racing_a_flush_ends_with_the_lifetime_equal_to_old_cycle_plus_cost(
|
||||
gateway: Gateway, provider: SharedProvider
|
||||
) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario)
|
||||
team_id: Final = scenario.team(models=[model])
|
||||
legacy: Final = _member_with_budget(gateway, scenario, team_id, max_budget=10_000)
|
||||
_seed_membership(legacy, team_id, spend=_LEGACY_SPEND, total_spend=_LEGACY_TOTAL)
|
||||
_expire_member_budget(legacy, team_id, budget_duration="1d")
|
||||
key: Final = scenario.key(user_id=legacy, team_id=team_id, models=[model])
|
||||
|
||||
_served(gateway, provider, model, key)
|
||||
asyncio.run(_run_scheduled_reset())
|
||||
|
||||
spend, total = _membership_reaches(legacy, team_id, total_spend=_LEGACY_SPEND + _CALL_COST)
|
||||
assert spend in (pytest.approx(0.0), pytest.approx(_CALL_COST)), (spend, total)
|
||||
assert total >= spend
|
||||
|
|
@ -113,6 +113,9 @@ class MockBatcher:
|
|||
self.litellm_projecttable = _Table("project", self)
|
||||
self.litellm_endusertable = _Table("enduser", self)
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None:
|
||||
self.calls.append({"table": "raw", "op": "execute_raw", "query": query, "args": args})
|
||||
|
||||
async def commit(self):
|
||||
self.committed = True
|
||||
return self.calls
|
||||
|
|
@ -2112,6 +2115,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo
|
|||
assert {(call["table"], call["op"]) for call in enduser_page.calls} == {("enduser", "update_many")}
|
||||
assert batcher.committed is True
|
||||
assert {(call["table"], call["op"]) for call in batcher.calls} == {
|
||||
("raw", "execute_raw"),
|
||||
("team_membership", "update_many"),
|
||||
("key", "update_many"),
|
||||
("org", "update_many"),
|
||||
|
|
@ -3393,6 +3397,35 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled(
|
|||
counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1")
|
||||
|
||||
|
||||
def test_budget_cascade_floors_membership_total_spend_before_zeroing_spend(
|
||||
rollover_enabled: None,
|
||||
reset_budget_job: ResetBudgetJob,
|
||||
mock_prisma_client: MockPrismaClient,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Membership rows written before total_spend existed hold their history only in
|
||||
spend. The cascade lifts total_spend to spend for the expiring tier's rows, in the
|
||||
same batch and ahead of the zero/decrement, so the reset cannot erase it."""
|
||||
_make_counter_invalidation_job(monkeypatch)
|
||||
budget: Final = _budget_row(budget_id="budget-legacy", budget_duration="7d", max_budget=10.0)
|
||||
mock_prisma_client.data["budget"] = [budget]
|
||||
membership: Final = type(
|
||||
"Membership",
|
||||
(),
|
||||
{"user_id": "member-1", "team_id": "team-1", "spend": 15.0, "budget_id": "budget-legacy"},
|
||||
)
|
||||
mock_prisma_client.db.litellm_teammembership.set_find_many_results([membership])
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
|
||||
|
||||
writes: Final = mock_prisma_client.db.batch_calls
|
||||
repair_index: Final = next(i for i, call in enumerate(writes) if call["op"] == "execute_raw")
|
||||
first_membership_reset: Final = next(i for i, call in enumerate(writes) if call["table"] == "team_membership")
|
||||
assert repair_index < first_membership_reset, writes
|
||||
assert writes[repair_index]["args"] == (["budget-legacy"],)
|
||||
assert mock_prisma_client.db.batchers[-1].committed
|
||||
|
||||
|
||||
def test_budget_cascade_carries_enduser_overage_when_rollover_enabled(
|
||||
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
|
||||
):
|
||||
|
|
|
|||
|
|
@ -1128,7 +1128,10 @@ async def test_commit_spend_updates_to_db_writes_team_member_spend_in_one_roster
|
|||
assert "members_with_roles @> jsonb_build_array(jsonb_build_object('user_id', member.user_id))" in statement
|
||||
assert "ON CONFLICT (user_id, team_id) DO UPDATE" in statement
|
||||
assert 'spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend' in statement
|
||||
assert 'total_spend = "LiteLLM_TeamMembership".total_spend + EXCLUDED.total_spend' in statement
|
||||
assert (
|
||||
'total_spend = GREATEST("LiteLLM_TeamMembership".total_spend, "LiteLLM_TeamMembership".spend)'
|
||||
" + EXCLUDED.total_spend"
|
||||
) in statement
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -12561,6 +12561,36 @@ async def test_team_member_me_returns_caller_membership(mock_db_client):
|
|||
assert call_kwargs["team_id"] == team_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"spend, stored_total, reported_total",
|
||||
[(612.4, 148.1, 612.4), (12.25, 97.8, 97.8)],
|
||||
ids=["legacy_row_floors_to_cycle_spend", "healthy_row_keeps_its_total"],
|
||||
)
|
||||
async def test_team_member_me_reports_total_spend_no_lower_than_cycle_spend(
|
||||
mock_db_client, spend, stored_total, reported_total
|
||||
):
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.management_endpoints.team_endpoints import team_member_me
|
||||
|
||||
team_id = "team-me-floor"
|
||||
caller_id = "alice@example.com"
|
||||
team = _build_team_for_me(team_id, [{"user_id": caller_id, "user_email": None, "role": "user"}])
|
||||
membership = _build_membership_for_me(caller_id, team_id, spend=spend).model_copy(
|
||||
update={"total_spend": stored_total}
|
||||
)
|
||||
p_team, p_membership, p_user = _patch_member_me_helpers(team=team, membership=membership, user=None)
|
||||
with p_team, p_membership, p_user:
|
||||
response = await team_member_me(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id=team_id,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id=caller_id),
|
||||
)
|
||||
|
||||
assert (response.spend, response.total_spend) == (spend, reported_total)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_me_matches_email_only_member(mock_db_client):
|
||||
"""
|
||||
|
|
@ -16608,7 +16638,7 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
|
|||
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b1", max_budget=50.0),
|
||||
)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
|
||||
|
|
@ -16631,15 +16661,76 @@ async def test_reset_team_member_spend_fn_success(monkeypatch):
|
|||
assert response["spend"] == 0.0
|
||||
assert response["previous_spend"] == 10.0
|
||||
assert response["max_budget"] == 50.0
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
|
||||
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
|
||||
data={"spend": 0.0},
|
||||
)
|
||||
statement, reset_user_id, reset_team_id, reset_to = mock_prisma_client.db.execute_raw.await_args.args
|
||||
assert (reset_user_id, reset_team_id, reset_to) == ("member-1", "team-1", 0.0)
|
||||
assert "total_spend = GREATEST(total_spend, spend)" in statement
|
||||
assert "spend >= $3" in statement
|
||||
assert await real_cache.async_get_cache(key="team-1_member-1") is None
|
||||
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
|
||||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_spend_fn_stale_spend_returns_409_and_keeps_caches(monkeypatch):
|
||||
"""The reset validates reset_to against a spend it read a moment ago. If the row moved
|
||||
underneath it (a flush landed, another admin reset first) the guarded UPDATE matches no
|
||||
row, and the endpoint must say so instead of reporting a reset that never happened."""
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
real_cache = UserApiKeyCache()
|
||||
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="cached-membership")
|
||||
real_spend_counter_cache = DualCache()
|
||||
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:team_member:member-1:team-1", value=999.0)
|
||||
|
||||
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache)
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1")),
|
||||
), pytest.raises(HTTPException) as exc_info:
|
||||
await reset_team_member_spend_fn(
|
||||
team_id="team-1",
|
||||
user_id="member-1",
|
||||
data=ResetSpendRequest(reset_to=5.0),
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user"
|
||||
),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") == "cached-membership"
|
||||
assert real_spend_counter_cache.in_memory_cache.get_cache(key="spend:team_member:member-1:team-1") == 999.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_team_memberships_floors_total_spend_at_current_cycle_spend():
|
||||
"""Rows written before total_spend existed carry their history only in spend. /team/info
|
||||
and /team/list both read through this helper, so a legacy row reports its lifetime spend
|
||||
as at least the current cycle, while a healthy row keeps its own total."""
|
||||
from litellm.proxy.management_endpoints.team_endpoints import get_all_team_memberships
|
||||
|
||||
legacy_row = MagicMock()
|
||||
legacy_row.model_dump = lambda: {"user_id": "legacy", "team_id": "team-1", "spend": 612.4, "total_spend": 148.1}
|
||||
healthy_row = MagicMock()
|
||||
healthy_row.model_dump = lambda: {"user_id": "healthy", "team_id": "team-1", "spend": 12.25, "total_spend": 97.8}
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[legacy_row, healthy_row])
|
||||
|
||||
legacy, healthy = await get_all_team_memberships(mock_prisma_client, ["team-1"])
|
||||
|
||||
assert (legacy.spend, legacy.total_spend) == (612.4, 612.4)
|
||||
assert (healthy.spend, healthy.total_spend) == (12.25, 97.8)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_spend_fn_membership_not_found(monkeypatch):
|
||||
mock_prisma_client = MagicMock()
|
||||
|
|
@ -16755,7 +16846,7 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey
|
|||
|
||||
membership_row = LiteLLM_TeamMembership(user_id="admin-user", team_id="team-1", spend=10.0)
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.execute_raw = AsyncMock(return_value=1)
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
|
|
@ -17994,6 +18085,7 @@ async def test_update_team_team_admin_resending_stored_member_defaults_is_not_an
|
|||
disable_audit_logging_for_mocked_team,
|
||||
):
|
||||
import contextlib
|
||||
|
||||
from litellm.models.budget import LiteLLM_BudgetTable
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ import pytest
|
|||
|
||||
from litellm.repositories.unit_of_work import (
|
||||
LinkedSpendResetWrites,
|
||||
SpendResetWrites,
|
||||
budget_cascade_unit_of_work,
|
||||
enduser_spend_reset_unit_of_work,
|
||||
spend_reset_unit_of_work,
|
||||
|
|
@ -39,6 +40,9 @@ class FakeBatch:
|
|||
self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls)
|
||||
self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
|
||||
|
||||
def execute_raw(self, query: str, *args: object) -> None:
|
||||
self.calls.append(("execute_raw", {"query": query}, {"args": args}))
|
||||
|
||||
async def commit(self) -> None:
|
||||
self.commit_count += 1
|
||||
|
||||
|
|
@ -166,7 +170,7 @@ async def test_every_cascade_dependent_writes_to_its_own_table_on_the_one_batch(
|
|||
async with budget_cascade_unit_of_work(_new_batch) as uow:
|
||||
writes = [getattr(uow, field.name) for field in fields(uow)]
|
||||
for write in writes:
|
||||
if isinstance(write, LinkedSpendResetWrites):
|
||||
if isinstance(write, SpendResetWrites):
|
||||
write.queue_spend_zero(where={"budget_id": "budget-1"})
|
||||
else:
|
||||
write.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
|
||||
|
|
@ -192,3 +196,18 @@ async def test_budget_cascade_raising_inside_block_skips_commit():
|
|||
await _blow_up_mid_transaction()
|
||||
|
||||
assert batch.commit_count == 0
|
||||
|
||||
|
||||
async def test_membership_total_spend_repair_is_queued_on_the_cascade_batch_ahead_of_the_reset():
|
||||
batch = FakeBatch()
|
||||
|
||||
async with budget_cascade_unit_of_work(lambda: batch) as uow:
|
||||
uow.team_memberships.queue_total_spend_repair(["budget-1", "budget-2"])
|
||||
uow.team_memberships.queue_spend_zero(where={"budget_id": "budget-1"})
|
||||
assert batch.commit_count == 0
|
||||
|
||||
assert batch.commit_count == 1
|
||||
repair, reset = batch.calls
|
||||
assert repair[0] == "execute_raw"
|
||||
assert repair[2] == {"args": (["budget-1", "budget-2"],)}
|
||||
assert reset == ("litellm_teammembership.update_many", {"budget_id": "budget-1"}, {"spend": 0})
|
||||
|
|
|
|||
|
|
@ -45,5 +45,30 @@ describe("MyUserTab", () => {
|
|||
expect(screen.getByText("member@example.com")).toBeInTheDocument();
|
||||
expect(screen.getByText("model-one")).toBeInTheDocument();
|
||||
expect(screen.getByText("TPM: 1,000")).toBeInTheDocument();
|
||||
expect(screen.getByText("Never resets")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows the next reset time when the member budget has a reset window", () => {
|
||||
vi.mocked(useMyTeamMember).mockReturnValue({
|
||||
data: {
|
||||
user_id: "user-1",
|
||||
user_email: "member@example.com",
|
||||
team_id: "team-1",
|
||||
role: "user",
|
||||
spend: 12.5,
|
||||
total_spend: 30,
|
||||
litellm_budget_table: {
|
||||
max_budget: 100,
|
||||
budget_duration: "30d",
|
||||
budget_reset_at: new Date(2026, 6, 15, 12, 0, 0).toISOString(),
|
||||
},
|
||||
},
|
||||
isLoading: false,
|
||||
error: null,
|
||||
} as ReturnType<typeof useMyTeamMember>);
|
||||
|
||||
renderWithProviders(<MyUserTab teamId="team-1" />);
|
||||
|
||||
expect(screen.getByText("Resets Jul 15, 2026 12:00")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import { formatBudgetReset } from "@/utils/budgetUtils";
|
||||
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
|
|
@ -67,7 +67,7 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
const totalSpend = data.total_spend ?? 0;
|
||||
const tpmLimit = budgetTable?.tpm_limit ?? null;
|
||||
const rpmLimit = budgetTable?.rpm_limit ?? null;
|
||||
const budgetReset = formatBudgetReset(budgetTable?.budget_reset_at);
|
||||
const budgetReset = describeBudgetResetSchedule(budgetTable?.budget_duration, budgetTable?.budget_reset_at);
|
||||
const allowedModels = budgetTable?.allowed_models ?? null;
|
||||
|
||||
return (
|
||||
|
|
@ -93,17 +93,16 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
<div className="grid grid-cols-1 gap-4 md:grid-cols-2">
|
||||
<Card>
|
||||
<CardContent>
|
||||
{labelWithTooltip(
|
||||
"Current Cycle Spend (USD)",
|
||||
"Spend for the current budget cycle. Resets to $0 when the budget window rolls over.",
|
||||
)}
|
||||
{labelWithTooltip("Current Cycle Spend (USD)", budgetReset.tooltip)}
|
||||
<div className="mt-2">
|
||||
<h3 className="text-2xl font-semibold">${formatNumber(spend, 4)}</h3>
|
||||
<span className="text-muted-foreground">
|
||||
of {maxBudget === null ? "Unlimited" : `$${formatNumber(maxBudget, 4)}`}
|
||||
</span>
|
||||
</div>
|
||||
{budgetReset && <div className="mt-1 text-muted-foreground">Resets {budgetReset}</div>}
|
||||
<div className="mt-1 text-muted-foreground">
|
||||
{budgetReset.resetAt ? `Resets ${budgetReset.label}` : budgetReset.label}
|
||||
</div>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
|
|
@ -120,7 +119,10 @@ export default function MyUserTab({ teamId }: MyUserTabProps) {
|
|||
|
||||
<Card>
|
||||
<CardContent>
|
||||
{labelWithTooltip("Total Spend (USD)", "Cumulative spend across all budget cycles within this team.")}
|
||||
{labelWithTooltip(
|
||||
"Total Spend (USD)",
|
||||
"Cumulative spend across budget cycles within this team. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.",
|
||||
)}
|
||||
<h4 className="mt-2 text-xl font-semibold">${formatNumber(totalSpend, 4)}</h4>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
|
|
|||
|
|
@ -313,7 +313,47 @@ describe("TeamMembersComponent", () => {
|
|||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument();
|
||||
expect(screen.getByText("Jul 15, 2026 12:00")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows Never resets for a member whose budget has no reset window", () => {
|
||||
const teamData = createMockTeamData();
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
|
||||
renderWithProviders(
|
||||
<TeamMembersComponent
|
||||
teamData={teamData}
|
||||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.queryByText("Jul 15, 2026 12:00")).not.toBeInTheDocument();
|
||||
expect(screen.getAllByText("Never resets").length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("shows Unknown for a member whose budget has a duration but no scheduled reset yet", () => {
|
||||
const teamData = createMockTeamData();
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_duration = "30d";
|
||||
teamData.team_memberships[0].litellm_budget_table!.budget_reset_at = null;
|
||||
renderWithProviders(
|
||||
<TeamMembersComponent
|
||||
teamData={teamData}
|
||||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Unknown")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("shows each member the values from its own membership row", () => {
|
||||
|
|
|
|||
|
|
@ -10,8 +10,9 @@ import MemberTable from "@/components/common_components/MemberTable";
|
|||
import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload";
|
||||
import { Member } from "@/components/networking";
|
||||
import { parseErrorMessage } from "@/components/shared/errorUtils";
|
||||
import { DateCell, MoneyCell } from "@/components/shared/table_cells";
|
||||
import { MoneyCell } from "@/components/shared/table_cells";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { describeBudgetResetSchedule } from "@/utils/budgetUtils";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles";
|
||||
import { CircleHelp } from "lucide-react";
|
||||
|
|
@ -152,6 +153,16 @@ export default function TeamMemberTab({
|
|||
return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null;
|
||||
};
|
||||
|
||||
const renderBudgetReset = (userId: string | null) => {
|
||||
const budget = findMembership(userId)?.litellm_budget_table;
|
||||
const schedule = describeBudgetResetSchedule(budget?.budget_duration ?? null, budget?.budget_reset_at ?? null);
|
||||
return (
|
||||
<SimpleTooltip content={schedule.tooltip}>
|
||||
<span>{schedule.label}</span>
|
||||
</SimpleTooltip>
|
||||
);
|
||||
};
|
||||
|
||||
const extraColumns: NonNullable<ComponentProps<typeof MemberTable>["extraColumns"]> = [
|
||||
{
|
||||
title: (
|
||||
|
|
@ -190,7 +201,7 @@ export default function TeamMemberTab({
|
|||
title: (
|
||||
<span className="flex items-center gap-1">
|
||||
Current Cycle Spend (USD)
|
||||
<SimpleTooltip content="Spend for the current budget cycle. Resets to $0 when the member's budget window rolls over. This is the value checked against the member's budget.">
|
||||
<SimpleTooltip content="Spend in this member's current budget cycle, the value checked against their budget. The Budget Reset column shows when this cycle ends, or Never resets. Spend above a rollover cap at reset carries into the next cycle.">
|
||||
<CircleHelp className="size-4" aria-label="Current cycle spend information" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
|
|
@ -204,7 +215,7 @@ export default function TeamMemberTab({
|
|||
title: (
|
||||
<span className="flex items-center gap-1">
|
||||
Total Spend (USD)
|
||||
<SimpleTooltip content="Cumulative spend by this member within this team, across all budget cycles. Tracking began 2026-04-21; spend from before that date is not included.">
|
||||
<SimpleTooltip content="Cumulative spend by this member within this team, across budget cycles. This is a lower bound: lifetime tracking began 2026-04-21, so cycles completed before then are not included. It is never less than Current Cycle Spend.">
|
||||
<CircleHelp className="size-4" aria-label="Total spend information" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
|
|
@ -255,7 +266,7 @@ export default function TeamMemberTab({
|
|||
title: "Budget Reset",
|
||||
key: "budget_reset",
|
||||
sortValue: (record: Member) => getUserBudgetReset(record.user_id),
|
||||
render: (record: Member) => <DateCell value={getUserBudgetReset(record.user_id)} precision="date" />,
|
||||
render: (record: Member) => renderBudgetReset(record.user_id),
|
||||
},
|
||||
{
|
||||
title: (
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import { afterEach, describe, expect, it } from "vitest";
|
||||
|
||||
import { formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
|
||||
import { describeBudgetResetSchedule, formatBudgetReset, formatBudgetResetUtc } from "./budgetUtils";
|
||||
|
||||
const originalTimezone: string | undefined = process.env.TZ;
|
||||
|
||||
|
|
@ -21,3 +21,26 @@ describe("formatBudgetResetUtc", () => {
|
|||
expect(formatBudgetReset(resetAt)).toBe("Oct 19, 2026");
|
||||
});
|
||||
});
|
||||
|
||||
describe("describeBudgetResetSchedule", () => {
|
||||
it("reports Never resets when the budget has no duration and no scheduled reset", () => {
|
||||
const schedule = describeBudgetResetSchedule(null, null);
|
||||
expect(schedule.label).toBe("Never resets");
|
||||
expect(schedule.resetAt).toBeNull();
|
||||
});
|
||||
|
||||
it("reports Unknown when a duration is set but no reset has been scheduled yet", () => {
|
||||
const schedule = describeBudgetResetSchedule("30d", null);
|
||||
expect(schedule.label).toBe("Unknown");
|
||||
expect(schedule.tooltip).toContain("every 30d");
|
||||
});
|
||||
|
||||
it("shows the next reset with its time zone and the rollover carryover note", () => {
|
||||
const resetAt = new Date(2026, 6, 15, 12, 0, 0);
|
||||
const schedule = describeBudgetResetSchedule("30d", resetAt.toISOString());
|
||||
expect(schedule.label).toBe("Jul 15, 2026 12:00");
|
||||
expect(schedule.resetAt).toBe(resetAt.toISOString());
|
||||
expect(schedule.tooltip).toContain(Intl.DateTimeFormat().resolvedOptions().timeZone);
|
||||
expect(schedule.tooltip).toContain("carries into the next cycle");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -16,3 +16,41 @@ export function formatBudgetResetUtc(iso: string | null | undefined): string | n
|
|||
if (!resetDate.isValid()) return null;
|
||||
return resetDate.format("MMM D, YYYY");
|
||||
}
|
||||
|
||||
export interface BudgetResetSchedule {
|
||||
label: string;
|
||||
resetAt: string | null;
|
||||
tooltip: string;
|
||||
}
|
||||
|
||||
const ROLLOVER_NOTE = "Spend above a rollover cap at reset carries into the next cycle.";
|
||||
|
||||
export function describeBudgetResetSchedule(
|
||||
budgetDuration: string | null | undefined,
|
||||
budgetResetAt: string | null | undefined,
|
||||
): BudgetResetSchedule {
|
||||
const resetDate = budgetResetAt ? dayjs(budgetResetAt) : null;
|
||||
if (resetDate?.isValid()) {
|
||||
const cadence = budgetDuration ? `every ${budgetDuration}` : "on a schedule";
|
||||
const timeZone = Intl.DateTimeFormat().resolvedOptions().timeZone;
|
||||
const label = resetDate.format("MMM D, YYYY HH:mm");
|
||||
return {
|
||||
label,
|
||||
resetAt: resetDate.toISOString(),
|
||||
tooltip: `Current Cycle Spend resets ${cadence}. Next reset ${label} (${timeZone}). ${ROLLOVER_NOTE}`,
|
||||
};
|
||||
}
|
||||
if (!budgetDuration) {
|
||||
return {
|
||||
label: "Never resets",
|
||||
resetAt: null,
|
||||
tooltip:
|
||||
"This budget has no reset window, so Current Cycle Spend accumulates for as long as the membership exists.",
|
||||
};
|
||||
}
|
||||
return {
|
||||
label: "Unknown",
|
||||
resetAt: null,
|
||||
tooltip: `Resets every ${budgetDuration}, but the next reset time has not been scheduled yet. ${ROLLOVER_NOTE}`,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue