From 2cc0c79d074761b7f7ada2e9f5ec61b24d532865 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 14:54:13 -0700 Subject: [PATCH] fix(proxy): drop spend counters on global and organization spend resets `/global/spend/reset` and `/organization/update` with `spend` wrote only the DB `.spend` column, so with the longer counter TTL a blocked key or organization stayed blocked until its counter expired. Both now drop the affected counters (every `spend:key:*` and `spend:team:*` counter through a bounded Redis scan, the one `spend:org:` counter) so the reset takes effect on the next request. `RedisCache.async_scan_iter` now returns decoded `str` keys and stops at `count` without a mutable accumulator, and `delete_cache_keys` takes an `Iterable[str]`, so callers no longer see `bytes` keys. The audit's integration cells for the TTL (16 counter scopes and surfaces, 7 settings edges, 4 chaos cells) are added under tests/integration/spend. The leaked-reservation chaos cell asserts the shipped behavior: a reservation left by a hard-killed worker stays on the counter for the counter TTL. --- litellm/caching/redis_cache.py | 32 +- .../organization_endpoints.py | 5 + litellm/proxy/proxy_server.py | 26 + .../spend_management_endpoints.py | 4 +- tests/integration/_support/spend_counters.py | 33 ++ tests/integration/_support/upstream.py | 91 +++- .../spend/test_spend_counter_redis_ttl.py | 488 ++++++++++++++++++ .../spend/test_spend_counter_ttl_chaos.py | 296 +++++++++++ .../spend/test_spend_counter_ttl_settings.py | 228 ++++++++ tests/unit/caching/test_redis_cache.py | 19 + .../common_utils/test_reset_budget_job.py | 32 ++ .../test_organization_endpoints.py | 36 ++ .../test_spend_management_endpoints.py | 54 ++ 13 files changed, 1315 insertions(+), 29 deletions(-) create mode 100644 tests/integration/_support/spend_counters.py create mode 100644 tests/integration/spend/test_spend_counter_redis_ttl.py create mode 100644 tests/integration/spend/test_spend_counter_ttl_chaos.py create mode 100644 tests/integration/spend/test_spend_counter_ttl_settings.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3a847044e3b..108cc29b509 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -18,7 +18,7 @@ import json import logging import threading import time -from collections.abc import Awaitable, Callable, Iterator, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -498,6 +498,14 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1) +async def _first_keys(scan: AsyncIterator[object], count: int) -> AsyncIterator[str]: + taken: Final = itertools.count(1) + async for key in scan: + yield key.decode() if isinstance(key, bytes) else str(key) + if next(taken) >= count: + return + + class RedisCircuitBreakerOpenError(Exception): pass @@ -1007,10 +1015,9 @@ class RedisCache(BaseCache): return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) @_redis_circuit_breaker_guard - async def async_scan_iter(self, pattern: str, count: int = 100) -> list: + async def async_scan_iter(self, pattern: str, count: int = 100) -> list[str]: start_time: Final = time.time() try: - keys: Final = [] _redis_client: Final = self.init_async_client() if not hasattr(_redis_client, "scan_iter"): verbose_logger.debug( @@ -1018,11 +1025,13 @@ class RedisCache(BaseCache): ) return [] - pattern = self.check_and_fix_namespace(key=pattern) - async for key in _redis_client.scan_iter(match=pattern + "*", count=count): - keys.append(key) - if len(keys) >= count: - break + namespaced_pattern: Final = self.check_and_fix_namespace(key=pattern) + keys: Final = [ + key + async for key in _first_keys( + _redis_client.scan_iter(match=namespaced_pattern + "*", count=count), count + ) + ] ## LOGGING ## end_time = time.time() @@ -1904,11 +1913,10 @@ class RedisCache(BaseCache): raise e @_redis_circuit_breaker_guard - async def delete_cache_keys(self, keys): + async def delete_cache_keys(self, keys: Iterable[str]) -> None: _redis_client: Final = self._async_commands() - keys = [self.check_and_fix_namespace(key=key) for key in keys] - # keys is a list, unpack it so it gets passed as individual elements to delete - await _redis_client.delete(*keys) + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in keys] + await _redis_client.delete(*namespaced_keys) def client_list(self) -> list: client_list: Final[list] = self.redis_client.client_list() diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index b8ffec639c0..8928832626e 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -816,6 +816,11 @@ async def update_organization( include={"members": True, "teams": True, "litellm_budget_table": True}, ) + if data.spend is not None: + from litellm.proxy.proxy_server import invalidate_spend_counter + + await invalidate_spend_counter(counter_key=f"spend:org:{data.organization_id}") + return response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6c09cbeccf4..2e63584f60e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -51,6 +51,7 @@ import websockets import websockets.exceptions from pydantic import BaseModel, ConfigDict, Json, JsonValue, TypeAdapter, ValidationError, with_config from pydantic.fields import FieldInfo, PydanticUndefined +from redis.exceptions import RedisError from typing_extensions import NotRequired, ReadOnly, TypedDict, assert_never from litellm._uuid import uuid @@ -3948,6 +3949,31 @@ async def invalidate_spend_counter(counter_key: str) -> None: _invalidate_spend_counter: Final = invalidate_spend_counter +_IN_MEMORY_COUNTER_KEYS: Final = TypeAdapter(tuple[str, ...]) + + +_SPEND_COUNTER_STORE_FAILURES: Final = (RedisError, RedisCircuitBreakerOpenError, OSError, ValidationError) + + +@with_service_target(SPEND_COUNTERS_TARGET) +async def invalidate_spend_counters_with_prefix(prefix: str) -> None: + for counter_key in _IN_MEMORY_COUNTER_KEYS.validate_python(tuple(spend_counter_cache.in_memory_cache.cache_dict)): + if counter_key.startswith(prefix): + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) + redis_cache: Final = spend_counter_cache.redis_cache + if redis_cache is None: + return + try: + while scanned := await redis_cache.async_scan_iter(pattern=prefix, count=500): + await redis_cache.delete_cache_keys(scanned) + except _SPEND_COUNTER_STORE_FAILURES: + verbose_proxy_logger.warning( + "Unable to delete the spend counters %s* after the spend reset; budget checks use the stale counters until " + "they expire", + prefix, + exc_info=True, + ) + async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: if _defer_spend_counter_increments(pending): diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 5482661c3de..8384e5e9177 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -3684,7 +3684,7 @@ async def global_spend_reset(): 3. LiteLLM_TeamTable spend will be set = 0 """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import invalidate_spend_counters_with_prefix, prisma_client if prisma_client is None: raise ProxyException( @@ -3696,6 +3696,8 @@ async def global_spend_reset(): await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + await invalidate_spend_counters_with_prefix(prefix="spend:key:") + await invalidate_spend_counters_with_prefix(prefix="spend:team:") return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/tests/integration/_support/spend_counters.py b/tests/integration/_support/spend_counters.py new file mode 100644 index 00000000000..a8c8b738df5 --- /dev/null +++ b/tests/integration/_support/spend_counters.py @@ -0,0 +1,33 @@ +import os +from collections.abc import Iterator +from contextlib import contextmanager +from hashlib import sha256 +from typing import Final + +from redis import Redis + +SPEND_COUNTER_TTL_SECONDS: Final = 3600 +TTL_SLACK_SECONDS: Final = 600 + + +@contextmanager +def rig_cache() -> Iterator[Redis]: + with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache: + yield cache + + +def key_counter(key: str) -> str: + return f"spend:key:{sha256(key.encode()).hexdigest()}" + + +def counter_ttl(cache: Redis, counter_key: str) -> int: + return int(str(cache.ttl(counter_key))) + + +def counter_value(cache: Redis, counter_key: str) -> float | None: + raw: Final = cache.get(counter_key) + return None if raw is None else float(str(raw, "utf-8") if isinstance(raw, bytes) else str(raw)) + + +def outlives(ttl: int, expected: int = SPEND_COUNTER_TTL_SECONDS) -> bool: + return expected - TTL_SLACK_SECONDS < ttl <= expected diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 6b36d56a757..831a4543a55 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -6,6 +6,7 @@ import base64 import json import os import struct +import time import uuid import zlib from collections import deque @@ -60,6 +61,46 @@ INTERNAL_FIELDS: Final = frozenset( ) +RESPONSES_USAGE: Final = {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40} + + +def responses_body(model: str) -> dict[str, JsonValue]: + return { + "id": f"resp_{uuid.uuid4().hex[:24]}", + "object": "response", + "created_at": int(time.time()), + "status": "completed", + "model": model, + "output": [ + { + "id": f"msg_{uuid.uuid4().hex[:24]}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Hello.", "annotations": []}], + } + ], + "usage": dict(RESPONSES_USAGE), + } + + +async def responses_stream(completed: Mapping[str, JsonValue]) -> AsyncIterator[str]: + item: Final = cast(dict[str, JsonValue], cast(list[JsonValue], completed["output"])[0]) + events: Final = ( + ("response.created", {"response": {**completed, "status": "in_progress", "output": [], "usage": None}}), + ("response.output_item.added", {"output_index": 0, "item": {**item, "status": "in_progress", "content": []}}), + ( + "response.output_text.delta", + {"output_index": 0, "content_index": 0, "item_id": item["id"], "delta": "Hello."}, + ), + ("response.output_item.done", {"output_index": 0, "item": item}), + ("response.completed", {"response": dict(completed)}), + ) + for sequence_number, (event_type, payload) in enumerate(events): + data: Final = json.dumps({"type": event_type, "sequence_number": sequence_number, **payload}) + yield f"event: {event_type}\ndata: {data}\n\n" + + def error_type(status: int) -> str: if status == 429: return "rate_limit_error" @@ -183,24 +224,41 @@ class Provider: for message in messages ): return JSONResponse({"error": {"message": "Invalid selected message contract"}}, status_code=400) - script: Final = self.scripts.get(str(body["model"])) - if script is not None: - if not script: - return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) - status: Final = script.popleft() - if status != 200: - return JSONResponse( - { - "error": { - "message": "Controlled provider failure", - "type": error_type(status), - "code": str(status), - } - }, - status_code=status, - ) + failure: Final = self._scripted_failure(str(body["model"])) + if failure is not None: + return failure return await chat_completions(request) + def _scripted_failure(self, model: str) -> Response | None: + script: Final = self.scripts.get(model) + if script is None: + return None + if not script: + return JSONResponse({"error": {"message": "Script exhausted", "type": "api_error"}}, status_code=500) + status: Final = script.popleft() + if status == 200: + return None + return JSONResponse( + {"error": {"message": "Controlled provider failure", "type": error_type(status), "code": str(status)}}, + status_code=status, + ) + + async def responses(self, request: Request) -> Response: + body: Final = JSON_OBJECT.validate_json(await request.body()) + self.observations.put( + Observation(request.url.path, request.headers.get("authorization", ""), body, request.method) + ) + model: Final = body.get("model") + if not isinstance(model, str) or "input" not in body: + return JSONResponse({"error": {"message": "model and input are required"}}, status_code=400) + failure: Final = self._scripted_failure(model) + if failure is not None: + return failure + completed: Final = responses_body(model) + if body.get("stream") is True: + return StreamingResponse(responses_stream(completed), media_type="text/event-stream") + return JSONResponse(completed) + async def vector_store_search(self, request: Request) -> Response: body: Final = JSON_OBJECT.validate_json(await request.body()) self.observations.put( @@ -560,6 +618,7 @@ class Provider: Route("/_cost_map", self.cost_map, methods=["GET"]), Route("/_oauth/token", self.oauth_token, methods=["POST"]), Route("/v1/chat/completions", self.chat, methods=["POST"]), + Route("/v1/responses", self.responses, methods=["POST"]), Route("/v1/completions", completions, methods=["POST"]), Route("/v1/embeddings", embeddings, methods=["POST"]), Route("/v1/moderations", moderations, methods=["POST"]), diff --git a/tests/integration/spend/test_spend_counter_redis_ttl.py b/tests/integration/spend/test_spend_counter_redis_ttl.py new file mode 100644 index 00000000000..ac4e338993b --- /dev/null +++ b/tests/integration/spend/test_spend_counter_redis_ttl.py @@ -0,0 +1,488 @@ +import asyncio +import json +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from hashlib import sha256 +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} + + +@dataclass(frozen=True, slots=True) +class Entities: + model: str + key: str + team: str + user: str + organization: str + end_user: str + tag: str + + def counters(self) -> Mapping[str, str]: + return { + "key": key_counter(self.key), + "team": f"spend:team:{self.team}", + "team_member": f"spend:team_member:{self.user}:{self.team}", + "user": f"spend:user:{self.user}", + "end_user": f"spend:end_user:{self.end_user}", + "tag": f"spend:tag:{self.tag}", + "org": f"spend:org:{self.organization}", + } + + +def _entities(gateway: Gateway, scenario: Scenario) -> Entities: + model: Final = scenario.model(**RATES) + models: Final = [model] + organization: Final = scenario.organization(max_budget=50.0) + team: Final = scenario.team(models=models, max_budget=20.0, organization_id=organization) + user: Final = scenario.user(max_budget=20.0) + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 20.0}, + ) + assert added.status_code == 200, added.text + tag: Final = f"integration-tag-{uuid.uuid4().hex}" + key: Final = scenario.key(team_id=team, user_id=user, models=models, max_budget=5.0, tags=[tag]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + return Entities(model, key, team, user, organization, end_user, tag) + + +def _base(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _chat_openai_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with openai.OpenAI(base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60) as client: + response: Final = client.chat.completions.create( + model=entities.model, messages=[{"role": "user", "content": text}], user=entities.end_user + ) + assert response.usage is not None and response.usage.total_tokens == 40, response + return response.id + + +def _chat_openai_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with openai.AsyncOpenAI( + base_url=f"{_base(gateway)}/v1", api_key=entities.key, max_retries=0, timeout=60 + ) as client: + stream: Final = await client.chat.completions.create( + model=entities.model, + messages=[{"role": "user", "content": text}], + user=entities.end_user, + stream=True, + stream_options={"include_usage": True}, + ) + ids: Final = [chunk.id async for chunk in stream] + assert len(set(ids)) == 1, ids + return ids[0] + + return asyncio.run(run()) + + +def _chat_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": entities.model, + "messages": [{"role": "user", "content": text}], + "user": entities.end_user, + "stream": True, + "stream_options": {"include_usage": True}, + }, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + frames: Final = [line.removeprefix("data: ") for line in response.iter_lines() if line.startswith("data: ")] + ids: Final = {json.loads(frame)["id"] for frame in frames if frame != "[DONE]"} + assert len(ids) == 1, frames + return ids.pop() + + +def _messages_anthropic_sync(gateway: Gateway, entities: Entities, text: str) -> str: + with anthropic.Anthropic(base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60) as client: + message: Final = client.messages.create( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) + assert message.usage.output_tokens > 0, message + return message.id + + +def _messages_anthropic_async_stream(gateway: Gateway, entities: Entities, text: str) -> str: + async def run() -> str: + async with anthropic.AsyncAnthropic( + base_url=_base(gateway), api_key=entities.key, max_retries=0, timeout=60 + ) as client: + async with client.messages.stream( + model=entities.model, + max_tokens=32, + messages=[{"role": "user", "content": text}], + metadata={"user_id": entities.end_user}, + ) as stream: + async for _ in stream.text_stream: + pass + final: Final = await stream.get_final_message() + assert final.usage.output_tokens > 0, final + return final.id + + return asyncio.run(run()) + + +def _messages_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": entities.model, + "max_tokens": 32, + "messages": [{"role": "user", "content": text}], + "metadata": {"user_id": entities.end_user}, + }, + key=entities.key, + headers=ANTHROPIC_HEADERS, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx(gateway: Gateway, entities: Entities, text: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": entities.model, "input": text, "user": entities.end_user}, + key=entities.key, + ) + assert response.status_code == 200, response.text + return str(response.json()["id"]) + + +def _responses_httpx_stream(gateway: Gateway, entities: Entities, text: str) -> str: + with gateway.client.stream( + "POST", + "/v1/responses", + json={"model": entities.model, "input": text, "user": entities.end_user, "stream": True}, + headers={"Authorization": f"Bearer {entities.key}"}, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = [line for line in response.iter_lines() if line.startswith("data: ")] + events: Final = [json.loads(line.removeprefix("data: ")) for line in lines if line != "data: [DONE]"] + completed: Final = [event for event in events if event.get("type") == "response.completed"] + assert len(completed) == 1, [event.get("type") for event in events] + return str(completed[0]["response"]["id"]) + + +@dataclass(frozen=True, slots=True) +class Caller: + call: Callable[[Gateway, Entities, str], str] + upstream_path: str + exact_cost: float | None + + +CALLERS: Final[Mapping[str, Caller]] = { + "chat-openai-sync": Caller(_chat_openai_sync, "/v1/chat/completions", ONE_CALL_COST), + "chat-openai-async-stream": Caller(_chat_openai_async_stream, "/v1/chat/completions", ONE_CALL_COST), + "chat-httpx-stream": Caller(_chat_httpx_stream, "/v1/chat/completions", ONE_CALL_COST), + "messages-anthropic-sync": Caller(_messages_anthropic_sync, "/v1/responses", ONE_CALL_COST), + "messages-anthropic-async-stream": Caller(_messages_anthropic_async_stream, "/v1/responses", ONE_CALL_COST), + "messages-httpx": Caller(_messages_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx": Caller(_responses_httpx, "/v1/responses", ONE_CALL_COST), + "responses-httpx-stream": Caller(_responses_httpx_stream, "/v1/responses", ONE_CALL_COST), +} + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (sha256(key.encode()).hexdigest(),), + ) + + +def _one_spend_row(key: str) -> dict[str, JsonValue]: + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) + return rows[0] + + +def _observed(cache: Redis, counters: Mapping[str, str]) -> Mapping[str, tuple[float | None, int]]: + return {name: (counter_value(cache, counter), counter_ttl(cache, counter)) for name, counter in counters.items()} + + +def _assert_every_counter_outlives( + cache: Redis, counters: Mapping[str, str], spend: float +) -> Mapping[str, tuple[float | None, int]]: + observed: Final = eventually( + lambda: _observed(cache, counters), + lambda states: all( + value is not None and value == pytest.approx(spend, abs=1e-9) for value, _ in states.values() + ), + seconds=20, + ) + for name, (_, ttl) in observed.items(): + assert outlives(ttl), (name, observed) + return observed + + +@pytest.mark.parametrize("caller", sorted(CALLERS)) +def test_every_counter_scope_outlives_the_batch_write(gateway: Gateway, caller: str) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + text: Final = f"counter ttl {caller} {uuid.uuid4().hex}" + upstream.get("/__observations").raise_for_status() + response_id: Final = CALLERS[caller].call(gateway, entities, text) + received: Final = upstream.get("/__observations").json()["requests"] + paths: Final = [request["path"] for request in received if text in json.dumps(request["body"])] + assert len(paths) == 1 and paths[0].endswith(CALLERS[caller].upstream_path), received + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + assert spend == pytest.approx(CALLERS[caller].exact_cost or spend) and spend > 0, row + assert response_id, row + _assert_every_counter_outlives(cache, entities.counters(), spend) + + +def test_a_response_cache_hit_keeps_the_counter_and_its_ttl(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + rig_cache() as cache, + ): + entities: Final = _entities(gateway, scenario) + body: Final = { + "model": entities.model, + "messages": [{"role": "user", "content": f"cache twin {uuid.uuid4().hex}"}], + "user": entities.end_user, + } + first: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert first.status_code == 200, first.text + row: Final = _one_spend_row(entities.key) + spend: Final = float(str(row["spend"])) + before: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + upstream.get("/__observations").raise_for_status() + second: Final = gateway.request("POST", "/v1/chat/completions", body, key=entities.key) + assert second.status_code == 200, second.text + assert second.json()["id"] == first.json()["id"], (first.text, second.text) + assert upstream.get("/__observations").json()["requests"] == [] + rows: Final = eventually(lambda: _spend_rows(entities.key), lambda values: len(values) == 2, seconds=70) + assert str(rows[1]["cache_hit"]) == "True" and float(str(rows[1]["spend"])) == 0.0, rows + after: Final = _assert_every_counter_outlives(cache, entities.counters(), spend) + assert {name: value for name, (value, _) in after.items()} == pytest.approx( + {name: value for name, (value, _) in before.items()}, abs=1e-9 + ), (before, after) + + +def _exhausted_key(gateway: Gateway, scenario: Scenario) -> tuple[str, str]: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=ONE_CALL_COST) + served: Final = gateway.chat(model, key=key, text=f"exhaust {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + return model, key + + +def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def _assert_denied(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert response.json()["error"]["type"] == "budget_exceeded", response.text + + +def test_a_stale_low_counter_is_repaired_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key( + models=[model], budget_limits=[{"max_budget": ONE_CALL_COST, "budget_duration": "1d"}] + ) + counter: Final = f"{key_counter(key)}:window:1d" + assert gateway.chat(model, key=key, text=f"exhaust window {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + assert counter_value(cache, counter) == pytest.approx(ONE_CALL_COST) + + served: Final[list[int]] = [] + + def stale_then_request() -> int: + assert cache.set(counter, "0.01") + code: Final = _chat(gateway, model, key).status_code + served.append(code) + return code + + assert eventually(stale_then_request, lambda code: code == 422, seconds=30) == 422 + assert set(served) == {200, 422}, served + repaired: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + floor: Final = repaired[0] + assert floor is not None and ONE_CALL_COST - 1e-9 <= floor <= served.count(200) * ONE_CALL_COST + 1e-9, ( + repaired, + served, + ) + assert outlives(repaired[1]), repaired + + +def test_a_lost_counter_is_reseeded_from_the_db_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + assert cache.delete(counter) == 1 + _assert_denied(_chat(gateway, model, key)) + reseeded: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] >= ONE_CALL_COST and outlives(state[1]), + seconds=15, + ) + assert reseeded[0] == pytest.approx(ONE_CALL_COST), reseeded + + +def test_key_update_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before update {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + updated: Final = gateway.request("POST", "/key/update", {"key": key, "spend": 1.25}) + assert updated.status_code == 200, updated.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (1.25, True), updated.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(1.25 + ONE_CALL_COST), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) + + +def test_key_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model, key = _exhausted_key(gateway, scenario) + counter: Final = key_counter(key) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["key_hash"] == sha256(key.encode()).hexdigest(), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_team_member_reset_spend_sets_the_counter_with_the_long_ttl(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + team: Final = scenario.team(models=[model]) + user: Final = scenario.user() + added: Final = gateway.request( + "POST", + "/team/member_add", + {"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": ONE_CALL_COST}, + ) + assert added.status_code == 200, added.text + key: Final = scenario.key(team_id=team, user_id=user, models=[model]) + counter: Final = f"spend:team_member:{user}:{team}" + assert gateway.chat(model, key=key, text=f"member {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', (team, user) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + _assert_denied(_chat(gateway, model, key)) + reset: Final = gateway.request("POST", f"/team/{team}/member/{user}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert reset.json()["previous_spend"] == pytest.approx(ONE_CALL_COST), reset.text + assert (counter_value(cache, counter), outlives(counter_ttl(cache, counter))) == (0.0, True), reset.text + served: Final = _chat(gateway, model, key) + assert served.status_code == 200, served.text + + +def test_malformed_spend_writes_are_rejected_and_leave_the_counter_alone(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + bystander: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert gateway.chat(model, key=key, text=f"before bad writes {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40 + _one_spend_row(key) + before: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert before[0] == pytest.approx(ONE_CALL_COST) and outlives(before[1]), before + rejected: Final = tuple( + gateway.request(*attempt) + for attempt in ( + ("POST", "/key/update", {"key": key, "spend": "a" * 5000}), + ("POST", "/key/update", {"key": key, "spend": [1.0]}), + ("POST", "/key/update", {"key": key, "spend": {"spend": 1.0}}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": "abc"}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": True}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": [0.0]}), + ("POST", f"/key/{key}/reset_spend", {"reset_to": ONE_CALL_COST * 10}), + ) + ) + for response in rejected: + assert response.status_code in (400, 422), response.text + assert "error" in response.json() or "detail" in response.json(), response.text + assert (counter_value(cache, counter), counter_ttl(cache, counter) <= before[1]) == (before[0], True) + assert outlives(counter_ttl(cache, counter)) + served: Final = _chat(gateway, model, bystander) + assert served.status_code == 200, served.text + + +def _landed_chat(gateway: Gateway, model: str, key: str, index: int) -> str: + served: Final = gateway.chat(model, key=key, text=f"repeat {index} {uuid.uuid4().hex}") + eventually(lambda: _spend_rows(key), lambda rows: len(rows) == index + 1, seconds=70) + return str(served["id"]) + + +def test_repeated_requests_log_once_each_and_the_counter_matches_their_sum(gateway: Gateway) -> None: + with gateway.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + ids: Final = tuple(_landed_chat(gateway, model, key, index) for index in range(3)) + assert len(set(ids)) == 3, ids + rows: Final = _spend_rows(key) + assert len(rows) == 3, rows + assert sorted(str(row["request_id"]) for row in rows) == sorted(ids), rows + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(3 * ONE_CALL_COST), rows + eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value == pytest.approx(total), + seconds=15, + ) + assert outlives(counter_ttl(cache, counter)) diff --git a/tests/integration/spend/test_spend_counter_ttl_chaos.py b/tests/integration/spend/test_spend_counter_ttl_chaos.py new file mode 100644 index 00000000000..b22f268b595 --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_chaos.py @@ -0,0 +1,296 @@ +import os +import signal +import uuid +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, group_members, owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.spend_counters import ( + counter_ttl, + counter_value, + key_counter, + outlives, + rig_cache, +) +from pydantic import JsonValue +from redis import Redis + +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ANTHROPIC_HEADERS: Final = {"anthropic-version": "2023-06-01"} +BURST: Final = 24 +UNBOUNDED_BUDGET: Final = 100000.0 +RESERVATION_ESTIMATE_CEILING: Final = 40.0 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +LEAK_KEY_BUDGET: Final = 20.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +RESPAWN_SECONDS: Final = graceful_stop_seconds() + 60 + + +def _chat(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"burst chat {index} {uuid.uuid4().hex}"}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + + +def _messages(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 32, + "messages": [{"role": "user", "content": f"burst messages {index} {uuid.uuid4().hex}"}], + "stream": stream, + }, + key=key, + headers=ANTHROPIC_HEADERS, + ) + + +def _responses(gateway: Gateway, model: str, key: str, index: int, stream: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": f"burst responses {index} {uuid.uuid4().hex}", "stream": stream}, + key=key, + ) + + +SURFACES: Final[tuple[Callable[[Gateway, str, str, int, bool], httpx.Response], ...]] = ( + _chat, + _messages, + _responses, +) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, spend, status FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + + +def _successes(key: str, count: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda values: len(values) >= count, + seconds=90, + ) + assert len(rows) == count, rows + assert len({str(row["request_id"]) for row in rows}) == count, rows + return rows + + +def _counter_matches(cache: Redis, counter: str, total: float) -> tuple[float | None, int]: + state: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda values: values[0] is not None and values[0] == pytest.approx(total, abs=1e-9), + seconds=20, + ) + assert outlives(state[1]), state + return state + + +def test_a_burst_with_upstream_failures_lands_each_success_once_with_a_long_ttl(gateway: Gateway) -> None: + with ( + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + gateway.scenario() as scenario, + rig_cache() as cache, + ): + provider_model: Final = f"burst-{uuid.uuid4().hex}" + model: Final = scenario.model(model=f"openai/{provider_model}", **RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + statuses: Final = [500] + [200, 200, 200] * (BURST // 4 + 2) + + def remove_script() -> None: + response: Final = upstream.delete(f"/__scripts/{provider_model}") + assert response.status_code in (200, 404), response.text + + scenario.cleanups.callback(remove_script) + configured: Final = upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}) + assert configured.status_code == 200, configured.text + + def attempt(index: int) -> httpx.Response: + stream: Final = index % 2 == 1 + return SURFACES[index % len(SURFACES)](gateway, model, key, index, stream) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + responses: Final = tuple(pool.map(attempt, range(BURST))) + codes: Final = [response.status_code for response in responses] + assert all(code in (200, 500) for code in codes), codes + succeeded: Final = codes.count(200) + assert 0 < succeeded < BURST, codes + rows: Final = _successes(key, succeeded) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx(succeeded * 0.06), rows + _counter_matches(cache, key_counter(key), total) + + +def _workers(owned: OwnedProxy) -> tuple[int, ...]: + return tuple( + member.pid + for member in group_members(owned.process.pid) + if member.pid != owned.process.pid + and member.is_running() + and member.status() != psutil.STATUS_ZOMBIE + and any("spawn_main" in part for part in member.cmdline()) + ) + + +def _attempt(gateway: Gateway, model: str, key: str, index: int) -> int: + try: + return _chat(gateway, model, key, index, False).status_code + except httpx.HTTPError: + return 0 + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_killed_worker_does_not_shorten_or_lose_the_counter(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=12) as pool: + first: Final = pool.submit(_attempt, owned.gateway, model, key, 0) + assert first.result(timeout=60) == 200 + in_flight: Final = [pool.submit(_attempt, owned.gateway, model, key, index) for index in range(1, 12)] + os.kill(before[0], signal.SIGKILL) + survivor: Final = _attempt(owned.gateway, model, key, 99) + codes: Final = [future.result(timeout=120) for future in in_flight] + assert survivor == 200, survivor + assert all(code in (0, 200) for code in codes), codes + replaced: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and before[0] not in pids and before[1] in pids, + seconds=RESPAWN_SECONDS, + ) + assert replaced + after: Final = [_attempt(owned.gateway, model, key, index) for index in range(12, 18)] + assert after == [200] * 6, after + answered: Final = 1 + 1 + codes.count(200) + 6 + landed: Final = eventually( + lambda: [row for row in _spend_rows(key) if str(row["status"]) == "success"], + lambda rows: len(rows) >= 1 + 1 + 6, + seconds=90, + ) + assert len({str(row["request_id"]) for row in landed}) == len(landed) <= answered, landed + ceiling: Final = answered * 0.06 + len(in_flight) * RESERVATION_ESTIMATE_CEILING + state: Final = eventually( + lambda: (counter_value(cache, key_counter(key)), counter_ttl(cache, key_counter(key))), + lambda values: values[0] is not None and 8 * 0.06 - 1e-9 <= values[0] <= ceiling, + seconds=20, + ) + assert outlives(state[1]), state + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reservation_leaked_by_a_killed_worker_stays_for_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy_process(gateway, tmp_path, {}, workers=2) as owned, + owned.gateway.scenario() as scenario, + rig_cache() as cache, + ): + slow: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[slow], max_budget=LEAK_KEY_BUDGET) + counter: Final = key_counter(key) + before: Final = _workers(owned) + assert len(before) == 2, before + with ThreadPoolExecutor(max_workers=1) as pool: + in_flight: Final = pool.submit(_attempt, owned.gateway, slow, key, 0) + reserved: Final = eventually( + lambda: counter_value(cache, counter), + lambda value: value is not None and value >= LEAK_KEY_BUDGET - 1e-9, + seconds=SLOW_UPSTREAM_SECONDS, + ) + for pid in before: + os.kill(pid, signal.SIGKILL) + interrupted: Final = in_flight.result(timeout=60) + assert interrupted != 200, interrupted + leaked: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert leaked[0] == pytest.approx(reserved), (reserved, leaked) + respawned: Final = eventually( + lambda: _workers(owned), + lambda pids: len(pids) == 2 and not set(pids) & set(before), + seconds=RESPAWN_SECONDS, + ) + assert respawned + blocked: Final = eventually( + lambda: _attempt(owned.gateway, slow, key, 1), lambda code: code == 422, seconds=RESPAWN_SECONDS + ) + assert blocked == 422, (blocked, leaked) + held: Final = (counter_value(cache, counter), counter_ttl(cache, counter)) + assert held[0] == pytest.approx(reserved), (reserved, held) + assert outlives(held[1]), held + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_redis_outage_mid_burst_keeps_serving_and_rebuilds_the_counter_long(gateway: Gateway, tmp_path: Path) -> None: + with owned_redis(tmp_path) as redis: + overrides: Final[Mapping[str, str]] = { + "REDIS_HOST": redis.host, + "REDIS_PORT": str(redis.port), + "REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT": "2", + } + with ( + owned_proxy_process(gateway, tmp_path, overrides) as owned, + owned.gateway.scenario() as scenario, + Redis(host=redis.host, port=redis.port) as cache, + ): + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=UNBOUNDED_BUDGET) + counter: Final = key_counter(key) + + def burst(start: int) -> list[int]: + with ThreadPoolExecutor(max_workers=6) as pool: + return list( + pool.map(lambda index: _attempt(owned.gateway, model, key, index), range(start, start + 6)) + ) + + assert burst(0) == [200] * 6 + _counter_matches(cache, counter, 6 * 0.06) + redis.stop() + during: Final = burst(6) + assert during == [200] * 6, during + redis.start() + assert cache.ping() + probes: Final[list[int]] = [] + + def probe_until_the_counter_is_back() -> float | None: + probes.append(_attempt(owned.gateway, model, key, 100 + len(probes))) + return counter_value(cache, counter) + + eventually(probe_until_the_counter_is_back, lambda value: value is not None, seconds=30) + assert set(probes) == {200}, probes + after: Final = burst(12) + assert after == [200] * 6, after + rows: Final = _successes(key, 18 + len(probes)) + total: Final = sum(float(str(row["spend"])) for row in rows) + assert total == pytest.approx((18 + len(probes)) * 0.06), rows + rebuilt: Final = eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and 7 * 0.06 - 1e-9 <= state[0] <= total + 1e-9, + seconds=20, + ) + assert outlives(rebuilt[1]), rebuilt diff --git a/tests/integration/spend/test_spend_counter_ttl_settings.py b/tests/integration/spend/test_spend_counter_ttl_settings.py new file mode 100644 index 00000000000..ec9d725bf5c --- /dev/null +++ b/tests/integration/spend/test_spend_counter_ttl_settings.py @@ -0,0 +1,228 @@ +import json +import os +import time +import uuid +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import psycopg +import pytest +import yaml +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.spend_counters import counter_ttl, counter_value, key_counter, outlives, rig_cache +from pydantic import JsonValue + +ROOT: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3]) +RATES: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} +ONE_CALL_COST: Final = 0.06 +SLOW_UPSTREAM_SECONDS: Final = 3.0 +SHORT_REDIS_TTL: Final = 4 +LONG_REDIS_TTL: Final = 86400 +ENV_REDIS_TTL: Final = "7200" +LONG_BATCH_WRITE_SECONDS: Final = 900 +STRICT_KEY_BUDGET: Final = 50.0 +CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +FAST_RESET_TICK: Final = MappingProxyType( + {"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "2"} +) + + +def _config( + directory: Path, + *, + litellm_settings: Mapping[str, JsonValue] = MappingProxyType({}), + general_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + shared: Final = object_value(yaml.safe_load((ROOT / "tests/integration/proxy_config.yaml").read_text())) + configured: Final = { + **shared, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + "general_settings": {**object_value(shared["general_settings"]), **general_settings}, + } + path: Final = directory / f"proxy_config_{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(configured, sort_keys=False)) + return path + + +def _chat(gateway: Gateway, model: str, key: str) -> str: + served: Final = gateway.chat(model, key=key, text=f"settings {uuid.uuid4().hex}") + assert served["usage"]["total_tokens"] == 40, served + return str(served["id"]) + + +def _counter_state(counter: str) -> tuple[float | None, int]: + with rig_cache() as cache: + return counter_value(cache, counter), counter_ttl(cache, counter) + + +def _settled(counter: str, spend: float, expected_ttl: int) -> tuple[float | None, int]: + return eventually( + lambda: _counter_state(counter), + lambda state: state[0] is not None and state[0] == pytest.approx(spend) and outlives(state[1], expected_ttl), + seconds=15, + ) + + +def _spend_rows(key: str) -> list[dict[str, JsonValue]]: + return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_null_default_redis_ttl_falls_back_to_the_hour_floor(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": None}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_short_default_redis_ttl_never_shortens_the_counter_and_the_lease_refresh_keeps_it_long( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": SHORT_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(model="openai/slow-endpoint", **RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + with ThreadPoolExecutor(max_workers=1) as pool: + slow: Final = pool.submit(lambda: (_chat(owned, model, key), time.monotonic())[1]) + eventually(lambda: cache.exists(counter), lambda present: present == 1, seconds=SLOW_UPSTREAM_SECONDS) + assert cache.expire(counter, 1000) + refreshed_at: Final = eventually( + lambda: (time.monotonic(), counter_ttl(cache, counter)), + lambda state: outlives(state[1]), + seconds=SLOW_UPSTREAM_SECONDS + 10, + )[0] + finished_at: Final = slow.result(timeout=60) + assert refreshed_at < finished_at, (refreshed_at, finished_at) + assert _settled(counter, ONE_CALL_COST, 3600) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_default_redis_ttl_wins_on_every_write_path(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": LONG_REDIS_TTL}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, LONG_REDIS_TTL) + updated: Final = owned.request("POST", "/key/update", {"key": key, "spend": 2.5}) + assert updated.status_code == 200, updated.text + assert _settled(counter, 2.5, LONG_REDIS_TTL) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_default_redis_ttl_given_as_an_environment_string_is_read_as_seconds( + gateway: Gateway, tmp_path: Path +) -> None: + config: Final = _config(tmp_path, litellm_settings={"default_redis_ttl": "os.environ/SPEND_COUNTER_REDIS_TTL"}) + overrides: Final = {"SPEND_COUNTER_REDIS_TTL": ENV_REDIS_TTL} + with owned_proxy(gateway, tmp_path, overrides, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + assert _settled(counter, ONE_CALL_COST, int(ENV_REDIS_TTL)) + reset: Final = owned.request("POST", f"/key/{key}/reset_spend", {"reset_to": 0.0}) + assert reset.status_code == 200, reset.text + assert _settled(counter, 0.0, int(ENV_REDIS_TTL)) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_long_batch_write_interval_lifts_the_counter_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"proxy_batch_write_at": LONG_BATCH_WRITE_SECONDS}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=5.0) + assert _chat(owned, model, key) + assert _settled(key_counter(key), ONE_CALL_COST, 10 * LONG_BATCH_WRITE_SECONDS) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_strict_enforcement_repairs_a_stale_low_key_counter_with_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + config: Final = _config(tmp_path, general_settings={"fail_closed_budget_enforcement": True}) + with owned_proxy(gateway, tmp_path, {}, config=config) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], max_budget=STRICT_KEY_BUDGET) + counter: Final = key_counter(key) + assert _chat(owned, model, key) + eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= ONE_CALL_COST, + seconds=70, + ) + attempts: Final[list[str]] = [] + + def stale_then_request() -> tuple[float | None, int]: + assert cache.set(counter, "0.01") + attempts.append(_chat(owned, model, key)) + return eventually( + lambda: (counter_value(cache, counter), counter_ttl(cache, counter)), + lambda state: state[0] is not None and state[0] < STRICT_KEY_BUDGET / 2, + seconds=15, + ) + + repaired: Final = eventually( + stale_then_request, + lambda state: state[0] is not None and state[0] > 0.01 + ONE_CALL_COST + 1e-6, + seconds=30, + ) + lifted: Final = repaired[0] + assert ( + lifted is not None and 2 * ONE_CALL_COST - 1e-9 <= lifted <= (len(attempts) + 1) * ONE_CALL_COST + 1e-9 + ), ( + repaired, + attempts, + ) + assert outlives(repaired[1]), repaired + + +def _window_reset_at(key: str) -> datetime: + rows: Final = read_rows( + 'SELECT budget_limits FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert len(rows) == 1, rows + raw: Final = rows[0]["budget_limits"] + windows: Final = json.loads(raw) if isinstance(raw, str) else raw + assert isinstance(windows, list) and len(windows) == 1, raw + return datetime.fromisoformat(str(object_value(windows[0])["reset_at"]).replace("Z", "+00:00")) + + +def _make_window_due(key: str) -> None: + due: Final = (datetime.now(UTC) - timedelta(days=1)).replace(tzinfo=None).isoformat() + windows: Final = [{"max_budget": 1.0, "budget_duration": "1d", "reset_at": due}] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection: + connection.execute( + 'UPDATE "LiteLLM_VerificationToken" SET budget_limits = %s::jsonb WHERE token = %s', + (json.dumps(windows), sha256(key.encode()).hexdigest()), + ) + + +@pytest.mark.timeout(CELL_TIMEOUT) +def test_a_reset_budget_window_keeps_the_long_ttl(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, FAST_RESET_TICK) as owned, owned.scenario() as scenario, rig_cache() as cache: + model: Final = scenario.model(**RATES) + key: Final = scenario.key(models=[model], budget_limits=[{"max_budget": 1.0, "budget_duration": "1d"}]) + window_counter: Final = f"{key_counter(key)}:window:1d" + assert _chat(owned, model, key) + assert _settled(window_counter, ONE_CALL_COST, 3600) + _make_window_due(key) + rolled: Final = eventually( + lambda: (_window_reset_at(key), counter_value(cache, window_counter), counter_ttl(cache, window_counter)), + lambda state: state[0].replace(tzinfo=None) > datetime.now(UTC).replace(tzinfo=None), + seconds=45, + ) + assert rolled[1] is not None and rolled[1] < ONE_CALL_COST, rolled + assert outlives(rolled[2]), rolled diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 57c32ace509..296bf678284 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -402,6 +402,25 @@ async def test_async_scan_iter_namespaces_pattern(namespace, expected_match, mon assert captured["match"] == expected_match +@pytest.mark.asyncio +async def test_async_scan_iter_decodes_bytes_keys_and_stops_at_count(monkeypatch, redis_no_ping): + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + def scan_iter(match, count): + async def gen(): + for key in (b"spend:key:a", "spend:key:b", b"spend:key:c", b"spend:key:d"): + yield key + + return gen() + + mock_redis_instance = MagicMock() + mock_redis_instance.scan_iter = scan_iter + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + keys = await redis_cache.async_scan_iter(pattern="spend:key:", count=3) + assert keys == ["spend:key:a", "spend:key:b", "spend:key:c"] + + @pytest.mark.parametrize("namespace, expected", [(None, "k"), ("ns", "ns:k")]) def test_increment_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "https://my-test-host") diff --git a/tests/unit/proxy/common_utils/test_reset_budget_job.py b/tests/unit/proxy/common_utils/test_reset_budget_job.py index c17aed49014..afb5dce44b7 100644 --- a/tests/unit/proxy/common_utils/test_reset_budget_job.py +++ b/tests/unit/proxy/common_utils/test_reset_budget_job.py @@ -4704,3 +4704,35 @@ async def test_reset_budget_for_litellm_team_members_called(): assert len(team_member_writes) == 1 assert team_member_writes[0]["where"]["budget_id"]["in"] == ["budget1"] assert team_member_writes[0]["data"] == {"spend": 0} + + +@pytest.mark.asyncio +async def test_reset_expired_window_writes_the_redis_counter_with_the_spend_counter_ttl(monkeypatch): + """The window reset is a counter write like every other, so it must carry the spend counter TTL, not the 60s cache default.""" + from datetime import datetime + + from litellm.proxy.common_utils.reset_budget_job import BudgetResetSettings, Litellm_EntityType + + fake_module = types.ModuleType("litellm.proxy.proxy_server") + fake_module.spend_counter_redis_ttl = lambda: 4321 + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", fake_module) + spend_counter_cache = MagicMock() + spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + window = {"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-01-01T00:00:00Z"} + + reset = await ResetBudgetJob._reset_expired_window( + window=window, + counter_key="spend:key:abc:window:1d", + spend_counter_cache=spend_counter_cache, + now=datetime(2026, 1, 2), + reset_settings=BudgetResetSettings(), + prisma_client=AsyncMock(), + entity_type=Litellm_EntityType.KEY, + entity_id="abc", + ) + + assert reset is True + spend_counter_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key="spend:key:abc:window:1d", value=0.0, ttl=4321 + ) + spend_counter_cache.in_memory_cache.set_cache.assert_called_once_with(key="spend:key:abc:window:1d", value=0.0) diff --git a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py index a98f1344279..92b2bf44a9f 100644 --- a/tests/unit/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_organization_endpoints.py @@ -1573,3 +1573,39 @@ async def test_delete_organization_evicts_the_cache_of_the_keys_it_deletes(monke assert all(cache.get_cache(key=cache_key) is None for cache_key in doomed_cache_keys) assert all(cache.get_cache(key=cache_key) == {"retained": True} for cache_key in kept_cache_keys) assert jwt_table.rows == (kept_row,) + + +def _org_spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + from litellm.caching.in_memory_cache import InMemoryCache + + spend_counter_cache = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend:org:org-1", value=4.2) + spend_counter_cache.redis_cache.async_delete_cache = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", spend_counter_cache) + return spend_counter_cache + + +@pytest.mark.asyncio +async def test_legacy_update_with_spend_drops_the_org_spend_counter(monkeypatch): + """A direct spend write must drop the cross-pod counter, or the org stays at the stale value until the counter expires.""" + spend_counter_cache = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "spend": 0.0}, existing_budget_id="budget-1" + ) + + spend_counter_cache.redis_cache.async_delete_cache.assert_awaited_once_with(key="spend:org:org-1") + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") is None + + +@pytest.mark.asyncio +async def test_legacy_update_without_spend_keeps_the_org_spend_counter(monkeypatch): + spend_counter_cache = _org_spend_counter_cache(monkeypatch) + + await _run_legacy_update_organization( + monkeypatch, body={"organization_id": "org-1", "organization_alias": "renamed"}, existing_budget_id="budget-1" + ) + + spend_counter_cache.redis_cache.async_delete_cache.assert_not_awaited() + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:org:org-1") == 4.2 diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 25180038e2a..69058542ffd 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -8059,3 +8059,57 @@ async def test_global_spend_refresh_runs_the_refresh_on_the_writer_when_a_reader assert [statement for statement in journal.statements if statement[1].startswith("REFRESH")] == [ ("writer", 'REFRESH MATERIALIZED VIEW "MonthlyGlobalSpend";') ] + + +def _spend_counter_cache_for_reset(monkeypatch: pytest.MonkeyPatch, *, scan_results: object) -> MagicMock: + from litellm.caching.in_memory_cache import InMemoryCache + + mock_prisma_client: Final = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + mock_prisma_client.db.litellm_teamtable.update_many = AsyncMock() + monkeypatch.setattr(ps, "prisma_client", mock_prisma_client) + spend_counter_cache: Final = MagicMock() + spend_counter_cache.in_memory_cache = InMemoryCache() + spend_counter_cache.in_memory_cache.set_cache(key="spend:key:aaa", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend:team:t1", value=0.5) + spend_counter_cache.in_memory_cache.set_cache(key="spend:user:u1", value=0.5) + spend_counter_cache.redis_cache.async_scan_iter = AsyncMock(side_effect=scan_results) + spend_counter_cache.redis_cache.delete_cache_keys = AsyncMock() + monkeypatch.setattr(ps, "spend_counter_cache", spend_counter_cache) + return spend_counter_cache + + +@pytest.mark.asyncio +async def test_global_spend_reset_drops_the_key_and_team_spend_counters(monkeypatch: pytest.MonkeyPatch) -> None: + spend_counter_cache: Final = _spend_counter_cache_for_reset( + monkeypatch, scan_results=[["spend:key:aaa", "spend:key:bbb"], [], ["spend:team:t1"], []] + ) + + response: Final = await spend_management_endpoints.global_spend_reset() + + assert response["status"] == "success" + assert [c.kwargs["pattern"] for c in spend_counter_cache.redis_cache.async_scan_iter.await_args_list] == [ + "spend:key:", + "spend:key:", + "spend:team:", + "spend:team:", + ] + assert [c.args for c in spend_counter_cache.redis_cache.delete_cache_keys.await_args_list] == [ + (["spend:key:aaa", "spend:key:bbb"],), + (["spend:team:t1"],), + ] + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:team:t1") is None + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:user:u1") == 0.5 + + +@pytest.mark.asyncio +async def test_global_spend_reset_still_resets_the_rows_when_redis_is_down(monkeypatch: pytest.MonkeyPatch) -> None: + spend_counter_cache: Final = _spend_counter_cache_for_reset(monkeypatch, scan_results=ConnectionError("redis down")) + + response: Final = await spend_management_endpoints.global_spend_reset() + + assert response["status"] == "success" + ps.prisma_client.db.litellm_teamtable.update_many.assert_awaited_once_with(data={"spend": 0.0}, where={}) + spend_counter_cache.redis_cache.delete_cache_keys.assert_not_awaited() + assert spend_counter_cache.in_memory_cache.get_cache(key="spend:key:aaa") is None