scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

perf(proxy): read budget-window spend from the maintained window table

Browse source 
Per-window budget enforcement aggregated LiteLLM_SpendLogs on every cold
counter and on every 5s authoritative floor check. SpendLogs has no index on
api_key or team_id, so each check range-scanned the highest-volume table.

Reads now hit the LiteLLM_BudgetWindowSpend row by primary key and only fall
back to the aggregate when no row exists for the window being enforced. A row
is current when its window_start is at or past the caller's expected start, so
a pod holding a stale reset_at trusts a window another pod already rolled
instead of summing the previous window back in.

window_duration is threaded from the budget_limits entry through to the read
rather than parsed back out of the counter key. The reader never writes rows.
This commit is contained in:
ryan-crabbe-berri 2026-08-04 17:39:51 -07:00
parent 187e0fab60
commit 124d08d592
8 changed files with 478 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3710,6 +3710,7 @@ async def _virtual_key_multi_budget_check(
max_budget=w["max_budget"],
window_entity_type="Key",
window_entity_id=valid_token.token,
window_duration=str(w["budget_duration"]),
window_start=get_budget_window_start(w),
)
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:
@ -4083,6 +4084,7 @@ async def _team_multi_budget_check(
max_budget=w["max_budget"],
window_entity_type="Team",
window_entity_id=team_object.team_id,
window_duration=str(w["budget_duration"]),
window_start=get_budget_window_start(w),
)
if math.isfinite(w["max_budget"]) and window_spend >= w["max_budget"]:

134
litellm/proxy/db/spend_counter_reseed.py
View file

