diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ccf3ea9952..271000bbe96 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4608,6 +4608,7 @@ class DBSpendUpdateTransactions(TypedDict): team_list_transactions: dict[str, float] | None team_member_list_transactions: dict[str, float] | None org_list_transactions: dict[str, float] | None + project_list_transactions: dict[str, float] | None # mutable-ok: matches sibling buffers tag_list_transactions: dict[str, float] | None agent_list_transactions: dict[str, float] | None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6b1d845ec95..9d80bd8a08d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3888,6 +3888,7 @@ async def _check_team_member_budget( and team_object.team_id is not None and valid_token is not None and valid_token.user_id is not None + and valid_token.project_id is None ): team_membership = await get_team_membership( user_id=valid_token.user_id, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 286837c8909..51d10a0bd3a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1037,6 +1037,63 @@ def _ensure_parent_otel_span_on_request_state(request: Request) -> None: request.state.parent_otel_span = parent_otel_span +async def _team_member_budget_check_for_key( + valid_token: UserAPIKeyAuth, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, +) -> None: + if valid_token.team_member_spend is None or valid_token.project_id is not None: + return + if prisma_client is None: + return + + _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" + + team_member_info = await user_api_key_cache.async_get_cache( + key=_cache_key, + model_type=LiteLLM_TeamMembership, + ) + if team_member_info is None and valid_token.user_id is not None and valid_token.team_id is not None: + _db_member = await TeamMembershipRepository(prisma_client).table.find_first( + where={ # mutable-ok: prisma query argument shape + "user_id": valid_token.user_id, + "team_id": valid_token.team_id, + }, + include={"litellm_budget_table": True}, # mutable-ok: prisma query argument shape + ) + if _db_member is not None: + team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) + await user_api_key_cache.async_set_cache( + key=_cache_key, + value=team_member_info, + model_type=LiteLLM_TeamMembership, + ttl=5, + ) + + if team_member_info is None or team_member_info.litellm_budget_table is None: + return + team_member_budget = team_member_info.litellm_budget_table.max_budget + if team_member_budget is None or team_member_budget <= 0: + return + + from litellm.proxy.proxy_server import get_current_spend + + team_member_spend = valid_token.team_member_spend + if valid_token.user_id is not None and valid_token.team_id is not None: + team_member_spend = await get_current_spend( + counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", + fallback_spend=team_member_spend, + max_budget=team_member_budget, + ) + if team_member_spend > team_member_budget: + raise litellm.BudgetExceededError( + current_cost=team_member_spend, + max_budget=team_member_budget, + entity_type=Litellm_EntityType.TEAM_MEMBER.value, + entity_id=f"{valid_token.user_id}:{valid_token.team_id}", + ) + + async def _user_api_key_auth_builder( request: Request, api_key: str, @@ -1778,56 +1835,12 @@ async def _user_api_key_auth_builder( verbose_proxy_logger.info(f"Skipping all budget checks for zero-cost model: {model}") # Check 3. Check if user is in their team budget - if not skip_budget_checks and valid_token.team_member_spend is not None: - if prisma_client is not None: - _cache_key = f"{valid_token.team_id}_{valid_token.user_id}" - - team_member_info = await user_api_key_cache.async_get_cache( - key=_cache_key, - model_type=LiteLLM_TeamMembership, - ) - if team_member_info is None: - # read from DB - _user_id = valid_token.user_id - _team_id = valid_token.team_id - - if _user_id is not None and _team_id is not None: - _db_member = await TeamMembershipRepository(prisma_client).table.find_first( - where={ - "user_id": _user_id, - "team_id": _team_id, - }, # type: ignore - include={"litellm_budget_table": True}, - ) - if _db_member is not None: - team_member_info = LiteLLM_TeamMembership(**_db_member.dict()) - await user_api_key_cache.async_set_cache( - key=_cache_key, - value=team_member_info, - model_type=LiteLLM_TeamMembership, - ttl=5, - ) - - if team_member_info is not None and team_member_info.litellm_budget_table is not None: - team_member_budget = team_member_info.litellm_budget_table.max_budget - if team_member_budget is not None and team_member_budget > 0: - # Read from cross-pod counter (Redis-first) if available - from litellm.proxy.proxy_server import get_current_spend - - team_member_spend = valid_token.team_member_spend - if valid_token.user_id is not None and valid_token.team_id is not None: - team_member_spend = await get_current_spend( - counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}", - fallback_spend=team_member_spend, - max_budget=team_member_budget, - ) - if team_member_spend > team_member_budget: - raise litellm.BudgetExceededError( - current_cost=team_member_spend, - max_budget=team_member_budget, - entity_type=Litellm_EntityType.TEAM_MEMBER.value, - entity_id=f"{valid_token.user_id}:{valid_token.team_id}", - ) + if not skip_budget_checks: + await _team_member_budget_check_for_key( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) # Check 3. If token is expired if valid_token.expires is not None: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 17410698aed..2f8316edba6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -131,6 +131,7 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, + project_id: str | None = None, ): from litellm.proxy.proxy_server import ( disable_spend_logs, @@ -197,6 +198,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, @@ -346,6 +348,7 @@ class DBSpendUpdateWriter: hashed_token: str | None, team_id: str | None, org_id: str | None, + project_id: str | None, end_user_id: str | None, prisma_client: PrismaClient | None, litellm_proxy_budget_name: str | None, @@ -412,6 +415,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: # noqa: BLE001 # one failing spend helper must not block the others + 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, @@ -656,6 +671,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, @@ -832,13 +874,16 @@ 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, tags=%d, agents=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, projects=%d, tags=%d, agents=%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("project_list_transactions") or {} + ), # mutable-ok: matches sibling transaction count reads len(db_spend_update_transactions.get("tag_list_transactions") or {}), len(db_spend_update_transactions.get("agent_list_transactions") or {}), ) @@ -1298,6 +1343,23 @@ class DBSpendUpdateWriter: e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) + ### UPDATE PROJECT TABLE ### + project_list_transactions = 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, + ) + if project_list_transactions and proxy_logging_obj is not None: + user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") + if user_api_key_cache is not None: + for project_id in project_list_transactions: + await user_api_key_cache.async_delete_cache(key=f"project_id:{project_id}") + ### UPDATE TAG TABLE ### tag_list_transactions = db_spend_update_transactions["tag_list_transactions"] await DBSpendUpdateWriter._update_entity_spend_in_db( 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 69614defb9e..e53c5ca952f 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -337,6 +337,10 @@ class RedisUpdateBuffer: Litellm_EntityType.ORGANIZATION, db_spend_update_transactions.get("org_list_transactions"), ), + ( + Litellm_EntityType.PROJECT, + db_spend_update_transactions.get("project_list_transactions"), + ), ( Litellm_EntityType.TAG, db_spend_update_transactions.get("tag_list_transactions"), @@ -715,6 +719,7 @@ class RedisUpdateBuffer: team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, + project_list_transactions={}, # mutable-ok: matches sibling transaction buffers combined in place tag_list_transactions={}, agent_list_transactions={}, ) @@ -727,6 +732,7 @@ class RedisUpdateBuffer: "team_list_transactions", "team_member_list_transactions", "org_list_transactions", + "project_list_transactions", "tag_list_transactions", "agent_list_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 43383e7b5d5..20293846ef1 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -136,6 +136,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_list_transactions={}, team_member_list_transactions={}, org_list_transactions={}, + project_list_transactions={}, # mutable-ok: matches sibling transaction buffers aggregated in place tag_list_transactions={}, agent_list_transactions={}, ) @@ -148,6 +149,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM: "team_list_transactions", Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", + Litellm_EntityType.PROJECT: "project_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", Litellm_EntityType.AGENT: "agent_list_transactions", } @@ -185,6 +187,8 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions["team_member_list_transactions"] elif dict_key == "org_list_transactions": transactions_dict = db_spend_update_transactions["org_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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0319a680714..b440cbfb4f5 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -176,6 +176,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 @@ -210,6 +211,7 @@ class _ProxyDBLogger(CustomLogger): user_id = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id = cast(str | None, metadata.get("user_api_key_org_id", None)) + project_id = cast(str | None, metadata.get("user_api_key_project_id", None)) # cast-ok: untyped metadata key_alias = cast(str | None, metadata.get("user_api_key_alias", None)) end_user_max_budget = metadata.get("user_api_end_user_max_budget", None) sl_object: StandardLoggingPayload | None = kwargs.get("standard_logging_object", None) @@ -247,6 +249,7 @@ class _ProxyDBLogger(CustomLogger): end_user_id=end_user_id, team_id=team_id, org_id=org_id, + project_id=project_id, kwargs=kwargs, completion_response=completion_response, start_time=start_time, @@ -483,6 +486,7 @@ async def _update_database_and_spend_counters( end_user_id: str | None, team_id: str | None, org_id: str | None, + project_id: str | None, kwargs: dict, completion_response: litellm.ModelResponse | Any | None, start_time: Any, @@ -503,6 +507,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: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 50ea349b4b7..c0e23fccae1 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -528,6 +528,9 @@ async def _get_team_member_budget_counter( if team_object is None or team_object.team_id is None or user_object is None or valid_token.user_id is None: return None + if valid_token.project_id is not None: + return None + membership_cache_key = f"team_membership:{valid_token.user_id}:{team_object.team_id}" cached_team_membership = await user_api_key_cache.async_get_cache(key=membership_cache_key) team_membership: LiteLLM_TeamMembership | None = None diff --git a/tests/test_litellm/proxy/auth/test_team_member_budget.py b/tests/test_litellm/proxy/auth/test_team_member_budget.py index b38a953d189..e4ec62810fe 100644 --- a/tests/test_litellm/proxy/auth/test_team_member_budget.py +++ b/tests/test_litellm/proxy/auth/test_team_member_budget.py @@ -436,3 +436,76 @@ async def test_team_member_budget_check_personal_key_not_team(): # Should pass and get_team_membership should not be called assert result is True mock_get_team_membership.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_member_budget_check_skipped_for_project_scoped_key(): + """ + Regression for project-scoped keys being blocked by the team member budget: + when a key carries a project_id, the project budget governs and the team + member budget check must not raise, even if the member is over budget. + """ + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "test"}], + } + + team_object = LiteLLM_TeamTable( + team_id="test-team-1", + team_alias="Test Team", + spend=0.0, + max_budget=None, + ) + + user_object = LiteLLM_UserTable( + user_id="test-user-1", + spend=0.0, + max_budget=None, + ) + + valid_token = UserAPIKeyAuth( + token="test-token", + user_id="test-user-1", + team_id="test-team-1", + project_id="test-project-1", + models=["gpt-3.5-turbo"], + ) + + team_membership = LiteLLM_TeamMembership( + user_id="test-user-1", + team_id="test-team-1", + spend=0.0000002, # Exceeds budget + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=0.0000001, + ), + ) + + mock_request = MagicMock(spec=Request) + mock_prisma_client = MagicMock() + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + with ( + patch( + "litellm.proxy.auth.auth_checks.get_team_membership", + new_callable=AsyncMock, + return_value=team_membership, + ), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", mock_user_api_key_cache), + ): + result = await common_checks( + request_body=request_body, + team_object=team_object, + user_object=user_object, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging_obj, + valid_token=valid_token, + request=mock_request, + ) + + assert result is True diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 3177fc5ba44..742c5771a8f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -5439,3 +5439,69 @@ async def test_temp_budget_increase_applied_for_cached_key(): cached_after = await user_api_key_cache.async_get_cache(key=hashed_token) assert cached_after.max_budget == 2.0 + + +class TestTeamMemberBudgetCheckForKey: + """Tests for _team_member_budget_check_for_key: virtual keys with a + team member budget must be blocked once over budget, except project-scoped + keys, which are governed by the project budget instead.""" + + def _valid_token(self, project_id=None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + token="hashed-token", + user_id="member-user", + team_id="member-team", + team_member_spend=0.03, + project_id=project_id, + ) + + def _cache_with_membership(self): + from litellm.proxy._types import LiteLLM_TeamMembership + + membership = LiteLLM_TeamMembership( + user_id="member-user", + team_id="member-team", + spend=0.03, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.01), + ) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=membership) + return cache + + @pytest.mark.asyncio + async def test_raises_when_member_over_budget(self): + from litellm.proxy.auth.user_api_key_auth import ( + _team_member_budget_check_for_key, + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=0.03, + ): + with pytest.raises(litellm.BudgetExceededError): + await _team_member_budget_check_for_key( + valid_token=self._valid_token(), + prisma_client=MagicMock(), + user_api_key_cache=self._cache_with_membership(), + ) + + @pytest.mark.asyncio + async def test_skips_for_project_scoped_key(self): + from litellm.proxy.auth.user_api_key_auth import ( + _team_member_budget_check_for_key, + ) + + cache = self._cache_with_membership() + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=0.03, + ): + await _team_member_budget_check_for_key( + valid_token=self._valid_token(project_id="project-1"), + prisma_client=MagicMock(), + user_api_key_cache=cache, + ) + + cache.async_get_cache.assert_not_called() diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py index 0ed5940dd75..1ba5e0498f2 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py @@ -296,3 +296,27 @@ async def test_queue_size_reduction_with_large_volume(monkeypatch, spend_queue): ) assert aggregated["user_list_transactions"]["user1"] == 200 * 0.5 assert aggregated["key_list_transactions"]["key1"] == 300 * 1.0 + + +@pytest.mark.asyncio +async def test_project_entity_aggregation(spend_queue): + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.PROJECT, + "entity_id": "project-1", + "response_cost": 0.5, + } + ) + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.PROJECT, + "entity_id": "project-1", + "response_cost": 0.25, + } + ) + + aggregated = ( + await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) + + assert aggregated["project_list_transactions"]["project-1"] == 0.75 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 191080e3a48..0e00418fce6 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 @@ -1638,6 +1638,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer._update_user_db = AsyncMock() db_writer._update_team_db = AsyncMock() db_writer._update_org_db = AsyncMock() + db_writer._update_project_db = AsyncMock() db_writer._update_tag_db = AsyncMock() db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() @@ -1653,6 +1654,7 @@ async def test_batch_database_updates_isolation_on_failure(): hashed_token="t1", team_id="team1", org_id="org1", + project_id="project1", end_user_id="eu1", prisma_client=MagicMock(), litellm_proxy_budget_name="budget", @@ -1664,6 +1666,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer._update_key_db.assert_awaited_once() db_writer._update_team_db.assert_awaited_once() db_writer._update_org_db.assert_awaited_once() + 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() @@ -2236,3 +2239,178 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 assert transaction["prompt_caching_savings_spend"] == 0 + + +@pytest.mark.asyncio +async def test_update_project_db_enqueues_project_spend(): + """ + Regression for project budgets never being enforced: _update_project_db must + enqueue a SpendUpdateQueueItem with entity_type=PROJECT so LiteLLM_ProjectTable.spend + gets incremented. + """ + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + project_id = "project-123" + response_cost = 0.1 + + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_project_db( + response_cost=response_cost, + project_id=project_id, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_called_once() + call_args = writer.spend_update_queue.add_update.call_args[1] + assert call_args["update"]["entity_type"] == Litellm_EntityType.PROJECT + assert call_args["update"]["entity_id"] == project_id + assert call_args["update"]["response_cost"] == response_cost + + +@pytest.mark.asyncio +async def test_update_project_db_skips_when_project_id_none(): + """_update_project_db does not enqueue when project_id is None.""" + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_project_db( + response_cost=0.05, + project_id=None, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_database_updates_tracks_project_spend(): + """ + _batch_database_updates must call _update_project_db with the project_id it + was given, otherwise project spend silently stops being tracked. + """ + writer = DBSpendUpdateWriter() + writer._update_project_db = AsyncMock() + writer._update_user_db = AsyncMock() + writer._update_key_db = AsyncMock() + writer._update_team_db = AsyncMock() + writer._update_org_db = AsyncMock() + writer._update_tag_db = AsyncMock() + writer._update_agent_db = AsyncMock() + writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() + writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() + writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() + writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock() + writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() + writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + + mock_prisma = MagicMock() + await writer._batch_database_updates( + response_cost=0.42, + user_id="user-1", + hashed_token="hashed-token", + team_id="team-1", + org_id=None, + project_id="project-xyz", + end_user_id=None, + prisma_client=mock_prisma, + litellm_proxy_budget_name=None, + payload={"request_tags": None}, + ) + + writer._update_project_db.assert_called_once() + call_kwargs = writer._update_project_db.call_args[1] + assert call_kwargs["project_id"] == "project-xyz" + assert call_kwargs["response_cost"] == 0.42 + + +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_increments_project_spend_and_invalidates_cache(): + """ + _commit_spend_updates_to_db must increment LiteLLM_ProjectTable.spend and + invalidate the cached project object so the next auth check sees fresh spend. + """ + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + mock_cache = MagicMock() + mock_cache.async_delete_cache = AsyncMock() + mock_proxy_logging = MagicMock() + mock_proxy_logging.call_details.get = MagicMock(return_value=mock_cache) + + project_id = "project-789" + response_cost = 0.25 + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "project_list_transactions": {project_id: response_cost}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=db_spend_update_transactions, + ) + + mock_batcher.litellm_projecttable.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_projecttable.update_many.call_args[1] + assert call_kwargs["where"] == {"project_id": project_id} + assert call_kwargs["data"] == {"spend": {"increment": response_cost}} + mock_cache.async_delete_cache.assert_any_call(key=f"project_id:{project_id}") + + +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_handles_missing_project_transactions_key(): + """ + Transactions buffered by an older pod (rolling upgrade) have no + project_list_transactions key; commit must not raise. + """ + db_writer = DBSpendUpdateWriter() + + mock_prisma_client = MagicMock() + mock_proxy_logging = MagicMock() + mock_proxy_logging.call_details.get = MagicMock(return_value=None) + + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=db_spend_update_transactions, + ) + + mock_prisma_client.db.tx.assert_not_called() 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 f289148101a..d49528a0745 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 @@ -362,6 +362,7 @@ async def test_update_database_and_spend_counters_releases_reservation_when_db_u end_user_id=None, team_id="test_team_id", org_id="test_org_id", + project_id=None, kwargs={}, completion_response=None, start_time=datetime.now(), @@ -411,6 +412,7 @@ async def test_update_database_and_spend_counters_preserves_db_exception_when_re end_user_id=None, team_id="test_team_id", org_id="test_org_id", + project_id=None, kwargs={}, completion_response=None, start_time=datetime.now(), @@ -452,6 +454,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda end_user_id="test_end_user_id", team_id="test_team_id", org_id="test_org_id", + project_id=None, kwargs={}, completion_response=None, start_time=datetime.now(), @@ -497,6 +500,7 @@ async def test_update_database_and_spend_counters_invalidates_reservation_when_c end_user_id=None, team_id="test_team_id", org_id="test_org_id", + project_id=None, kwargs={}, completion_response=None, start_time=datetime.now(), @@ -543,6 +547,7 @@ async def test_update_database_and_spend_counters_preserves_counter_exception_wh end_user_id=None, team_id="test_team_id", org_id="test_org_id", + project_id=None, kwargs={}, completion_response=None, start_time=datetime.now(), diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 1a584423fac..9ba4469a21d 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2542,3 +2542,46 @@ async def test_streaming_slow_path_processes_and_yields_chunk(spend_counter_stat assert received == [{"content": "hi"}] streaming_logging_obj.async_post_call_streaming_hook.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_team_member_budget_counter_skipped_for_project_scoped_key(): + """Project-scoped keys are governed by the project budget, so no team + member budget counter should be reserved for them.""" + from litellm.proxy.spend_tracking.budget_reservation import ( + _get_team_member_budget_counter, + ) + + membership = LiteLLM_TeamMembership( + user_id="member-user", + team_id="member-team", + spend=0.03, + litellm_budget_table=LiteLLM_BudgetTable(max_budget=0.01), + ) + cache = MagicMock() + cache.async_get_cache = AsyncMock(return_value=membership) + team_object = LiteLLM_TeamTable(team_id="member-team") + user_object = LiteLLM_UserTable(user_id="member-user") + + counter = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth( + token="hashed", user_id="member-user", team_id="member-team" + ), + team_object=team_object, + user_object=user_object, + user_api_key_cache=cache, + ) + assert counter is not None + + counter = await _get_team_member_budget_counter( + valid_token=UserAPIKeyAuth( + token="hashed", + user_id="member-user", + team_id="member-team", + project_id="project-1", + ), + team_object=team_object, + user_object=user_object, + user_api_key_cache=cache, + ) + assert counter is None