From 1a749d84bdd66706bb41cafdba28e7a8b6a20fa9 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 01:56:00 +0000 Subject: [PATCH 1/9] 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> --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 34 ++-- .../proxy/common_utils/reset_budget_job.py | 24 +++ .../proxy/common_utils/user_api_key_cache.py | 10 ++ litellm/proxy/db/db_spend_update_writer.py | 79 ++++++++- .../redis_update_buffer.py | 7 + .../spend_update_queue.py | 4 + litellm/proxy/db/spend_counter_reseed.py | 5 + .../proxy/hooks/proxy_track_cost_callback.py | 6 + litellm/proxy/proxy_server.py | 30 ++++ .../spend_tracking/budget_reservation.py | 41 +++++ .../spend_tracking/spend_counter_batch.py | 11 +- litellm/repositories/prisma_protocols.py | 3 + litellm/repositories/unit_of_work.py | 2 + .../proxy/auth/test_auth_checks.py | 61 +++++++ .../common_utils/test_reset_budget_job.py | 30 +++- .../proxy/db/test_db_spend_update_writer.py | 76 +++++++++ .../proxy/db/test_spend_counter_reseed.py | 19 +++ .../hooks/test_proxy_track_cost_callback.py | 1 + .../test_spend_tracking_utils.py | 1 + .../proxy/test_budget_reservation.py | 156 ++++++++++++++++++ .../repositories/test_unit_of_work.py | 3 + 22 files changed, 583 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 321f8190f13..228a91ad446 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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] diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e3783c94dc7..ef17913d9ec 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index acb51e73daf..2baefa89943 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 1c7a379897f..2187ed63ea5 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -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__" diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index a90d1351fd7..51d00b9789c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 6f49a00b763..cead63795a2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -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( diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index bc068d10daf..2b8535cb113 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -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": diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index 89a07234c6c..2dd028454d6 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -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: diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be390..0c562cf37ef 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7964556531..1e84e5f56e2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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], diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 373f2d0fe36..24b8470145c 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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, diff --git a/litellm/proxy/spend_tracking/spend_counter_batch.py b/litellm/proxy/spend_tracking/spend_counter_batch.py index 7106d88c655..ddb074ae023 100644 --- a/litellm/proxy/spend_tracking/spend_counter_batch.py +++ b/litellm/proxy/spend_tracking/spend_counter_batch.py @@ -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( diff --git a/litellm/repositories/prisma_protocols.py b/litellm/repositories/prisma_protocols.py index 93b8c5c7cd7..60c16fbd746 100644 --- a/litellm/repositories/prisma_protocols.py +++ b/litellm/repositories/prisma_protocols.py @@ -152,4 +152,7 @@ class PrismaBatch(Protocol): @property def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ... + @property + def litellm_projecttable(self) -> BatchTable: ... + async def commit(self) -> None: ... diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index 0cdce307f9b..c09e5eb75d4 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -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), ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 8c8b755195f..0d5c0dd5d72 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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 diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..8c254686385 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -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"), } diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index c547d06904b..da5879a375a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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(): """ diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index bca6344b3f7..53e91b8792e 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -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 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index dfc95db3e14..965e134772d 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -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, ) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 8b105e94d19..834cf8d100d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -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 ( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 032722d3259..014f240d9cc 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -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. diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index b52b8ced31e..9eff248b917 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -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}), ] From b00cd15bd73a380c77aad0d04672d27ee4bc13cf Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:07:11 +0000 Subject: [PATCH 2/9] test(proxy): import project_cache_key from user_api_key_cache Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_endpoints/test_key_management_endpoints.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4e70063015d..b57f8b7857d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -38,11 +38,10 @@ from litellm.proxy._types import ( from litellm.models.object_permission import LiteLLM_ObjectPermissionTable from litellm.proxy.auth.auth_checks import ( _delete_cache_key_object, - _project_cache_key, jwt_key_mapping_cache_key, ) from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy.management_endpoints.key_management_endpoints import ( _check_org_key_limits, @@ -18951,7 +18950,7 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( - key=_project_cache_key(project_id), + key=project_cache_key(project_id), value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models), model_type=LiteLLM_ProjectTableCachedObj, ) From b20f1422eb204c2cb2fba26912a548454cd18b02 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 16 Sep 2026 02:28:38 +0000 Subject: [PATCH 3/9] fix(proxy): carry project_id through key metadata enrichment and drop docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/user_api_key_cache.py | 2 -- litellm/proxy/hooks/proxy_track_cost_callback.py | 2 ++ tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- .../proxy/common_utils/test_reset_budget_job.py | 1 - tests/test_litellm/proxy/db/test_db_spend_update_writer.py | 4 ---- tests/test_litellm/proxy/db/test_spend_counter_reseed.py | 2 -- .../proxy/hooks/test_proxy_track_cost_callback.py | 3 +++ tests/test_litellm/proxy/test_budget_reservation.py | 7 ------- 8 files changed, 5 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 2187ed63ea5..0386b58070d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -307,12 +307,10 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str: 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}" diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0c562cf37ef..5e525108ade 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -504,6 +504,8 @@ class _ProxyDBLogger(CustomLogger): metadata["user_api_key_team_id"] = key_obj.team_id if metadata.get("user_api_key_org_id") is None: metadata["user_api_key_org_id"] = key_obj.org_id + if metadata.get("user_api_key_project_id") is None: + metadata["user_api_key_project_id"] = key_obj.project_id except Exception: verbose_proxy_logger.debug( "Failed to enrich failure metadata with key info for api_key=%s", diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 0d5c0dd5d72..fe469fc574e 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7161,9 +7161,6 @@ def _project_with_budget(spend: float, max_budget: float): ], ) 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 diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 8c254686385..48b06649237 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1666,7 +1666,6 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( 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( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index da5879a375a..60cc742577b 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1062,10 +1062,6 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us @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, diff --git a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py index 53e91b8792e..ff0b67d426b 100644 --- a/tests/test_litellm/proxy/db/test_spend_counter_reseed.py +++ b/tests/test_litellm/proxy/db/test_spend_counter_reseed.py @@ -432,8 +432,6 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys(): @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 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 965e134772d..202495517ad 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1372,6 +1372,7 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): mock_key_obj.user_id = "fetched-user-id" mock_key_obj.team_id = "fetched-team-id" mock_key_obj.org_id = "fetched-org-id" + mock_key_obj.project_id = "fetched-project-id" mock_team_obj = MagicMock() mock_team_obj.team_alias = "fetched-team-alias" @@ -1395,12 +1396,14 @@ async def test_enrich_failure_metadata_with_full_key_lookup(): "user_api_key_team_id": None, "user_api_key_team_alias": None, "user_api_key_org_id": None, + "user_api_key_project_id": None, } result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) assert result["user_api_key_alias"] == "fetched-key-alias" assert result["user_api_key_user_id"] == "fetched-user-id" assert result["user_api_key_team_id"] == "fetched-team-id" assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_project_id"] == "fetched-project-id" assert result["user_api_key_team_alias"] == "fetched-team-alias" diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 014f240d9cc..c834ac05f0a 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -672,10 +672,6 @@ async def _seed_project_scoped_budgets( @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( @@ -724,8 +720,6 @@ async def test_should_reserve_project_and_team_member_counters_for_project_scope @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( @@ -755,7 +749,6 @@ async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spen @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( From 7e5b3b49d46ae1773e3f0450adbda334c198f108 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 22:33:54 +0000 Subject: [PATCH 4/9] fix(proxy): treat non-positive project max_budget as unbudgeted Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 2 +- .../proxy/auth/test_auth_checks.py | 23 +++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e80a4a677ce..e6313790380 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -5612,7 +5612,7 @@ 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 None or not math.isfinite(max_budget): + if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget): return from litellm.proxy.proxy_server import get_current_spend diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 56ee0aac249..08ed1575fdd 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7378,6 +7378,29 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" +@pytest.mark.asyncio +@pytest.mark.parametrize("max_budget", [0.0, -1.0]) +async def test_project_max_budget_check_treats_non_positive_budget_as_unbudgeted(max_budget): + from litellm.caching.dual_cache import DualCache + from litellm.proxy.auth.auth_checks import _project_max_budget_check + + real_spend_counter_cache = DualCache() + real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=12.5) + 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 + ): + await _project_max_budget_check( + project_object=_project_with_budget(spend=12.5, max_budget=max_budget), + valid_token=UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget"), + proxy_logging_obj=proxy_logging_obj, + ) + + proxy_logging_obj.budget_alerts.assert_not_awaited() + + 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 From 1c0342332beb96ba3c13ede34563c7e1ccd3c8a2 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 22:40:07 +0000 Subject: [PATCH 5/9] refactor(proxy): keep project spend enqueue within lint ceilings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 67 ++++++++-------------- 1 file changed, 25 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4f04a1b3a1a..e0985d5b5ae 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -706,17 +706,11 @@ 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(), - ) + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) try: await self._update_tag_db( @@ -985,27 +979,16 @@ class DBSpendUpdateWriter: 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, - ) + ) -> None: + 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, @@ -1246,17 +1229,17 @@ class DBSpendUpdateWriter: "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, " "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 {}), - len(db_spend_update_transactions.get("org_list_transactions") or {}), - 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 {}), + 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 ()), + len(db_spend_update_transactions.get("org_list_transactions") or ()), + 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 ()), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, From 70164dabd2795613a814f3e6d76c344e5f5ce618 Mon Sep 17 00:00:00 2001 From: ryan Date: Thu, 17 Sep 2026 23:20:10 +0000 Subject: [PATCH 6/9] fix(proxy): isolate project spend enqueue failures from sibling spend writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 16 +++++++---- .../proxy/db/test_db_spend_update_writer.py | 27 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e0985d5b5ae..2d6c7a0a74c 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -706,11 +706,17 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) - await self._update_project_db( - response_cost=response_cost, - project_id=project_id, - prisma_client=prisma_client, - ) + try: + await self._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=prisma_client, + ) + except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes + verbose_proxy_logger.debug( + "_batch_database_updates: _update_project_db failed: %s", + traceback.format_exc(), + ) try: await self._update_tag_db( diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index f9494cd6332..9a2b15d8a30 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1175,6 +1175,33 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( assert transactions["project_list_transactions"] == {} +@pytest.mark.asyncio +async def test_project_enqueue_failure_does_not_stop_sibling_spend_updates(): + db_writer: Final = DBSpendUpdateWriter() + db_writer._update_project_db = AsyncMock(side_effect=RuntimeError("project queue boom")) + db_writer._update_tag_db = AsyncMock() + db_writer._update_agent_db = AsyncMock() + db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + + await db_writer._batch_database_updates( + response_cost=0.1, + 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.1, "request_tags": ["t"]}, + project_id="proj-1", + ) + + db_writer._update_project_db.assert_awaited_once() + db_writer._update_tag_db.assert_awaited_once() + db_writer._update_agent_db.assert_awaited_once() + db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ From e4778cd2d1d14a5efed40546031d474fe53eab63 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 17 Sep 2026 16:24:51 -0700 Subject: [PATCH 7/9] test(proxy): drive the real spend queue in the project enqueue isolation test Substitutes a queue that rejects project items instead of replacing writer methods with mocks, so the test asserts the sibling key, team and tag spend actually landed in the queue rather than that a mock was awaited. --- .../proxy/db/test_db_spend_update_writer.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 9a2b15d8a30..3d61b1dba12 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -15,8 +15,9 @@ import pytest from redis.exceptions import DataError import litellm -from litellm.proxy._types import Litellm_EntityType +from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( build_window_spend_transaction, ) @@ -1176,30 +1177,34 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( @pytest.mark.asyncio -async def test_project_enqueue_failure_does_not_stop_sibling_spend_updates(): +async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): + class _ProjectRejectingQueue(SpendUpdateQueue): + async def add_update(self, update: SpendUpdateQueueItem): + if update.get("entity_type") is Litellm_EntityType.PROJECT: + raise RuntimeError("project enqueue boom") + await super().add_update(update) + db_writer: Final = DBSpendUpdateWriter() - db_writer._update_project_db = AsyncMock(side_effect=RuntimeError("project queue boom")) - db_writer._update_tag_db = AsyncMock() - db_writer._update_agent_db = AsyncMock() - db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + db_writer.spend_update_queue = _ProjectRejectingQueue() await db_writer._batch_database_updates( - response_cost=0.1, + response_cost=0.25, user_id="u1", hashed_token="t1", team_id="team-1", - org_id=None, + org_id="org-1", end_user_id=None, prisma_client=MagicMock(), litellm_proxy_budget_name=None, - payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1, "request_tags": ["t"]}, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, project_id="proj-1", ) - db_writer._update_project_db.assert_awaited_once() - db_writer._update_tag_db.assert_awaited_once() - db_writer._update_agent_db.assert_awaited_once() - db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() + transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert transactions["project_list_transactions"] == {} + assert transactions["tag_list_transactions"] == {"tag-1": 0.25} + assert transactions["key_list_transactions"] == {"t1": 0.25} + assert transactions["team_list_transactions"] == {"team-1": 0.25} @pytest.mark.asyncio From 89289d4f7753d3d780ce0cf3b679c409220c2b88 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 17 Sep 2026 16:26:26 -0700 Subject: [PATCH 8/9] fix(proxy): surface a failed project spend enqueue at error level The call-site guard kept a project enqueue failure from skipping the sibling spend writes, but logged it at debug only, so a dropped project charge was invisible on a default log level. Log it through spend_log_error inside the helper and re-raise, the way the org and agent helpers already do. --- litellm/proxy/db/db_spend_update_writer.py | 22 +++++++++---- .../proxy/db/test_db_spend_update_writer.py | 33 +++++++++++-------- 2 files changed, 36 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 2d6c7a0a74c..f715897b314 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -988,13 +988,23 @@ class DBSpendUpdateWriter: ) -> None: 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, + try: + 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, diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 3d61b1dba12..2efc0d99b8a 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -1,6 +1,7 @@ import asyncio import copy import json +import logging import re @@ -15,6 +16,7 @@ import pytest from redis.exceptions import DataError import litellm +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue @@ -1177,7 +1179,9 @@ async def test_batch_database_updates_without_project_id_touches_no_project_row( @pytest.mark.asyncio -async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): +async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_the_batch( + caplog: pytest.LogCaptureFixture, +): class _ProjectRejectingQueue(SpendUpdateQueue): async def add_update(self, update: SpendUpdateQueueItem): if update.get("entity_type") is Litellm_EntityType.PROJECT: @@ -1187,18 +1191,21 @@ async def test_failed_project_enqueue_does_not_drop_the_rest_of_the_batch(): db_writer: Final = DBSpendUpdateWriter() db_writer.spend_update_queue = _ProjectRejectingQueue() - await db_writer._batch_database_updates( - response_cost=0.25, - user_id="u1", - hashed_token="t1", - team_id="team-1", - org_id="org-1", - end_user_id=None, - prisma_client=MagicMock(), - litellm_proxy_budget_name=None, - payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, - project_id="proj-1", - ) + with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): + await db_writer._batch_database_updates( + response_cost=0.25, + user_id="u1", + hashed_token="t1", + team_id="team-1", + org_id="org-1", + end_user_id=None, + prisma_client=MagicMock(), + litellm_proxy_budget_name=None, + payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]}, + project_id="proj-1", + ) + + assert any("proj-1" in record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR) transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() assert transactions["project_list_transactions"] == {} From 595768b54b7ccfefdee9234c0bb70fbe219ee255 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 18 Sep 2026 14:33:41 -0700 Subject: [PATCH 9/9] fix(proxy): narrow project_id without a cast and fold the unbudgeted cases into the budget matrix test The lint job failed on one new typing.cast (LIT006) in the cost callback, and the test-quality gate behind it would have failed next on a test whose only assertion inspected a mock (TQ002). project_id is now narrowed with isinstance, and the zero and negative max_budget cases run through the existing parametrized budget test, which asserts the raised error or a clean admit with no alert --- .../proxy/hooks/proxy_track_cost_callback.py | 6 ++- .../proxy/auth/test_auth_checks.py | 45 ++++++------------- 2 files changed, 19 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 5e525108ade..903255c7b6c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -319,7 +319,11 @@ 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)) + project_id: Final = ( + project_id_value + if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str) + else 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) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index beadaa27831..85673df57ba 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7559,15 +7559,19 @@ def _project_with_budget(spend: float, max_budget: float): @pytest.mark.asyncio @pytest.mark.parametrize( - "counter_spend, db_spend, blocks", + "counter_spend, db_spend, max_budget, 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"), + pytest.param(5.0, 0.0, 5.0, True, id="counter-at-budget-blocks-despite-stale-db-row"), + pytest.param(4.99, 0.0, 5.0, False, id="counter-under-budget-admits"), + pytest.param(None, 5.0, 5.0, True, id="no-counter-falls-back-to-persisted-spend"), + pytest.param(None, 0.0, 5.0, False, id="no-counter-and-no-persisted-spend-admits"), + pytest.param(12.5, 12.5, 0.0, False, id="zero-budget-is-unbudgeted"), + pytest.param(12.5, 12.5, -1.0, False, id="negative-budget-is-unbudgeted"), ], ) -async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks): +async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_positive_budget( + counter_spend, db_spend, max_budget, blocks +): from litellm.caching.dual_cache import DualCache from litellm.proxy.auth.auth_checks import _project_max_budget_check @@ -7583,14 +7587,16 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, ): if not blocks: await _project_max_budget_check( - project_object=_project_with_budget(spend=db_spend, max_budget=5.0), + project_object=_project_with_budget(spend=db_spend, max_budget=max_budget), valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) + await asyncio.sleep(0) + proxy_logging_obj.budget_alerts.assert_not_awaited() 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), + project_object=_project_with_budget(spend=db_spend, max_budget=max_budget), valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, ) @@ -7603,29 +7609,6 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget" -@pytest.mark.asyncio -@pytest.mark.parametrize("max_budget", [0.0, -1.0]) -async def test_project_max_budget_check_treats_non_positive_budget_as_unbudgeted(max_budget): - from litellm.caching.dual_cache import DualCache - from litellm.proxy.auth.auth_checks import _project_max_budget_check - - real_spend_counter_cache = DualCache() - real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=12.5) - 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 - ): - await _project_max_budget_check( - project_object=_project_with_budget(spend=12.5, max_budget=max_budget), - valid_token=UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget"), - proxy_logging_obj=proxy_logging_obj, - ) - - proxy_logging_obj.budget_alerts.assert_not_awaited() - - 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