From 6eda17aadbc144b6bf878b0e4e7313c80a034c59 Mon Sep 17 00:00:00 2001 From: ishaan-j-berri Date: Fri, 13 Nov 2026 23:33:24 +0800 Subject: [PATCH] fix(proxy): enforce settled customer budgets across workers (#51568) * fix(proxy): enforce settled customer budgets across workers Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(proxy): sync customer budget reset API types Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> * fix(proxy): reseed customer spend through shared counters Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> --- litellm/proxy/auth/customer_hierarchy_auth.py | 31 +++- litellm/proxy/auth/customer_spend.py | 53 ++++++ litellm/proxy/db/spend_counter_reseed.py | 38 +++++ .../proxy/hooks/proxy_track_cost_callback.py | 8 + litellm/proxy/litellm_pre_call_utils.py | 1 + .../redis-store/src/spend_counters/keys.rs | 23 +++ .../redis-store/src/spend_counters/mod.rs | 25 ++- .../types/proxy/auth/customer_hierarchy.py | 2 + litellm/types/utils.py | 1 + .../src/lib/http/schema.d.ts | 2 + .../auth/test_customer_hierarchy_auth.py | 117 +++++++++++++ tests/unit/proxy/auth/test_customer_spend.py | 157 ++++++++++++++++++ .../proxy/db/test_spend_counter_reseed.py | 42 +++++ .../unit/redis/stores/test_spend_counters.py | 6 + 14 files changed, 500 insertions(+), 6 deletions(-) create mode 100644 litellm/proxy/auth/customer_spend.py create mode 100644 tests/unit/proxy/auth/test_customer_spend.py diff --git a/litellm/proxy/auth/customer_hierarchy_auth.py b/litellm/proxy/auth/customer_hierarchy_auth.py index b9d12b00c1a..6065fdd636e 100644 --- a/litellm/proxy/auth/customer_hierarchy_auth.py +++ b/litellm/proxy/auth/customer_hierarchy_auth.py @@ -22,6 +22,7 @@ from litellm.proxy._types import ProxyErrorTypes, ProxyException from litellm.proxy.auth.auth_checks import can_object_call_model from litellm.proxy.auth.auth_object_prefetch import AUTH_OBJECTS_TARGET from litellm.proxy.auth.customer_member_auto_create import ensure_customer_member +from litellm.proxy.auth.customer_spend import current_customer_spend from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl @@ -61,6 +62,7 @@ class _OrganizationRow(BaseModel): customer_org_id: str spend: float = 0.0 + budget_reset_at: datetime | None = None budget: Mapping[str, object] | None = None @@ -69,6 +71,7 @@ class _TeamRow(BaseModel): customer_team_id: str spend: float = 0.0 + budget_reset_at: datetime | None = None budget: Mapping[str, object] | None = None parent_org: _OrganizationRow @@ -78,6 +81,7 @@ class _MemberRow(BaseModel): customer_member_id: str spend: float = 0.0 + budget_reset_at: datetime | None = None budget: Mapping[str, object] | None = None budget_override: Mapping[str, object] | None = None budget_override_expires_at: datetime | None = None @@ -129,13 +133,16 @@ def customer_identity_from_headers(headers: Mapping[str, object] | None) -> Cust return identity -def _scope(key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None) -> CustomerScopeAuth: +def _scope( + key: CustomerScopeKey, entity_id: str, spend: float, budget: Budget | None, budget_reset_at: datetime | None = None +) -> CustomerScopeAuth: if budget is None: - return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend) + return CustomerScopeAuth(key=key, entity_id=entity_id, spend=spend, budget_reset_at=budget_reset_at) return CustomerScopeAuth( key=key, entity_id=entity_id, spend=spend, + budget_reset_at=budget_reset_at, max_budget=budget.max_budget, tpm_limit=budget.tpm_limit, rpm_limit=budget.rpm_limit, @@ -163,13 +170,21 @@ def customer_hierarchy_from_member_row(row: Mapping[str, object], now: datetime) organization.customer_org_id, organization.spend, related_budget(organization.budget, now), + organization.budget_reset_at, + ), + _scope( + "customer_team", + team.customer_team_id, + team.spend, + related_budget(team.budget, now), + team.budget_reset_at, ), - _scope("customer_team", team.customer_team_id, team.spend, related_budget(team.budget, now)), _scope( "customer_member", member.customer_member_id, member.spend, effective_member_budget(template, override, member.budget_override_expires_at, now), + member.budget_reset_at, ), ), ) @@ -186,6 +201,7 @@ def customer_hierarchy_from_organization_row(row: Mapping[str, object], now: dat organization.customer_org_id, organization.spend, related_budget(organization.budget, now), + organization.budget_reset_at, ), ), ) @@ -300,7 +316,9 @@ def _not_found(message: str, param: str) -> ProxyException: async def _organization_named_by_alias( - customer_org_alias: str, prisma_client: "PrismaClient | None", user_api_key_cache: UserApiKeyCache + customer_org_alias: str, + prisma_client: "PrismaClient | None", + user_api_key_cache: UserApiKeyCache, ) -> CustomerHierarchyAuth: organization: Final = ( await get_customer_organization_hierarchy(customer_org_alias, prisma_client, user_api_key_cache) @@ -411,7 +429,10 @@ async def resolve_customer_hierarchy( ) -> CustomerHierarchyAuth: """Load the hierarchy the headers name and admit the request against it. ``request_model`` is the model the route will actually serve, which may differ from the body.""" - hierarchy: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register) + from litellm.proxy.proxy_server import get_current_spend + + cached: Final = await load_customer_hierarchy(identity, prisma_client, user_api_key_cache, register) + hierarchy: Final = await current_customer_spend(cached, get_current_spend) check_customer_hierarchy_key_scope(hierarchy, key_metadata or {}) if not RouteChecks.is_info_route(route): check_customer_hierarchy_budget(hierarchy) diff --git a/litellm/proxy/auth/customer_spend.py b/litellm/proxy/auth/customer_spend.py new file mode 100644 index 00000000000..0c28a9412e5 --- /dev/null +++ b/litellm/proxy/auth/customer_spend.py @@ -0,0 +1,53 @@ +import asyncio +from collections.abc import Awaitable, Callable, Mapping +from typing import Final + +from pydantic import ValidationError + +from litellm.redis.stores.spend_counters import SpendCounterStore +from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth + + +def customer_spend_counter(scope: CustomerScopeAuth) -> str: + period: Final = scope.budget_reset_at.isoformat() if scope.budget_reset_at is not None else "lifetime" + return f"spend:{scope.key}:{scope.entity_id}:{period}" + + +async def _current_scope( + scope: CustomerScopeAuth, current: float | None, store: SpendCounterStore +) -> CustomerScopeAuth: + spend: Final = ( + await store.seed(customer_spend_counter(scope), scope.spend, increment_if_lost=False) + if current is None + else current + ) + return scope.model_copy(update={"spend": max(scope.spend, spend)}) + + +async def current_customer_spend( + hierarchy: CustomerHierarchyAuth, + read_spend: Callable[[str, float, float | None], Awaitable[float]], +) -> CustomerHierarchyAuth: + readings: Final = await asyncio.gather( + *(read_spend(customer_spend_counter(scope), scope.spend, scope.max_budget) for scope in hierarchy.scopes) + ) + scopes: Final = tuple( + scope.model_copy(update={"spend": max(scope.spend, spend)}) + for scope, spend in zip(hierarchy.scopes, readings, strict=True) + ) + return hierarchy.model_copy(update={"scopes": tuple(scopes)}) + + +async def settle_customer_spend(metadata: Mapping[str, object], cost: float, store: SpendCounterStore) -> None: + snapshot: Final = metadata.get("user_api_key_customer_hierarchy") + if snapshot is None or cost == 0: + return + try: + hierarchy: Final = CustomerHierarchyAuth.model_validate(snapshot) + except ValidationError: + return + readings: Final = await asyncio.gather(*(store.read(customer_spend_counter(scope)) for scope in hierarchy.scopes)) + _ = await asyncio.gather( + *(_current_scope(scope, reading[0], store) for scope, reading in zip(hierarchy.scopes, readings, strict=True)) + ) + _ = await store.increment_many(tuple((customer_spend_counter(scope), cost) for scope in hierarchy.scopes)) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index ae45b43f15f..ef59e7d587f 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -13,6 +13,7 @@ memory in long-lived deployments. """ import asyncio +import re from collections import OrderedDict from collections.abc import Mapping from datetime import datetime, timezone @@ -34,6 +35,9 @@ from litellm.repositories.organization_repository import OrganizationRepository from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( BudgetWindowSpendRepository, + CustomerMemberRepository, + CustomerOrganizationRepository, + CustomerTeamRepository, EndUserRepository, SpendLogsRepository, TeamMembershipRepository, @@ -69,6 +73,23 @@ def _as_utc(value: datetime) -> datetime: return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc) +_CUSTOMER_COUNTER: Final = re.compile( + r"spend:(customer_org|customer_team|customer_member):(.+):" + r"(lifetime|\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d{1,6})?(?:Z|[+-]\d{2}:\d{2})?)" +) + + +def _customer_counter_period_matches(row: object, period: str) -> bool: + try: + expected: Final = None if period == "lifetime" else datetime.fromisoformat(period) + except ValueError: + return False + actual: Final = getattr(row, "budget_reset_at", None) + if expected is None or actual is None: + return expected is actual + return isinstance(actual, datetime) and _as_utc(actual) == _as_utc(expected) + + class SpendCounterReseed: """ Reseeds spend counters from the authoritative DB and warms the cache, @@ -147,11 +168,28 @@ class SpendCounterReseed: return None if row is None: return None + customer: Final = _CUSTOMER_COUNTER.fullmatch(counter_key) + if customer is not None and not _customer_counter_period_matches(row, customer.group(3)): + return None return float(getattr(row, "spend", 0.0) or 0.0) @staticmethod async def _counter_row(prisma_client: "PrismaClient", counter_key: str) -> object | None: async with db_lookup_gate.current(): + customer: Final = _CUSTOMER_COUNTER.fullmatch(counter_key) + if customer is not None: + family, entity_id, _ = customer.groups() + if family == "customer_org": + return await CustomerOrganizationRepository(prisma_client).table.find_unique( + where={"customer_org_id": entity_id} + ) + if family == "customer_team": + return await CustomerTeamRepository(prisma_client).table.find_unique( + where={"customer_team_id": entity_id} + ) + return await CustomerMemberRepository(prisma_client).table.find_unique( + where={"customer_member_id": entity_id} + ) if counter_key.startswith("spend:key:"): token: Final = counter_key[len("spend:key:") :] return await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": token}) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 554fc7ed1d5..23e596e41f1 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -36,6 +36,7 @@ from litellm.proxy.auth.auth_checks import ( get_key_object, get_team_object, ) +from litellm.proxy.auth.customer_spend import settle_customer_spend from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.db_lookup_gate import DBLookupDeadlineExceeded from litellm.proxy.db.db_spend_update_writer import ( @@ -1440,7 +1441,14 @@ async def _update_database_and_spend_counters_in_batch( await _release_budget_reservation(budget_reservation=budget_reservation) return False + from litellm.proxy.proxy_server import spend_counter_store + try: + await settle_customer_spend( + _OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=kwargs)), + response_cost, + spend_counter_store, + ) await increment_spend_counters( token=user_api_key, team_id=team_id, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 5da902ee4ec..6354c6c3628 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2140,6 +2140,7 @@ class LiteLLMProxyRequestSetup: user_api_key_customer_member_id=customer.customer_member_id if customer else None, user_api_key_customer_team_id=customer.customer_team_id if customer else None, user_api_key_customer_org_id=customer.customer_org_id if customer else None, + user_api_key_customer_hierarchy=customer.model_dump(mode="json") if customer else None, user_api_key_user_email=user_api_key_dict.user_email, user_api_key_request_route=user_api_key_dict.request_route, user_api_key_budget_reset_at=( diff --git a/litellm/rust/crates/redis-store/src/spend_counters/keys.rs b/litellm/rust/crates/redis-store/src/spend_counters/keys.rs index 5221f8ab66c..ff7269aee95 100644 --- a/litellm/rust/crates/redis-store/src/spend_counters/keys.rs +++ b/litellm/rust/crates/redis-store/src/spend_counters/keys.rs @@ -23,6 +23,9 @@ pub enum SpendCounter { Project(String), ModelAccessGroup(String), ShadowEval(String), + CustomerOrg(String), + CustomerTeam(String), + CustomerMember(String), } #[derive(Debug, thiserror::Error)] @@ -51,6 +54,9 @@ impl SpendCounter { Self::Project(id) => format!("spend:project:{id}"), Self::ModelAccessGroup(name) => format!("spend:model_access_group:{name}"), Self::ShadowEval(job_id) => format!("spend:shadow_eval:{job_id}"), + Self::CustomerOrg(id) => format!("spend:customer_org:{id}"), + Self::CustomerTeam(id) => format!("spend:customer_team:{id}"), + Self::CustomerMember(id) => format!("spend:customer_member:{id}"), } } @@ -95,7 +101,24 @@ impl SpendCounter { "project" => Self::Project(id.to_owned()), "model_access_group" => Self::ModelAccessGroup(id.to_owned()), "shadow_eval" => Self::ShadowEval(id.to_owned()), + "customer_org" => Self::CustomerOrg(id.to_owned()), + "customer_team" => Self::CustomerTeam(id.to_owned()), + "customer_member" => Self::CustomerMember(id.to_owned()), _ => return Err(invalid()), }) } } + +#[cfg(test)] +mod tests { + use super::SpendCounter; + use rstest::rstest; + + #[rstest] + #[case::organization("spend:customer_org:tenant:lifetime")] + #[case::team("spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00")] + #[case::member("spend:customer_member:user@example.com:lifetime")] + fn customer_counter_keys_round_trip(#[case] key: &str) { + assert_eq!(SpendCounter::parse(key).unwrap().key(), key); + } +} diff --git a/litellm/rust/crates/redis-store/src/spend_counters/mod.rs b/litellm/rust/crates/redis-store/src/spend_counters/mod.rs index 425f83ef85c..99ab9bea8f7 100644 --- a/litellm/rust/crates/redis-store/src/spend_counters/mod.rs +++ b/litellm/rust/crates/redis-store/src/spend_counters/mod.rs @@ -65,7 +65,10 @@ impl SpendCounters { fn ttl(&self, counter: &SpendCounter, default_ttl: Duration) -> Duration { match counter { - SpendCounter::Committed(_) => default_ttl.max(Duration::from_millis( + SpendCounter::Committed(_) + | SpendCounter::CustomerOrg(_) + | SpendCounter::CustomerTeam(_) + | SpendCounter::CustomerMember(_) => default_ttl.max(Duration::from_millis( self.committed_ttl_millis.load(Ordering::Relaxed), )), _ => default_ttl, @@ -418,3 +421,23 @@ impl SpendCounters { } } } + +#[cfg(test)] +mod tests { + use super::{SpendCounter, SpendCounters}; + use rstest::rstest; + use std::time::Duration; + + #[rstest] + #[case::organization(SpendCounter::CustomerOrg("tenant:lifetime".into()))] + #[case::team(SpendCounter::CustomerTeam("team:lifetime".into()))] + #[case::member(SpendCounter::CustomerMember("member:lifetime".into()))] + fn customer_counters_outlive_the_pending_database_flush(#[case] counter: SpendCounter) { + let store = SpendCounters::new(None); + store.set_committed_ttl(Duration::from_secs(120)); + assert_eq!( + store.ttl(&counter, Duration::from_secs(1)), + Duration::from_secs(120) + ); + } +} diff --git a/litellm/types/proxy/auth/customer_hierarchy.py b/litellm/types/proxy/auth/customer_hierarchy.py index a1141c9d113..09a4fcd2834 100644 --- a/litellm/types/proxy/auth/customer_hierarchy.py +++ b/litellm/types/proxy/auth/customer_hierarchy.py @@ -1,5 +1,6 @@ from collections.abc import Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType from typing import Final, Literal, TypeAlias @@ -26,6 +27,7 @@ class CustomerScopeAuth(BaseModel): key: CustomerScopeKey entity_id: str spend: float = 0.0 + budget_reset_at: datetime | None = None max_budget: float | None = None tpm_limit: int | None = None rpm_limit: int | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 7ae7bc92b01..d35c6f6379e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3148,6 +3148,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_customer_member_id: ReadOnly[str | None] user_api_key_customer_team_id: ReadOnly[str | None] user_api_key_customer_org_id: ReadOnly[str | None] + user_api_key_customer_hierarchy: NotRequired[ReadOnly[Mapping[str, object] | None]] user_api_key_request_route: str | None user_api_key_auth_metadata: dict[str, str] | None diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5b240a6fe3b..3245c4db55f 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -33236,6 +33236,8 @@ export interface components { CustomerScopeAuth: { /** Allowed Models */ allowed_models?: string[] | null; + /** Budget Reset At */ + budget_reset_at?: string | null; /** Entity Id */ entity_id: string; /** diff --git a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py index 9efb9167ecd..e950c48eb87 100644 --- a/tests/unit/proxy/auth/test_customer_hierarchy_auth.py +++ b/tests/unit/proxy/auth/test_customer_hierarchy_auth.py @@ -23,16 +23,126 @@ from litellm.proxy.auth.customer_hierarchy_auth import ( ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.management_helpers.customer_hierarchy import ErrorCode, Failure +from litellm.redis.stores.connection import connect +from litellm.redis.stores.spend_counters import SpendCounterStore from litellm.types.proxy.auth.customer_hierarchy import ( CustomerHierarchyAuth, CustomerRequestIdentity, CustomerScopeAuth, ) from litellm.types.proxy.management_endpoints.customer_hierarchy import CustomerMember +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters NOW: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) +@pytest.fixture(autouse=True) +def customer_spend_store(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "spend_counter_store", SpendCounterStore(connect(None).spend_counters())) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("alias_only", (False, True)) +async def test_a_missing_shared_counter_reseeds_spend_without_reloading_a_stale_cached_snapshot( + alias_only: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy import proxy_server + + cache: Final = UserApiKeyCache() + old: Final = _Prisma( + _MemberTable({"member-1": _row(org_budget={"max_budget": 1})}), + _OrganizationTable((_organization("tenant", max_budget=1),)), + ) + persisted: Final = _Prisma( + _MemberTable({"member-1": _row(org_budget={"max_budget": 1}, spend=(2, 0, 0))}), + _OrganizationTable((_organization("tenant", max_budget=1, spend=2),)), + ) + monkeypatch.setattr(proxy_server, "prisma_client", persisted) + if alias_only: + _ = await get_customer_organization_hierarchy( + "tenant", + old, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + ) + else: + _ = await get_customer_hierarchy( + "member-1", + old, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + ) + identity: Final = CustomerRequestIdentity( + customer_org_alias="tenant" if alias_only else None, + customer_member_id=None if alias_only else "member-1", + ) + with pytest.raises(litellm.BudgetExceededError) as rejected: + await resolve_customer_hierarchy( + identity, + persisted, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + None, + "model", + "/responses", + ) + assert rejected.value.current_cost == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("level", ("org", "team", "member")) +@pytest.mark.parametrize("alias_only", (False, True)) +async def test_settled_spend_reaches_another_workers_cached_hierarchy( + monkeypatch: pytest.MonkeyPatch, level: str, alias_only: bool +) -> None: + from litellm.proxy import proxy_server + + shared: Final[dict[str, object]] = {} + settling: Final = FakeSpendCounters(redis=shared).store() + reading: Final = FakeSpendCounters(redis=shared).store() + monkeypatch.setattr(proxy_server, "spend_counter_store", reading) + cache: Final = UserApiKeyCache() + row: Final = _row(template={"max_budget": 1.0}, team_budget={"max_budget": 1.0}, org_budget={"max_budget": 1.0}) + prisma: Final = _Prisma( + _MemberTable({"member-1": row}), _OrganizationTable((_organization("tenant", max_budget=1),)) + ) + identity: Final = CustomerRequestIdentity( + customer_org_alias="tenant" if alias_only else None, + customer_member_id=None if alias_only else "member-1", + ) + await resolve_customer_hierarchy( + identity, + prisma, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + None, + "model", + "/responses", + ) + counter: Final = f"spend:customer_{level}:{level}-1:lifetime" + await settling.increment(counter, 1.25) + if alias_only and level != "org": + allowed: Final = await resolve_customer_hierarchy( + identity, + prisma, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + None, + "model", + "/responses", + ) + assert allowed.customer_member_id is None + return + with pytest.raises(litellm.BudgetExceededError) as rejected: + await resolve_customer_hierarchy( + identity, + prisma, # pyright: ignore[reportArgumentType] # database boundary fixture + cache, + None, + "model", + "/responses", + ) + assert (rejected.value.entity_type, rejected.value.current_cost) == (f"customer_{level}", 1.25) + + def _row( *, template: Mapping[str, object] | None = None, @@ -187,6 +297,10 @@ class _Record: def model_dump(self) -> Mapping[str, object]: return self.values + @property + def spend(self): + return self.values.get("spend", 0.0) + class _OrganizationTable: def __init__(self, rows: Sequence[Mapping[str, object]] = ()) -> None: @@ -200,6 +314,9 @@ class _OrganizationTable: matching: Final = [_Record(row) for row in self.rows if row.get("alias") == where["alias"]] return matching[:take] if take is not None else matching + async def find_unique(self, where: Mapping[str, object]): + return next((_Record(row) for row in self.rows if row.get("customer_org_id") == where["customer_org_id"]), None) + class _Db: def __init__(self, table: _MemberTable, organizations: _OrganizationTable) -> None: diff --git a/tests/unit/proxy/auth/test_customer_spend.py b/tests/unit/proxy/auth/test_customer_spend.py new file mode 100644 index 00000000000..194e2761115 --- /dev/null +++ b/tests/unit/proxy/auth/test_customer_spend.py @@ -0,0 +1,157 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Final + +import pytest + +import litellm +from litellm.proxy.auth.customer_hierarchy_auth import check_customer_hierarchy_budget +from litellm.proxy.auth.customer_spend import current_customer_spend as read_customer_spend +from litellm.proxy.auth.customer_spend import customer_spend_counter, settle_customer_spend +from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters + + +async def current_customer_spend(hierarchy, store, prisma=None): + from litellm.proxy import proxy_server + + with pytest.MonkeyPatch.context() as context: + context.setattr(proxy_server, "spend_counter_store", store) + context.setattr(proxy_server, "prisma_client", prisma) + return await read_customer_spend(hierarchy, proxy_server.get_current_spend) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("scope", ("customer_org", "customer_team", "customer_member")) +async def test_a_lost_counter_reseeds_committed_spend_without_refreshing_the_customer_snapshot(scope: str) -> None: + from tests.unit.proxy.db.test_spend_counter_reseed import _FakeFindUniqueTable + + row: Final = SimpleNamespace(spend=2.5, budget_reset_at=None) + table: Final = _FakeFindUniqueTable(row) + table_name: Final = { + "customer_org": "litellm_customerorganizationtable", + "customer_team": "litellm_customerteamtable", + "customer_member": "litellm_customermembertable", + }[scope] + prisma: Final = SimpleNamespace(db=SimpleNamespace(**{table_name: table})) + snapshot: Final = CustomerHierarchyAuth( + customer_org_id="tenant", scopes=(CustomerScopeAuth(key=scope, entity_id="id:with:colons", max_budget=2),) + ) + shared: Final[dict[str, object]] = {} + first: Final = FakeSpendCounters(redis=shared).store() + second: Final = FakeSpendCounters(redis=shared).store() + for store in (first, second): + loaded: Final = await current_customer_spend(snapshot, store, prisma) + with pytest.raises(litellm.BudgetExceededError): + check_customer_hierarchy_budget(loaded) + assert snapshot.scopes[0].spend == 0 + assert len(table.where_clauses) == 1 + await first.delete((customer_spend_counter(snapshot.scopes[0]),)) + reseeded: Final = await current_customer_spend(snapshot, second, prisma) + with pytest.raises(litellm.BudgetExceededError): + check_customer_hierarchy_budget(reseeded) + assert reseeded.scopes[0].spend == row.spend + assert len(table.where_clauses) == 2 + + +@pytest.mark.asyncio +async def test_a_database_snapshot_ahead_of_settlement_does_not_count_the_pending_charge_twice() -> None: + store: Final = FakeSpendCounters().store() + initial: Final = CustomerHierarchyAuth( + customer_org_id="tenant", scopes=(CustomerScopeAuth(key="customer_org", entity_id="tenant"),) + ) + _ = await current_customer_spend(initial, store) + persisted: Final = initial.model_copy(update={"scopes": (initial.scopes[0].model_copy(update={"spend": 0.5}),)}) + _ = await current_customer_spend(persisted, store) + await settle_customer_spend({"user_api_key_customer_hierarchy": initial.model_dump(mode="json")}, 0.5, store) + assert (await current_customer_spend(persisted, store)).scopes[0].spend == 0.5 + + +@pytest.mark.asyncio +async def test_member_b_settlement_blocks_cached_member_a_and_an_alias_only_request_on_another_worker() -> None: + shared: Final[dict[str, object]] = {} + settling: Final = FakeSpendCounters(redis=shared).store() + reading: Final = FakeSpendCounters(redis=shared).store() + organization: Final = CustomerScopeAuth(key="customer_org", entity_id="tenant", max_budget=1) + member_a: Final = CustomerHierarchyAuth( + customer_org_id="tenant", + customer_member_id="a", + scopes=(organization, CustomerScopeAuth(key="customer_member", entity_id="a")), + ) + member_b: Final = CustomerHierarchyAuth( + customer_org_id="tenant", + customer_member_id="b", + scopes=(organization, CustomerScopeAuth(key="customer_member", entity_id="b")), + ) + alias_only: Final = CustomerHierarchyAuth(customer_org_id="tenant", scopes=(organization,)) + _ = await current_customer_spend(member_a, reading) + _ = await current_customer_spend(alias_only, reading) + await settle_customer_spend({"user_api_key_customer_hierarchy": member_b.model_dump(mode="json")}, 1.25, settling) + for cached in (member_a, alias_only): + current: Final = await current_customer_spend(cached, reading) + with pytest.raises(litellm.BudgetExceededError) as rejected: + check_customer_hierarchy_budget(current) + assert rejected.value.entity_type == "customer_org" + assert rejected.value.entity_id == "tenant" + assert (await current_customer_spend(member_a, reading)).scopes[1].spend == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("redis", (True, False)) +async def test_settlement_counts_once_on_each_scope_without_modifying_the_cached_snapshot(redis: bool) -> None: + handle: Final = FakeSpendCounters(redis={} if redis else None) + store: Final = handle.store() + hierarchy: Final = CustomerHierarchyAuth( + customer_org_id="tenant", + scopes=( + CustomerScopeAuth(key="customer_org", entity_id="tenant", spend=1, max_budget=5), + CustomerScopeAuth(key="customer_team", entity_id="team", spend=2), + CustomerScopeAuth(key="customer_member", entity_id="member", spend=3), + ), + ) + await settle_customer_spend({"user_api_key_customer_hierarchy": hierarchy.model_dump(mode="json")}, 0.5, store) + current: Final = await current_customer_spend(hierarchy, store) + assert tuple(scope.spend for scope in current.scopes) == (1.5, 2.5, 3.5) + assert tuple(scope.spend for scope in hierarchy.scopes) == (1.0, 2.0, 3.0) + + +@pytest.mark.asyncio +async def test_budget_reset_uses_a_new_counter_even_when_an_old_request_settles_late() -> None: + store: Final = FakeSpendCounters().store() + old: Final = CustomerHierarchyAuth( + customer_org_id="tenant", + scopes=(CustomerScopeAuth(key="customer_org", entity_id="tenant", spend=2, max_budget=1),), + ) + reset: Final = old.model_copy( + update={ + "scopes": ( + old.scopes[0].model_copy( + update={"spend": 0.0, "budget_reset_at": datetime(2026, 1, 1, tzinfo=timezone.utc)} + ), + ) + } + ) + await settle_customer_spend({"user_api_key_customer_hierarchy": old.model_dump(mode="json")}, 0.5, store) + current: Final = await current_customer_spend(reset, store) + check_customer_hierarchy_budget(current) + assert current.scopes[0].spend == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "metadata", ({}, {"user_api_key_customer_hierarchy": {}}, {"user_api_key_customer_hierarchy": "bad"}) +) +async def test_missing_or_invalid_customer_metadata_does_not_charge_a_counter(metadata: dict[str, object]) -> None: + handle: Final = FakeSpendCounters() + await settle_customer_spend(metadata, 1, handle.store()) + assert handle.redis == {} + + +@pytest.mark.asyncio +async def test_zero_cost_does_not_initialize_customer_counters() -> None: + handle: Final = FakeSpendCounters() + hierarchy: Final = CustomerHierarchyAuth(customer_org_id="tenant", scopes=()) + await settle_customer_spend( + {"user_api_key_customer_hierarchy": hierarchy.model_dump(mode="json")}, 0, handle.store() + ) + assert handle.redis == {} diff --git a/tests/unit/proxy/db/test_spend_counter_reseed.py b/tests/unit/proxy/db/test_spend_counter_reseed.py index 7314bfbb7c5..8406512c6e5 100644 --- a/tests/unit/proxy/db/test_spend_counter_reseed.py +++ b/tests/unit/proxy/db/test_spend_counter_reseed.py @@ -27,6 +27,48 @@ from litellm.redis.stores.spend_counters import SpendCounterStore WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "family,table_name,id_field", + ( + ("customer_org", "litellm_customerorganizationtable", "customer_org_id"), + ("customer_team", "litellm_customerteamtable", "customer_team_id"), + ("customer_member", "litellm_customermembertable", "customer_member_id"), + ), +) +@pytest.mark.parametrize( + "period,reset_at,expected", + ( + ("lifetime", None, 4.5), + (WINDOW_START.isoformat(), WINDOW_START, 4.5), + (WINDOW_START.isoformat(), WINDOW_START.replace(tzinfo=None), 4.5), + (WINDOW_START.isoformat(), WINDOW_START + timedelta(days=1), None), + (WINDOW_START.isoformat(), WINDOW_START - timedelta(days=1), None), + ("lifetime", WINDOW_START, None), + (WINDOW_START.isoformat(), None, None), + ("not-a-period", None, None), + ("2026-99-99T00:00:00+00:00", WINDOW_START, None), + ), +) +async def test_customer_counter_reseed_reads_only_the_matching_budget_period( + family, table_name, id_field, period, reset_at, expected +): + table: Final = _FakeFindUniqueTable(SimpleNamespace(spend=4.5, budget_reset_at=reset_at)) + prisma: Final = SimpleNamespace(db=SimpleNamespace(**{table_name: table})) + actual: Final = await SpendCounterReseed.from_db(prisma, f"spend:{family}:id:with:colons:{period}") + assert actual == expected + if period != "not-a-period": + assert table.where_clauses == [{id_field: "id:with:colons"}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", (None, ConnectionError("database is unavailable"))) +async def test_missing_or_unavailable_customer_spend_does_not_seed_a_counter(error): + table: Final = _FakeFindUniqueTable(None, error) + prisma: Final = SimpleNamespace(db=SimpleNamespace(litellm_customermembertable=table)) + assert await SpendCounterReseed.from_db(prisma, "spend:customer_member:member:lifetime") is None + + class _FakeFindUniqueTable: def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None: self._row = row diff --git a/tests/unit/redis/stores/test_spend_counters.py b/tests/unit/redis/stores/test_spend_counters.py index e7859873de7..a8f2c707631 100644 --- a/tests/unit/redis/stores/test_spend_counters.py +++ b/tests/unit/redis/stores/test_spend_counters.py @@ -47,6 +47,12 @@ _KEYS: Final = ( (project_counter("proj-1"), "spend:project:proj-1"), (model_access_group_counter("gpt-group"), "spend:model_access_group:gpt-group"), (shadow_eval_counter("job-7"), "spend:shadow_eval:job-7"), + ("spend:customer_org:tenant:lifetime", "spend:customer_org:tenant:lifetime"), + ( + "spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00", + "spend:customer_team:tenant:team:2026-01-01T00:00:00+00:00", + ), + ("spend:customer_member:user@example.com:lifetime", "spend:customer_member:user@example.com:lifetime"), (committed_counter(key_counter("88dc28d0")), "spend:committed:key:88dc28d0"), (committed_counter(team_counter("team-a")), "spend:committed:team:team-a"), (committed_counter(user_counter("user@example.com")), "spend:committed:user:user@example.com"),