fix(proxy): track project spend and enforce project budgets additively

Project-scoped keys never wrote spend to LiteLLM_ProjectTable, so
/project/info stayed at 0 and project budgets could not block. Wire the
PROJECT entity through the spend queue, redis buffer, and db writer,
reserve and increment a spend:project counter, reseed it from the
project row, reset project spend in the budget cascade, and read the
live counter in the project max budget check. Team member budgets keep
gating project-scoped keys alongside the project budget

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-16 01:56:00 +00:00
parent 67cb0bc089
commit 1a749d84bd
22 changed files with 583 additions and 21 deletions

View file

@ -5261,6 +5261,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_member_list_transactions: dict[str, float] | None
org_list_transactions: dict[str, float] | None
org_member_list_transactions: ReadOnly[dict[str, float] | None]
project_list_transactions: ReadOnly[dict[str, float] | None]
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]

View file

@ -92,6 +92,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
object_permission_cache_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
@ -5586,16 +5588,22 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if (
max_budget is not None
and project_object.spend is not None
and math.isfinite(max_budget)
and project_object.spend > max_budget
):
if max_budget is None or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend
project_spend: Final = await get_current_spend(
counter_key=project_spend_counter_key(project_object.project_id),
fallback_spend=project_object.spend or 0.0,
max_budget=max_budget,
)
if project_spend >= max_budget:
if valid_token:
call_info: Final = CallInfo(
token=valid_token.token,
spend=project_object.spend,
spend=project_spend,
max_budget=max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
@ -5611,9 +5619,9 @@ async def _project_max_budget_check(
)
raise litellm.BudgetExceededError(
current_cost=project_object.spend,
current_cost=project_spend,
max_budget=max_budget,
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}",
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}",
entity_type=Litellm_EntityType.PROJECT.value,
entity_id=project_object.project_id,
)
@ -5663,10 +5671,6 @@ async def _project_soft_budget_check(
)
def _project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
async def get_project_object(
project_id: str,
prisma_client: PrismaClient | None,
@ -5684,7 +5688,7 @@ async def get_project_object(
return None
# Check cache first
cache_key: Final = _project_cache_key(project_id)
cache_key: Final = project_cache_key(project_id)
deserialized_project: Final = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_ProjectTableCachedObj,
@ -5726,7 +5730,7 @@ async def delete_cached_project_object(
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
await evict_and_broadcast(
cache_keys=(_project_cache_key(project_id),),
cache_keys=(project_cache_key(project_id),),
user_api_key_cache=user_api_key_cache,
)

View file

@ -41,6 +41,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
@ -49,6 +51,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import SpendLinkedTable
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
EndUserRepository,
ModelAccessGroupBudgetRepository,
@ -115,6 +118,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol):
def access_group_name(self) -> str: ...
class _ProjectRow(_BudgetLinkedRow, Protocol):
@property
def project_id(self) -> str: ...
class _EndUserRow(_BudgetLinkedRow, Protocol):
@property
def user_id(self) -> str: ...
@ -185,6 +193,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
return (model_access_group_cache_key(row.access_group_name),)
def _project_counter_key(row: _ProjectRow) -> str:
return project_spend_counter_key(row.project_id)
def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
return (project_cache_key(row.project_id),)
def _enduser_counter_key(row: _EndUserRow) -> str:
return f"spend:end_user:{row.user_id}"
@ -661,6 +677,11 @@ class ResetBudgetJob:
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="model access groups",
)
projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows(
table=ProjectRepository(self.prisma_client).table,
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="projects",
)
rollover_caps: Final[Mapping[str, float]] = MappingProxyType(
{ # mutable-ok: MappingProxyType wraps a one-shot dict comprehension
b.budget_id: cap
@ -695,6 +716,7 @@ class ResetBudgetJob:
(_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
for row in model_access_groups
),
*((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects),
*((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers),
),
rollover_caps=rollover_caps,
@ -704,6 +726,7 @@ class ResetBudgetJob:
*(key for row in orgs for key in _org_cache_keys(row)),
*(key for row in tags for key in _tag_cache_keys(row)),
*(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
*(key for row in projects for key in _project_cache_keys(row)),
*(key for row in endusers for key in _enduser_cache_keys(row)),
),
)
@ -731,6 +754,7 @@ class ResetBudgetJob:
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
_queue_enduser_resets(uow.endusers, cascade)
for budget_id, budget_reset_at in cascade.budget_resets:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)

View file

@ -306,6 +306,16 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str:
return f"spend:model_access_group:{access_group_name}"
def project_cache_key(project_id: str) -> str:
"""Cache key one project row is stored under; shared by auth, spend tracking and the spend writer."""
return f"project_id:{project_id}"
def project_spend_counter_key(project_id: str) -> str:
"""Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it."""
return f"spend:project:{project_id}"
#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds
#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch.
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__"

View file

@ -44,6 +44,7 @@ from litellm.proxy._types import (
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
)
from litellm.proxy.common_utils.user_api_key_cache import project_cache_key
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
@ -116,6 +117,7 @@ class _SpendBatch(Protocol):
litellm_teammembership: BatchTable
litellm_organizationtable: BatchTable
litellm_organizationmembership: BatchTable
litellm_projecttable: BatchTable
litellm_tagtable: BatchTable
litellm_agentstable: BatchTable
litellm_modelaccessgroupbudgettable: BatchTable
@ -254,6 +256,7 @@ class DBSpendUpdateWriter:
start_time: datetime | None,
end_time: datetime | None,
response_cost: float | None,
project_id: str | None = None,
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
@ -335,6 +338,7 @@ class DBSpendUpdateWriter:
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
project_id=project_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
litellm_proxy_budget_name=litellm_proxy_budget_name,
@ -631,6 +635,7 @@ class DBSpendUpdateWriter:
litellm_proxy_budget_name: str | None,
payload: SpendLogsPayload,
request_model_access_groups: Sequence[str] = (),
project_id: str | None = None,
):
"""
Runs all 13 spend-update helpers sequentially inside a single asyncio task.
@ -694,6 +699,18 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
await self._update_project_db(
response_cost=response_cost,
project_id=project_id,
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_project_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_tag_db(
response_cost=response_cost,
@ -956,6 +973,33 @@ class DBSpendUpdateWriter:
)
raise e
async def _update_project_db(
self,
response_cost: float | None,
project_id: str | None,
prisma_client: PrismaClient | None,
):
try:
if project_id is None or prisma_client is None:
return
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.PROJECT,
entity_id=project_id,
response_cost=response_cost,
)
)
except Exception as e:
spend_log_error(
"Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s",
project_id,
response_cost,
str(e),
exc=e,
)
raise e
async def _update_agent_db(
self,
response_cost: float | None,
@ -1193,8 +1237,8 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
"agents=%d, model_access_groups=%d",
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, "
"projects=%d, tags=%d, agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
@ -1202,6 +1246,7 @@ class DBSpendUpdateWriter:
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
len(db_spend_update_transactions.get("project_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
@ -1762,6 +1807,22 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE PROJECT TABLE ###
project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions")
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Project",
transactions=project_list_transactions,
table_accessor="litellm_projecttable",
where_field="project_id",
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
await DBSpendUpdateWriter._invalidate_project_caches(
project_ids=tuple(project_list_transactions or ()),
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(
@ -1800,11 +1861,23 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod
async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None:
if not project_ids or proxy_logging_obj is None:
return
user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache")
if user_api_key_cache is None:
return
for project_id in project_ids:
await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id))
@staticmethod
async def _update_entity_spend_in_db(
entity_name: str,
transactions: dict[str, float] | None,
table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"],
table_accessor: Literal[
"litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable"
],
where_field: str,
n_retry_times: int,
prisma_client: PrismaClient,

View file

@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -418,6 +420,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION_MEMBER,
db_spend_update_transactions.get("org_member_list_transactions"),
),
(
Litellm_EntityType.PROJECT,
db_spend_update_transactions.get("project_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -885,6 +891,7 @@ class RedisUpdateBuffer:
org_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "org_member_list_transactions"
),
project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"),
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
model_access_group_list_transactions=_merged_entity_transactions(

View file

@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_member_list_transactions={},
org_list_transactions={},
org_member_list_transactions={},
project_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
Litellm_EntityType.PROJECT: "project_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["org_list_transactions"]
elif dict_key == "org_member_list_transactions":
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
elif dict_key == "project_list_transactions":
transactions_dict = db_spend_update_transactions["project_list_transactions"]
elif dict_key == "tag_list_transactions":
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
elif dict_key == "agent_list_transactions":

View file

@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_lookup_gate import db_lookup_gate
from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
BudgetWindowSpendRepository,
EndUserRepository,
@ -77,6 +78,7 @@ class SpendCounterReseed:
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
spend:user:{user_id} -> LiteLLM_UserTable.spend
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
spend:project:{project_id} -> LiteLLM_ProjectTable.spend
End-user and tag spend counters intentionally do not reseed here. Their
auth paths already load the corresponding objects via get_end_user_object()
@ -157,6 +159,9 @@ class SpendCounterReseed:
row = await OrganizationRepository(prisma_client).table.find_unique(
where={"organization_id": org_id}
)
elif counter_key.startswith("spend:project:"):
project_id: Final = counter_key[len("spend:project:") :]
row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id})
else:
return None
except Exception:

View file

@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger):
start_time=actual_start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
@log_db_metrics
@ -318,6 +319,7 @@ class _ProxyDBLogger(CustomLogger):
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None))
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
@ -368,6 +370,7 @@ class _ProxyDBLogger(CustomLogger):
budget_reservation=budget_reservation,
request_tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
)
if not charged:
return
@ -651,6 +654,7 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
) -> bool:
if budget_reservation is not None:
await _reconcile_budget_reservation_before_db_update(
@ -668,6 +672,7 @@ async def _update_database_and_spend_counters(
start_time=start_time,
end_time=end_time,
org_id=org_id,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:
@ -698,6 +703,7 @@ async def _update_database_and_spend_counters(
tags=request_tags,
request_started_at=start_time,
model_access_groups=model_access_groups,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:

View file

@ -417,6 +417,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.config_resolvers import resolve_fields
@ -2780,6 +2782,7 @@ async def increment_spend_counters(
tags: list[str] | None = None,
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -2801,6 +2804,7 @@ async def increment_spend_counters(
end_user_id=end_user_id,
tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
),
):
await _increment_spend_counters_batched(
@ -2814,6 +2818,7 @@ async def increment_spend_counters(
tags=tags,
request_started_at=request_started_at,
model_access_groups=model_access_groups,
project_id=project_id,
)
@ -2828,6 +2833,7 @@ async def _increment_spend_counters_batched(
tags: list[str] | None,
request_started_at: datetime | None,
model_access_groups: Sequence[str] | None,
project_id: str | None = None,
):
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET."""
reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update(
@ -3028,6 +3034,13 @@ async def _increment_spend_counters_batched(
)
if org_id is not None
else None,
_prepare_project_spend_increment(
project_id=project_id,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
)
if project_id is not None
else None,
)
if coro is not None
)
@ -3180,6 +3193,23 @@ async def _prepare_org_spend_increment(
return (pending,) if pending is not None else ()
async def _prepare_project_spend_increment(
project_id: str | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> tuple[PendingSpendIncrement, ...]:
if project_id is None:
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=project_spend_counter_key(project_id),
source_cache_key=project_cache_key(project_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
return (pending,) if pending is not None else ()
async def _prepare_unreserved_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],

View file

@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import json
import math
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@ -29,6 +30,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
team_membership_reservation_cache_key,
)
@ -62,6 +65,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = {
"Tag": Litellm_EntityType.TAG.value,
"Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value,
"Organization": Litellm_EntityType.ORGANIZATION.value,
"Project": Litellm_EntityType.PROJECT.value,
}
@ -542,6 +546,13 @@ async def _get_budget_counters(
if org_counter is not None:
counters.append(org_counter)
project_counter: Final = await _get_project_budget_counter(
valid_token=valid_token,
user_api_key_cache=user_api_key_cache,
)
if project_counter is not None:
counters.append(project_counter)
return counters
@ -751,6 +762,36 @@ async def _get_org_budget_counter(
)
async def _get_project_budget_counter(
valid_token: UserAPIKeyAuth,
user_api_key_cache: UserApiKeyCache,
) -> _BudgetCounter | None:
if valid_token.project_id is None:
return None
source_cache_key: Final = project_cache_key(valid_token.project_id)
project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key)
if project_object is None:
return None
project_budget_table: Final = _get_value(project_object, "litellm_budget_table")
if project_budget_table is None:
return None
project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget"))
if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget):
return None
return _BudgetCounter(
counter_key=project_spend_counter_key(valid_token.project_id),
source_cache_key=source_cache_key,
max_budget=project_max_budget,
fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0,
entity_type="Project",
entity_id=valid_token.project_id,
)
def _get_budget_limit_counters(
entity_prefix: str,
entity_type: str,

View file

@ -12,7 +12,10 @@ from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key
from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_spend_counter_key,
project_spend_counter_key,
)
_CounterValues: Final = TypeAdapter(dict[str, float | None])
_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
@ -154,6 +157,8 @@ def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None)
yield f"spend:end_user:{end_user_id}"
if token.org_id is not None:
yield f"spend:org:{token.org_id}"
if token.project_id is not None:
yield project_spend_counter_key(token.project_id)
def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]:
@ -168,10 +173,12 @@ def post_call_counter_keys(
end_user_id: str | None,
tags: Sequence[object] | None,
model_access_groups: Sequence[object] | None,
project_id: str | None = None,
) -> frozenset[str]:
"""Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read."""
entity_keys: Final = admission_counter_keys(
UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id
UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id, project_id=project_id),
end_user_id,
)
tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str))
group_keys: Final = frozenset(

View file

@ -152,4 +152,7 @@ class PrismaBatch(Protocol):
@property
def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ...
@property
def litellm_projecttable(self) -> BatchTable: ...
async def commit(self) -> None: ...

View file

@ -109,6 +109,7 @@ class BudgetCascadeUnitOfWork:
organizations: LinkedSpendResetWrites
tags: LinkedSpendResetWrites
model_access_groups: LinkedSpendResetWrites
projects: LinkedSpendResetWrites
endusers: LinkedSpendResetWrites
budgets: BudgetWindowWrites
@ -135,6 +136,7 @@ async def budget_cascade_unit_of_work(
organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),
model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable),
projects=LinkedSpendResetWrites(table=batch.litellm_projecttable),
endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable),
budgets=BudgetWindowWrites(table=batch.litellm_budgettable),
)