@ -14,14 +14,18 @@ memory in long-lived deployments.
import asyncio
from collections import OrderedDict
from datetime import datetime
from collections.abc import Mapping
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, Optional
from litellm._logging import verbose_proxy_logger
from litellm.constants import SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import Litellm_EntityType
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.table_repositories import (
BudgetWindowSpendRepository,
SpendLogsRepository,
TeamMembershipRepository,
)
@ -36,6 +40,25 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
_WINDOW_SPEND_ENTITY_TYPES: Final[Mapping[str, str]] = MappingProxyType(
{
"Key": Litellm_EntityType.KEY.value,
"Team": Litellm_EntityType.TEAM.value,
}
)
_WINDOW_SPEND_LOG_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{
"Key": "api_key",
"Team": "team_id",
}
)
def _as_utc(value: datetime) -> datetime:
return value if value.tzinfo is not None else value.replace(tzinfo=timezone.utc)
class SpendCounterReseed:
"""
Reseeds spend counters from the authoritative DB and warms the cache,
@ -205,6 +228,92 @@ class SpendCounterReseed:
raise
return current_value
@staticmethod
async def window_from_table(
prisma_client: Optional["PrismaClient"],
entity_type: str,
entity_id: str,
window_duration: str,
expected_window_start: datetime,
) -> float | None:
"""
Read the maintained per-window spend row by primary key.
Returns the row's spend only when the row belongs to the window the
caller is enforcing, i.e. ``row.window_start >= expected_window_start``.
A row at or past the expected start was rolled by a pod whose reset_at
was at least as fresh as this caller's, so it is trusted; an older row
means the window boundary was crossed and nothing has rolled the row
yet, so its spend belongs to a previous window.
Returns None for a missing, stale or unreadable row so the caller falls
back to the spend-logs aggregate. ``entity_type`` is the counter-facing
label ("Key"/"Team"); anything else has no row and returns None.
"""
if prisma_client is None:
return None
row_entity_type: Final = _WINDOW_SPEND_ENTITY_TYPES.get(entity_type)
if row_entity_type is None:
return None
try:
row: Final = await BudgetWindowSpendRepository(prisma_client).table.find_unique(
where={
"entity_type_entity_id_window_duration": {
"entity_type": row_entity_type,
"entity_id": entity_id,
"window_duration": window_duration,
}
}
)
except Exception:
verbose_proxy_logger.exception(
"SpendCounterReseed.window_from_table: failed for %s=%s window=%s",
entity_type,
entity_id,
window_duration,
)
return None
if row is None:
return None
if _as_utc(row.window_start) < _as_utc(expected_window_start):
return None
return float(row.spend or 0.0)
@staticmethod
async def window_from_db(
prisma_client: Optional["PrismaClient"],
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime,
) -> float | None:
"""
Authoritative window spend: the maintained row first, falling back to
the spend-logs aggregate only when no current row exists.
The aggregate range-scans an unindexed table, so it must stay a
transitional path (window configured before the row existed) rather
than a steady-state read.
"""
if window_duration is not None:
from_table: Final = await SpendCounterReseed.window_from_table(
prisma_client=prisma_client,
entity_type=entity_type,
entity_id=entity_id,
window_duration=window_duration,
expected_window_start=window_start,
)
if from_table is not None:
return from_table
return await SpendCounterReseed.window_from_spend_logs(
prisma_client=prisma_client,
entity_type=entity_type,
entity_id=entity_id,
window_start=window_start,
)
@staticmethod
async def window_from_spend_logs(
prisma_client: Optional["PrismaClient"],
@ -215,20 +324,13 @@ class SpendCounterReseed:
if prisma_client is None:
return None
if entity_type == "Key":
group_field = "api_key"
where = {
"api_key": entity_id,
"startTime": {"gte": window_start},
}
elif entity_type == "Team":
group_field = "team_id"
where = {
"team_id": entity_id,
"startTime": {"gte": window_start},
}
else:
group_field: Final = _WINDOW_SPEND_LOG_FIELDS.get(entity_type)
if group_field is None:
return None
where: Final = {
group_field: entity_id,
"startTime": {"gte": window_start},
}
try:
response: Final = await SpendLogsRepository(prisma_client).table.group_by(
@ -258,6 +360,7 @@ class SpendCounterReseed:
counter_key: str,
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime,
) -> float | None:
lock: Final = await SpendCounterReseed._get_lock(counter_key)
@ -276,10 +379,11 @@ class SpendCounterReseed:
if val is not None:
return float(val)
window_spend: Final = await SpendCounterReseed.window_from_spend_logs(
window_spend: Final = await SpendCounterReseed.window_from_db(
prisma_client=prisma_client,
entity_type=entity_type,
entity_id=entity_id,
window_duration=window_duration,
window_start=window_start,
)
if window_spend is None:

15
litellm/proxy/proxy_server.py
View file

@ -2184,6 +2184,7 @@ async def get_current_spend(
max_budget: float | None = None,
window_entity_type: str | None = None,
window_entity_id: str | None = None,
window_duration: str | None = None,
window_start: datetime | None = None,
fallback_authoritative: bool = False,
) -> float:
@ -2208,7 +2209,8 @@ async def get_current_spend(
runs and a key can leak spend past ``max_budget`` indefinitely. The
authoritative source depends on the counter: primary key/team/user/org
counters read the DB row; per-window counters (``window_start`` supplied)
aggregate spend logs; end-user/tag counters have no DB row, so the caller's
read the maintained window-spend row and only aggregate spend logs when
that row is missing or stale; end-user/tag counters have no DB row, so the caller's
``fallback_spend`` (loaded fresh in auth) is authoritative. The DB read is
skipped for healthy primary counters (counter at or above recorded spend)
and cached in-process for a few seconds, so a persistently stale counter
@ -2233,6 +2235,7 @@ async def get_current_spend(
counter_key=counter_key,
window_entity_type=window_entity_type,
window_entity_id=window_entity_id,
window_duration=window_duration,
window_start=window_start,
)
if authoritative is not None:
@ -2314,6 +2317,7 @@ async def _authoritative_floor_spend(
counter_key: str,
window_entity_type: str | None = None,
window_entity_id: str | None = None,
window_duration: str | None = None,
window_start: datetime | None = None,
) -> float | None:
marker_key: Final = f"spend_db_floor:{counter_key}"
@ -2328,10 +2332,11 @@ async def _authoritative_floor_spend(
and window_entity_id is not None
and window_start is not None
):
db_spend = await SpendCounterReseed.window_from_spend_logs(
db_spend = await SpendCounterReseed.window_from_db(
prisma_client=prisma_client,
entity_type=window_entity_type,
entity_id=window_entity_id,
window_duration=window_duration,
window_start=window_start,
)
if db_spend is None:
@ -2456,6 +2461,7 @@ async def increment_spend_counters(
counter_key=key_window_counter,
entity_type="Key",
entity_id=hashed_token,
window_duration=duration,
window_start=key_window_start,
increment=cost,
)
@ -2497,6 +2503,7 @@ async def increment_spend_counters(
counter_key=team_window_counter,
entity_type="Team",
entity_id=scope_team_id,
window_duration=duration,
window_start=team_window_start,
increment=cost,
)
@ -2749,6 +2756,7 @@ async def _init_and_increment_window_spend_counter(
counter_key: str,
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime | None,
increment: float,
):
@ -2763,6 +2771,7 @@ async def _init_and_increment_window_spend_counter(
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
window_duration=window_duration,
window_start=window_start,
)
if initialized is False:
@ -2808,6 +2817,7 @@ async def _ensure_window_spend_counter_initialized(
counter_key: str,
entity_type: str,
entity_id: str,
window_duration: str | None,
window_start: datetime,
) -> bool:
is_warm: Final = await _is_spend_counter_cache_warm(counter_key=counter_key)
@ -2820,6 +2830,7 @@ async def _ensure_window_spend_counter_initialized(
counter_key=counter_key,
entity_type=entity_type,
entity_id=entity_id,
window_duration=window_duration,
window_start=window_start,
)
if window_spend is None:

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

@ -37,6 +37,7 @@ class _BudgetCounter:
entity_id: str
source_cache_key: str | None = None
spend_log_entity_id: str | None = None
window_duration: str | None = None
window_start: datetime | None = None
@ -633,6 +634,7 @@ def _get_budget_limit_counters(
entity_type=entity_type,
entity_id=f"{entity_id}:{budget_duration}",
spend_log_entity_id=entity_id,
window_duration=str(budget_duration),
window_start=window_start,
)
)
@ -676,6 +678,7 @@ async def _reserve_counter(
counter_key=counter.counter_key,
entity_type=counter.entity_type,
entity_id=counter.spend_log_entity_id,
window_duration=counter.window_duration,
window_start=counter.window_start,
)
if initialized is False:

4
litellm/repositories/table_repositories.py
View file

@ -62,6 +62,10 @@ class SpendLogsRepository(PrismaTableRepository):
table_name = "litellm_spendlogs"
class BudgetWindowSpendRepository(PrismaTableRepository):
table_name = "litellm_budgetwindowspend"
class ClaudeCodePluginRepository(PrismaTableRepository):
table_name = "litellm_claudecodeplugintable"

250
tests/test_litellm/proxy/db/test_spend_counter_reseed.py Normal file
View file

@ -0,0 +1,250 @@
"""Window-spend reads in ``SpendCounterReseed``.
The maintained ``LiteLLM_BudgetWindowSpend`` row replaces a per-request
``LiteLLM_SpendLogs`` range scan, so these pin *when* the aggregate is still
allowed to run: only when the row is missing or belongs to an older window.
"""
from __future__ import annotations
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
import pytest
from litellm.caching.dual_cache import DualCache
from litellm.proxy.db.spend_counter_reseed import SpendCounterReseed
WINDOW_START = datetime(2026, 8, 1, tzinfo=timezone.utc)
class _FakeWindowSpendTable:
def __init__(self, row: SimpleNamespace | None, error: Exception | None = None) -> None:
self._row = row
self._error = error
self.where_clauses: list[dict] = []
async def find_unique(self, where: dict):
self.where_clauses.append(where)
if self._error is not None:
raise self._error
return self._row
class _FakeSpendLogsTable:
def __init__(self, total: float) -> None:
self._total = total
self.call_count = 0
async def group_by(self, by: list[str], where: dict, sum: dict):
self.call_count += 1
return [{by[0]: where.get(by[0]), "_sum": {"spend": self._total}}]
class _FakePrismaClient:
def __init__(
self,
row: SimpleNamespace | None = None,
spend_logs_total: float = 0.0,
error: Exception | None = None,
) -> None:
self.db = SimpleNamespace(
litellm_budgetwindowspend=_FakeWindowSpendTable(row=row, error=error),
litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total),
)
def _row(window_start: datetime, spend: float) -> SimpleNamespace:
return SimpleNamespace(window_start=window_start, spend=spend)
@pytest.mark.asyncio
async def test_window_from_table_reads_row_by_primary_key():
"""The lookup must use the table's own entity_type values ("key"), not the
"Key"/"Team" labels the counter keys and spend-log aggregates use."""
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5))
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
expected_window_start=WINDOW_START,
)
assert result == 4.5
assert prisma.db.litellm_budgetwindowspend.where_clauses == [
{
"entity_type_entity_id_window_duration": {
"entity_type": "key",
"entity_id": "tok-1",
"window_duration": "30d",
}
}
]
@pytest.mark.asyncio
async def test_window_from_table_maps_team_entity_type():
prisma = _FakePrismaClient(row=_row(WINDOW_START, 9.0))
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type="Team",
entity_id="team-1",
window_duration="1d",
expected_window_start=WINDOW_START,
)
assert result == 9.0
inner = prisma.db.litellm_budgetwindowspend.where_clauses[0]["entity_type_entity_id_window_duration"]
assert inner["entity_type"] == "team"
@pytest.mark.asyncio
async def test_window_from_table_trusts_row_newer_than_expected_window():
"""Regression: a pod holding a stale ``reset_at`` computes an expected start
behind a window another pod already rolled. Trusting only an exact match
would make it re-add the previous window's spend to the current one."""
prisma = _FakePrismaClient(row=_row(WINDOW_START + timedelta(days=1), 2.0))
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
expected_window_start=WINDOW_START,
)
assert result == 2.0
@pytest.mark.asyncio
async def test_window_from_table_rejects_row_from_previous_window():
prisma = _FakePrismaClient(row=_row(WINDOW_START - timedelta(seconds=1), 99.0))
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
expected_window_start=WINDOW_START,
)
assert result is None
@pytest.mark.asyncio
async def test_window_from_table_treats_naive_row_timestamp_as_utc():
"""The column is ``timestamp(3)``, so a driver that hands back a naive value
must still compare against the tz-aware expected start."""
prisma = _FakePrismaClient(row=_row(WINDOW_START.replace(tzinfo=None), 3.0))
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
expected_window_start=WINDOW_START,
)
assert result == 3.0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma, entity_type",
[
(_FakePrismaClient(row=None), "Key"),
(_FakePrismaClient(row=_row(WINDOW_START, 1.0)), "User"),
(_FakePrismaClient(error=RuntimeError("connection reset")), "Key"),
(None, "Key"),
],
)
async def test_window_from_table_returns_none_without_a_usable_row(prisma, entity_type):
result = await SpendCounterReseed.window_from_table(
prisma_client=prisma,
entity_type=entity_type,
entity_id="tok-1",
window_duration="30d",
expected_window_start=WINDOW_START,
)
assert result is None
@pytest.mark.asyncio
async def test_window_from_db_prefers_the_row_over_the_spend_logs_aggregate():
"""The aggregate range-scans an unindexed table; a current row must keep it
from running at all."""
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0)
result = await SpendCounterReseed.window_from_db(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
window_start=WINDOW_START,
)
assert result == 4.5
assert prisma.db.litellm_spendlogs.call_count == 0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"row",
[None, _row(WINDOW_START - timedelta(seconds=1), 99.0)],
ids=["missing_row", "previous_window_row"],
)
async def test_window_from_db_falls_back_to_spend_logs(row):
prisma = _FakePrismaClient(row=row, spend_logs_total=7.25)
result = await SpendCounterReseed.window_from_db(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
window_start=WINDOW_START,
)
assert result == 7.25
assert prisma.db.litellm_spendlogs.call_count == 1
@pytest.mark.asyncio
async def test_window_from_db_without_a_duration_skips_the_row_lookup():
"""Callers that cannot name the window (no PK) keep the pre-table behavior."""
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=7.25)
result = await SpendCounterReseed.window_from_db(
prisma_client=prisma,
entity_type="Key",
entity_id="tok-1",
window_duration=None,
window_start=WINDOW_START,
)
assert result == 7.25
assert prisma.db.litellm_budgetwindowspend.where_clauses == []
@pytest.mark.asyncio
async def test_coalesced_window_seeds_a_cold_counter_from_the_row():
prisma = _FakePrismaClient(row=_row(WINDOW_START, 4.5), spend_logs_total=100.0)
cache = DualCache()
counter_key = "spend:key:tok-1:window:30d"
result = await SpendCounterReseed.coalesced_window(
prisma_client=prisma,
spend_counter_cache=cache,
counter_key=counter_key,
entity_type="Key",
entity_id="tok-1",
window_duration="30d",
window_start=WINDOW_START,
)
assert result == 4.5
assert cache.in_memory_cache.get_cache(key=counter_key) == 4.5
assert prisma.db.litellm_spendlogs.call_count == 0

