scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): track project spend and let project budget govern project-scoped keys

Browse source 
Project spend was never written to LiteLLM_ProjectTable, so the project
max_budget check in auth could never fire. Keys created with user_id,
team_id, and project_id were instead blocked by the team member budget,
which surprised customers who set a larger project budget.

This wires Litellm_EntityType.PROJECT through the spend update pipeline
(queue aggregation, redis buffer, batched db commit with cache
invalidation) and passes project_id from the cost tracking callback.
Team member budget checks (auth hot path, common_checks, and budget
reservation counters) now skip project-scoped keys; the project budget
governs them instead. Non-project keys keep the existing team member
budget enforcement
This commit is contained in:
Shivam Rawat 2026-08-03 16:43:24 -07:00
parent 8ad5d144a1
commit 3d2932edb6
14 changed files with 535 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -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

1
litellm/proxy/auth/auth_checks.py
View file

@ -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,

113
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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:

64
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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(

6
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -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",
]

4
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -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":

5
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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:

3
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -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

73
tests/test_litellm/proxy/auth/test_team_member_budget.py
View file

@ -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

66
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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()

24
tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
View file

@ -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

178
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -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()

5
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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(),

43
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -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