View file

@ -7138,6 +7138,67 @@ async def test_project_allowlist_enforced_when_key_models_empty():
assert exc_info.value.code == "403"
def _project_with_budget(spend: float, max_budget: float):
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj
return LiteLLM_ProjectTableCachedObj(
project_id="p-budget",
team_id="t-1",
budget_id="b-1",
spend=spend,
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b-1", max_budget=max_budget),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"counter_spend, db_spend, blocks",
[
pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"),
pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"),
pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"),
],
)
async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks):
"""LIT-3269: project budget enforcement must read the cross-pod
``spend:project:{id}`` counter first and only fall back to the cached row's
spend, matching key/team/org checks. The boundary is inclusive (>=)."""
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.auth_checks import _project_max_budget_check
real_spend_counter_cache = DualCache()
if counter_spend is not None:
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend)
valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1")
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
if not blocks:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=5.0),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
return
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=5.0),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0)
assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value
assert exc_info.value.entity_id == "p-budget"
assert exc_info.value.current_cost == 5.0
proxy_logging_obj.budget_alerts.assert_awaited_once()
assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget"
def test_is_user_proxy_admin_rejects_view_only_admin():
"""This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an
Admin Viewer answering True here would gain every write route. Read parity for

View file

@ -78,6 +78,7 @@ class MockBatcher:
self.litellm_organizationtable = _Table("org", self)
self.litellm_tagtable = _Table("tag", self)
self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self)
self.litellm_projecttable = _Table("project", self)
self.litellm_endusertable = _Table("enduser", self)
async def commit(self):
@ -93,6 +94,7 @@ class MockDB:
self.litellm_organizationtable = MockTable()
self.litellm_tagtable = MockTable()
self.litellm_modelaccessgroupbudgettable = MockTable()
self.litellm_projecttable = MockTable()
self.batch_calls: List[Dict[str, Any]] = []
self.batchers: List[MockBatcher] = []
@ -1507,13 +1509,19 @@ _INVALIDATION_CASES = [
"spend:model_access_group:gpt-4-group",
{"model_access_group:gpt-4-group"},
),
(
"litellm_projecttable",
type("Project", (), {"project_id": "proj-1"}),
"spend:project:proj-1",
{"project_id:proj-1"},
),
]
@pytest.mark.parametrize(
"table_attr, linked_row, counter_key, cache_keys",
_INVALIDATION_CASES,
ids=["team_membership", "key", "org", "tag", "model_access_group"],
ids=["team_membership", "key", "org", "tag", "model_access_group", "project"],
)
def test_budget_table_reset_invalidates_counters_and_management_cache(
reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys
@ -1657,6 +1665,25 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first(
counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}")
def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch):
"""A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction."""
_make_counter_invalidation_job(monkeypatch)
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")]
mock_prisma_client.db.litellm_projecttable.set_find_many_results(
[type("Project", (), {"project_id": "proj-1", "spend": 12.0, "budget_id": "budget-due"})]
)
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}}
assert mock_prisma_client.db.litellm_projecttable.find_many_calls == [{"where": expected_where}]
writes = _batch_writes(mock_prisma_client, "project", op="update_many")
assert len(writes) == 1
assert writes[0]["where"] == expected_where
assert writes[0]["data"] == {"spend": 0}
assert mock_prisma_client.db.batchers[0].committed is True
def test_budget_cascade_carries_access_group_overage_when_rollover_enabled(
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
):
@ -1802,6 +1829,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo
("org", "update_many"),
("tag", "update_many"),
("model_access_group", "update_many"),
("project", "update_many"),
("enduser", "update_many"),
("budget", "update_many"),
}

View file

@ -1060,6 +1060,82 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us
assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1}
@pytest.mark.asyncio
async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted():
"""Regression for LIT-3269: a request made with a project-scoped key must
increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0
and the project budget never blocks. The cached project row is evicted so
the next auth check reads the fresh spend."""
db_writer: Final = DBSpendUpdateWriter()
await db_writer._batch_database_updates(
response_cost=0.25,
user_id="u1",
hashed_token="t1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25},
project_id="proj-1",
)
await db_writer._batch_database_updates(
response_cost=0.5,
user_id="u1",
hashed_token="t1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-2", "model": "gpt-4o-mini", "spend": 0.5},
project_id="proj-1",
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["project_list_transactions"] == {"proj-1": 0.75}
assert transactions["team_member_list_transactions"] == {"team_id::team-1::user_id::u1": 0.75}
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
user_api_key_cache: Final = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {"user_api_key_cache": user_api_key_cache}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_projecttable.update_many.assert_called_once_with(
where={"project_id": "proj-1"},
data={"spend": {"increment": 0.75}},
)
user_api_key_cache.async_delete_cache.assert_any_await(key="project_id:proj-1")
@pytest.mark.asyncio
async def test_batch_database_updates_without_project_id_touches_no_project_row():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id=None,
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1},
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["project_list_transactions"] == {}
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""

View file

@ -67,12 +67,14 @@ class _FakePrismaClient:
error: Exception | None = None,
end_user_row: SimpleNamespace | None = None,
end_user_error: Exception | None = None,
project_row: SimpleNamespace | None = None,
) -> None:
self.db = SimpleNamespace(
litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error),
litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total),
litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error),
litellm_verificationtoken=_InFlightCountingTable(),
litellm_projecttable=_FakeFindUniqueTable(row=project_row),
)
@ -428,6 +430,23 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys():
assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY
@pytest.mark.asyncio
async def test_from_db_reseeds_project_counter_from_the_project_row():
"""LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend,
so a fresh pod enforces the project budget against persisted spend rather than 0."""
prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25))
assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25
assert prisma.db.litellm_projecttable.where_clauses == [{"project_id": "proj-1"}]
@pytest.mark.asyncio
async def test_from_db_returns_none_for_a_missing_project_row():
prisma: Final = _FakePrismaClient(project_row=None)
assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None
@pytest.mark.asyncio
async def test_from_db_still_never_reads_the_end_user_row():
"""A cold end-user counter keeps seeding from the cached end-user object the auth