79
tests/test_litellm/proxy/proxy_server/test_spend_counters.py
View file

@ -271,6 +271,81 @@ async def test_get_current_spend_floors_window_against_spend_logs(monkeypatch):
)
def _make_window_spend_prisma(row=None, spend_logs_total=0.0):
prisma = MagicMock()
prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=row)
prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[{"api_key": "tok", "_sum": {"spend": spend_logs_total}}]
)
return prisma
@pytest.mark.asyncio
async def test_get_current_spend_floors_window_against_maintained_row(monkeypatch):
"""The floor re-check runs every few seconds per pod, so the window branch
must read the maintained row and leave the unindexed spend-logs scan alone."""
from datetime import timezone
from types import SimpleNamespace
window_start = datetime(2026, 1, 1, tzinfo=timezone.utc)
fake_prisma = _make_window_spend_prisma(
row=SimpleNamespace(window_start=window_start, spend=15.0),
spend_logs_total=100.0,
)
fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
monkeypatch.setattr(ps, "prisma_client", fake_prisma)
counter_key = "spend:key:tok:window:7d"
result = await ps.get_current_spend(
counter_key=counter_key,
fallback_spend=0.0,
max_budget=10.0,
window_entity_type="Key",
window_entity_id="tok",
window_duration="7d",
window_start=window_start,
)
assert result == 15.0
fake_prisma.db.litellm_spendlogs.group_by.assert_not_awaited()
fake_cache.redis_cache.async_set_max.assert_awaited_once_with(
key=counter_key, value=15.0
)
@pytest.mark.asyncio
async def test_get_current_spend_floors_window_against_logs_when_row_stale(monkeypatch):
"""A row left behind at a crossed window boundary must not be read as the
current window's spend; the aggregate stays the fallback."""
from datetime import timedelta, timezone
from types import SimpleNamespace
window_start = datetime(2026, 1, 8, tzinfo=timezone.utc)
fake_prisma = _make_window_spend_prisma(
row=SimpleNamespace(
window_start=window_start - timedelta(days=7), spend=999.0
),
spend_logs_total=15.0,
)
fake_cache = _make_spend_counter_cache(redis_get_value=2.0)
monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
monkeypatch.setattr(ps, "prisma_client", fake_prisma)
result = await ps.get_current_spend(
counter_key="spend:key:tok:window:7d",
fallback_spend=0.0,
max_budget=10.0,
window_entity_type="Key",
window_entity_id="tok",
window_duration="7d",
window_start=window_start,
)
assert result == 15.0
fake_prisma.db.litellm_spendlogs.group_by.assert_awaited_once()
@pytest.mark.asyncio
async def test_get_current_spend_fail_closed_rejects_when_unverifiable(monkeypatch):
"""With fail_closed_budget_enforcement on, an admit decision backed only by a
@ -895,6 +970,7 @@ async def test_init_and_increment_window_spend_counter_increments_when_initializ
counter_key="spend:key:k:window:1d",
entity_type="Key",
entity_id="k",
window_duration="1d",
window_start=datetime(2024, 1, 1),
increment=5.0,
)
@ -922,6 +998,7 @@ async def test_init_and_increment_window_spend_counter_missing_window_start_inva
counter_key="spend:key:k:window:1d",
entity_type="Key",
entity_id="k",
window_duration="1d",
window_start=None,
increment=5.0,
)
@ -1059,6 +1136,7 @@ async def test_ensure_window_spend_counter_initialized_warm_returns_true(monkeyp
counter_key="spend:key:k:window:1d",
entity_type="Key",
entity_id="k",
window_duration="1d",
window_start=datetime(2024, 1, 1),
)
@ -1091,6 +1169,7 @@ async def test_ensure_window_spend_counter_initialized_db_failure_invalid_return
counter_key="spend:key:k:window:1d",
entity_type="Key",
entity_id="k",
window_duration="1d",
window_start=datetime(2024, 1, 1),
)

8
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -7643,6 +7643,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss():
counter_cache = DualCache()
window_start = datetime.now(timezone.utc) - timedelta(hours=1)
fake_prisma = MagicMock()
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[{"api_key": "key-window", "_sum": {"spend": 2.25}}]
)
@ -7657,6 +7658,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss():
counter_key="spend:key:key-window:window:1h",
entity_type="Key",
entity_id="key-window",
window_duration="1h",
window_start=window_start,
increment=0.5,
)
@ -7766,6 +7768,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory():
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[{"api_key": "key-window-stale-local", "_sum": {"spend": 2.25}}]
)
@ -7780,6 +7783,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory():
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-stale-local",
window_duration="1h",
window_start=window_start,
increment=0.5,
)
@ -7830,6 +7834,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed()
counter_cache.redis_cache = fake_redis
fake_prisma = MagicMock()
fake_prisma.db.litellm_budgetwindowspend.find_unique = AsyncMock(return_value=None)
fake_prisma.db.litellm_spendlogs.group_by = AsyncMock(
return_value=[
{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}
@ -7846,6 +7851,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed()
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-concurrent-seed",
window_duration="1h",
window_start=window_start,
increment=0.5,
)
@ -7880,6 +7886,7 @@ async def test_window_spend_counter_skips_invalid_window_start():
counter_key="spend:key:key-invalid-window:window:not-a-duration",
entity_type="Key",
entity_id="key-invalid-window",
window_duration="not-a-duration",
window_start=None,
increment=0.5,
)
@ -7912,6 +7919,7 @@ async def test_window_spend_counter_does_not_seed_zero_when_db_unavailable():
counter_key=counter_key,
entity_type="Key",
entity_id="key-window-db-unavailable",
window_duration="1h",
window_start=datetime.now(timezone.utc) - timedelta(hours=1),
)