View file

@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
tags=["tag-a"],
request_started_at=start_time,
model_access_groups=("premium",),
project_id=None,
)

View file

@ -1187,6 +1187,7 @@ async def test_api_key_preserved_through_failure_hook_to_database():
start_time,
end_time,
org_id,
project_id=None,
):
"""Mock update_database and capture the payload it creates"""
from litellm.proxy.spend_tracking.spend_tracking_utils import (

View file

@ -22,6 +22,7 @@ from litellm.proxy._types import (
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
@ -631,6 +632,161 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_
await release_budget_reservation(reservation)
def _project_scoped_token() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
token="key-project-scoped",
spend=0.0,
user_id="user-proj",
team_id="team-proj",
project_id="proj-1",
)
async def _seed_project_scoped_budgets(
key_cache: DualCache,
team_member_spend: float,
team_member_max_budget: float,
project_spend: float,
project_max_budget: float,
) -> None:
await key_cache.async_set_cache(
key="team_membership:user-proj:team-proj",
value=LiteLLM_TeamMembership(
user_id="user-proj",
team_id="team-proj",
spend=team_member_spend,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_max_budget),
).model_dump(),
)
await key_cache.async_set_cache(
key="project_id:proj-1",
value=LiteLLM_ProjectTableCachedObj(
project_id="proj-1",
team_id="team-proj",
budget_id="project-budget-id",
spend=project_spend,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=project_max_budget),
).model_dump(),
)
@pytest.mark.asyncio
async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state):
"""LIT-3269: a key carrying user_id, team_id and project_id reserves against
both the team member counter and the project counter; neither replaces the
other. After the call the project counter reflects the real cost once, not
the reservation plus the post-call increment."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=0.1,
team_member_max_budget=1.0,
project_spend=0.2,
project_max_budget=1.0,
)
estimated = estimate_request_max_cost(request_body=_request_body(), route="/chat/completions", llm_router=None)
assert estimated is not None and estimated > 0
reservation = await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert reservation is not None
assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(
0.1 + estimated
)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated)
from litellm.proxy.proxy_server import increment_spend_counters
await increment_spend_counters(
token="key-project-scoped",
team_id="team-proj",
user_id="user-proj",
response_cost=0.05,
budget_reservation=reservation,
project_id="proj-1",
)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25)
assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15)
@pytest.mark.asyncio
async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state):
"""LIT-3269: the project budget is additive. A project with plenty of
headroom must not let a key through once its team member budget is spent."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=1.0,
team_member_max_budget=1.0,
project_spend=0.0,
project_max_budget=100.0,
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert "TeamMember=user-proj:team-proj" in str(exc_info.value)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0))
@pytest.mark.asyncio
async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state):
"""LIT-3269: with team member headroom left, the project budget alone blocks the key."""
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=0.0,
team_member_max_budget=100.0,
project_spend=5.0,
project_max_budget=5.0,
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert "Project=proj-1" in str(exc_info.value)
assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value
assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in (
None,
pytest.approx(0.0),
)
@pytest.mark.asyncio
async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state):
"""The reservation path mirrors the read path: no personal user counter for a team key.

View file

@ -35,6 +35,7 @@ class FakeBatch:
self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls)
self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls)
self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls)
self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls)
self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
async def commit(self) -> None:
@ -94,6 +95,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch():
uow.organizations.queue_spend_zero(where=linked)
uow.tags.queue_spend_zero(where=linked)
uow.model_access_groups.queue_spend_zero(where=linked)
uow.projects.queue_spend_zero(where=linked)
uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}})
uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
assert batch.commit_count == 0
@ -105,6 +107,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch():
("litellm_organizationtable.update_many", linked, {"spend": 0}),
("litellm_tagtable.update_many", linked, {"spend": 0}),
("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}),
("litellm_projecttable.update_many", linked, {"spend": 0}),
("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}),
("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}),
]