scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-02 07:13:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 48

fix(proxy): raise live budget counters for failure spend rows (#47741)

Browse source 
* fix(proxy): raise live budget counters for failure spend rows

A stream that breaks mid-flight writes its recovered cost to the spend log, but the failure hook
had already released the reservation and never touched the live counters, so an over-budget key
kept passing until the counter was next repaired from the database

Every charged failure row with a cost now raises the same counters the success path raises. The
error waits at most a second for them, so a hung Redis cannot hold it back

Fixes #35568

* test(proxy): cover failure spend counters across budgets, endpoints and outages

Integration cells for team, user, end-user, tag and window budgets,
concurrent streams, billed retries, native /v1/messages, Bedrock guardrail
cost, a Redis stall on two workers and disable_spend_updates

* test(proxy): check stalled failure counters and use the real spend writer

- The Redis stall cell now reads the burst key's live counter and requires at least the full failure spend
- Unit tests drop the charging writer double, bare dicts, list mutation and wall-clock timestamps
- Partial-stream test models are frozen

* test(proxy): move the stalled counter bound reason into the assert message

---------

Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
This commit is contained in:
yucheng-berri 2026-10-30 13:05:12 -07:00 • committed by GitHub
parent 8f968b6c8e
commit f4d1f45d87
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 1103 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import batch_cost_is_final
from litellm.caching.redis_batch import flush_post_call_redis_batches
from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
@ -134,6 +135,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset(
)
)
FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0
_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping)
@ -235,6 +238,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter:
return proxy_logging_obj.db_spend_update_writer
def _proxy_spend_counters() -> "_IncrementSpendCounters":
from litellm.proxy.proxy_server import increment_spend_counters
return increment_spend_counters
DeploymentLookup: TypeAlias = Callable[[str], Deployment | None]
@ -361,11 +370,16 @@ class _ProxyDBLogger(CustomLogger):
message_logging: bool = True,
spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer,
deployment_lookup: DeploymentLookup = _proxy_deployment,
spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters,
failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS,
) -> None:
super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging)
self.spend_event_producer = spend_event_producer
self._spend_writer: Final = spend_writer
self._deployment_lookup: Final = deployment_lookup
self._spend_counters: Final = spend_counters
self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds
self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done
async def async_log_success_event(
self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime
@ -590,9 +604,11 @@ class _ProxyDBLogger(CustomLogger):
completed_cost: bool = False,
request_tags: list[str] | None = None,
) -> None:
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair."""
for row_cost, row_kwargs in spend_rows:
if completed_cost:
"""One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the
live spend counters raised by every charged row that cost something. A completed cost settles through the
success path's write, which raises the counters against the kept reservation."""
if completed_cost:
for row_cost, row_kwargs in spend_rows:
from litellm.proxy.proxy_server import increment_spend_counters, proxy_logging_obj
await _update_database_and_spend_counters(
@ -623,12 +639,16 @@ class _ProxyDBLogger(CustomLogger):
router=get_llm_router(),
),
)
continue
return
token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id
charged: Final = [
await self._spend_writer().update_database(
token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict),
token=token,
response_cost=row_cost,
user_id=user_api_key_dict.user_id,
end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id,
end_user_id=end_user_id,
team_id=user_api_key_dict.team_id,
kwargs=dict(row_kwargs),
completion_response=original_exception,
@ -637,6 +657,91 @@ class _ProxyDBLogger(CustomLogger):
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
for row_cost, row_kwargs in spend_rows
]
counted_rows: Final = tuple(
row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0
)
if counted_rows and not ProxyUpdateSpend.disable_spend_updates():
await self._increment_failure_spend_counters(
user_api_key_dict, token, end_user_id, counted_rows, start_time
)
async def _increment_failure_spend_counters(
self,
user_api_key_dict: UserAPIKeyAuth,
token: str | None,
end_user_id: str | None,
counted_rows: Sequence[tuple[float, Mapping[str, object]]],
start_time: datetime,
) -> None:
"""The counters the success path raises for the same spend, sent before the client sees the error. A
counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of
holding the error back. The reservation was already released, so every counter takes the full row cost."""
increment: Final = asyncio.create_task(
self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time)
)
self._pending_increments.add(increment)
increment.add_done_callback(self._pending_increments.discard)
finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds)
if not finished:
verbose_proxy_logger.warning(
"Spend counters for failure spend rows took over %ss, finishing them in the background",
self._failure_counter_wait_seconds,
)
async def _raise_failure_spend_counters(
self,
user_api_key_dict: UserAPIKeyAuth,
token: str | None,
end_user_id: str | None,
counted_rows: Sequence[tuple[float, Mapping[str, object]]],
start_time: datetime,
) -> None:
for row_cost, row_kwargs in counted_rows:
await self._raise_row_spend_counters(
user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time
)
await flush_post_call_redis_batches()
async def _raise_row_spend_counters(
self,
user_api_key_dict: UserAPIKeyAuth,
token: str | None,
end_user_id: str | None,
row_cost: float,
row_kwargs: Mapping[str, object],
start_time: datetime,
) -> None:
standard_logging_object: Final = row_kwargs.get("standard_logging_object")
sl_object: Final = (
cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict
if isinstance(standard_logging_object, dict)
else None
)
metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata"))
served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields(
standard_logging_object
).get("model_id")
try:
await self._spend_counters()(
token=token,
team_id=user_api_key_dict.team_id,
user_id=user_api_key_dict.user_id,
response_cost=row_cost,
org_id=user_api_key_dict.org_id,
end_user_id=end_user_id,
tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)),
request_started_at=start_time,
model_access_groups=debitable_model_access_groups(
attributed=get_request_model_access_groups(row_kwargs),
served_model_id=served_model_id if isinstance(served_model_id, str) else None,
router=get_llm_router(),
),
project_id=user_api_key_dict.project_id,
)
except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing
verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row")
async def _PROXY_track_cost_callback(
self,
@ -1060,7 +1165,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]:
return tuple(sorted(str(key) for key in metadata))
def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:
def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None:
return budget_reservation_from_metadata(metadata)
@ -1095,6 +1200,7 @@ class _IncrementSpendCounters(Protocol):
tags: list[str] | None = None,
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
) -> None: ...
@ -1111,7 +1217,7 @@ async def _update_database_and_spend_counters(
start_time: datetime | None,
end_time: datetime | None,
response_cost: float,
budget_reservation: dict | None,
budget_reservation: dict[str, object] | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
@ -1174,7 +1280,7 @@ async def _update_database_and_spend_counters_in_batch(
start_time: datetime | None,
end_time: datetime | None,
response_cost: float,
budget_reservation: dict | None,
budget_reservation: dict[str, object] | None,
request_tags: list[str] | None,
model_access_groups: Sequence[str] | None,
project_id: str | None,

2
litellm/proxy/proxy_server.py
View file

@ -3510,7 +3510,7 @@ async def increment_spend_counters(
user_id: str | None,
response_cost: float | None,
org_id: str | None = None,
budget_reservation: dict | None = None,
budget_reservation: dict[str, object] | None = None,
end_user_id: str | None = None,
tags: list[str] | None = None,
request_started_at: datetime | None = None,

173
tests/integration/spend/_partial_stream.py Normal file
View file

@ -0,0 +1,173 @@
import json
import threading
import uuid
from collections.abc import Callable, Iterable, Iterator
from hashlib import sha256
from typing import Final
import httpx
from integration._support.client import Gateway, Scenario, eventually, held
from integration._support.wire import Reply, Request, Wire
from openai.types.chat import ChatCompletionChunk
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
BACKEND: Final = "integration-partial-stream"
PIECE: Final = "partial stream billed tokens "
DELIVERED_PIECES: Final = 12
_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"'
ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"'
class _SpendRow(BaseModel):
model_config = ConfigDict(frozen=True)
spend: float
class StatusRow(BaseModel):
model_config = ConfigDict(frozen=True)
status: str
spend: float
class _BudgetError(BaseModel):
model_config = ConfigDict(frozen=True)
type: str
class _ErrorResponse(BaseModel):
model_config = ConfigDict(frozen=True)
error: _BudgetError
def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes:
body: Final = {
"id": stream_id,
"object": "chat.completion.chunk",
"created": 1700000000,
"model": BACKEND,
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
}
return f"data: {json.dumps(body)}\n\n".encode()
def completion() -> bytes:
return json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1700000000,
"model": BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
).encode()
def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method != "POST":
return Reply(status=404)
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}"
return Reply(
content_type="text/event-stream",
chunks=(
*(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)),
_chunk(stream_id, {}, "stop"),
b"data: [DONE]\n\n",
),
abort_after=delivered,
pause_between_chunks=0.02,
gate_before_first=gate,
)
return Reply(body=completion())
return respond
def priced_model(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model=f"openai/{BACKEND}",
api_base=f"{wire.url}/v1",
api_key="synthetic-openai-key",
input_cost_per_token=0.001,
output_cost_per_token=0.002,
)
def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"max_tokens": 5,
"stream": True,
"messages": [{"role": "user", "content": "partial stream"}],
**fields,
},
key=key,
)
def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None:
response: Final = stream(gateway, model, key, **fields)
assert response.status_code == 200, response.text
events: Final = tuple(
line.removeprefix("data: ")
for line in response.text.splitlines()
if line.startswith("data: ") and line != "data: [DONE]"
)
assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text
delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1]))
assert delivered == PIECE * DELIVERED_PIECES, response.text
def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]:
for chunk in chunks:
yield from (choice.delta.content or "" for choice in chunk.choices)
def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response:
return gateway.request(
"POST",
"/v1/chat/completions",
{
"model": model,
"max_tokens": 5,
"messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}],
**fields,
},
key=key,
)
def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]:
digest: Final = sha256(key.encode()).hexdigest()
eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70)
found: Final = held(
lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2
)
return tuple(_SpendRow.model_validate(row).spend for row in found)
def chats_sent(wire: Wire) -> int:
return sum(1 for request in wire.drain() if request.method == "POST")
def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float:
key: Final = scenario.key(models=[model], max_budget=100)
broken_stream(gateway, model, key)
(spend,) = failure_spend(gateway, key)
assert spend > 0
return spend
def assert_budget_refusal(response: httpx.Response) -> None:
assert response.status_code == 422, response.text
assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text

239
tests/integration/spend/test_failure_spend_live_counters_chaos.py Normal file
View file

@ -0,0 +1,239 @@
import json
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha256
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.process import owned_proxy_process
from integration._support.redis_process import owned_redis
from integration._support.timing import assert_not_delayed_by, timed
from integration._support.wire import Reply, Request, Wire, wire_server
from integration.spend._partial_stream import (
DELIVERED_PIECES,
assert_budget_refusal,
broken_stream,
chat,
chats_sent,
failure_spend,
peer,
priced_model,
recovered_cost,
)
from pydantic import JsonValue
from redis import Redis
from litellm.constants import REDIS_SOCKET_TIMEOUT
_BURST: Final = 12
_SPREAD: Final = 6
_ATTACK_MARKER: Final = "synthetic-attack-marker"
_POLICY_ID: Final = "synthetic-priced-policy"
_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply"
def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path:
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
path: Final = directory / f"{name}.yaml"
path.write_text(yaml.safe_dump({**base, **sections}))
return path
def _redis_counter_config(directory: Path) -> Path:
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
return _proxy_config(
directory,
"redis-counters",
{
"router_settings": {
**base["router_settings"],
"redis_host": "os.environ/REDIS_HOST",
"redis_port": "os.environ/REDIS_PORT",
}
},
)
def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]:
return tuple(
gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]},
key=key,
headers={"connection": "close"},
)
for _ in range(_SPREAD)
)
@pytest.mark.timeout(600)
def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker(
gateway: Gateway, tmp_path: Path
) -> None:
release: Final = threading.Event()
with (
owned_redis(tmp_path) as counters,
wire_server(peer(DELIVERED_PIECES)) as wire,
wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire,
owned_proxy_process(
gateway,
tmp_path,
{"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)},
config=_redis_counter_config(tmp_path),
workers=2,
) as owned,
owned.gateway.scenario() as scenario,
):
scenario.cleanups.callback(release.set)
candidate: Final = owned.gateway
model: Final = priced_model(scenario, wire)
stalled_model: Final = priced_model(scenario, held_wire)
cost: Final = recovered_cost(candidate, scenario, model)
control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model])))
burst_key: Final = scenario.key(models=[stalled_model])
chats_sent(held_wire)
with ThreadPoolExecutor(max_workers=_BURST) as pool:
pending: Final = tuple(
pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST)
)
eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30)
with counters.paused():
release.set()
drained: Final = timed(lambda: tuple(future.result() for future in pending))
assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT)
assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST)
spent: Final = cost * _BURST
with Redis(host=counters.host, port=counters.port) as live_counters:
landed: Final = eventually(
lambda: float(live_counters.get(f"spend:key:{sha256(burst_key.encode()).hexdigest()}") or 0),
lambda counted: counted >= spent - 1e-9,
seconds=30,
return_last_on_timeout=True,
)
assert spent - 1e-9 <= landed <= 2 * spent + 1e-9, (
f"one broken stream costs {cost}. The cold counter is seeded from the database, and a spend flush"
" during the stall can land these rows there first, so up to twice their cost is expected, as for"
" completed streams held past a flush"
)
budgeted: Final = scenario.key(models=[model], max_budget=cost / 2)
broken_stream(candidate, model, budgeted)
before_kill: Final = _spread_requests(candidate, model, budgeted)
owned.kill_worker(min(owned.worker_pids()))
after_kill: Final = _spread_requests(candidate, model, budgeted)
for refused in (*before_kill, *after_kill):
assert_budget_refusal(refused)
def _policy(request: Request) -> Reply:
if request.method != "POST" or request.target != _POLICY_APPLY_PATH:
return Reply(status=404)
if _ATTACK_MARKER.encode() not in request.body:
return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode())
return Reply(
body=json.dumps(
{
"action": "GUARDRAIL_INTERVENED",
"outputs": [{"text": "blocked by synthetic policy"}],
"assessments": [
{"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}}
],
"usage": {"contentPolicyUnits": 1},
}
).encode()
)
def _policy_config(directory: Path, policy: Wire) -> Path:
return _proxy_config(
directory,
"priced-policy",
{
"guardrails": [
{
"guardrail_name": "priced-bedrock-policy",
"litellm_params": {
"guardrail": "bedrock",
"mode": "pre_call",
"default_on": True,
"guardrailIdentifier": _POLICY_ID,
"guardrailVersion": "DRAFT",
"aws_region_name": "us-east-1",
"aws_access_key_id": "AKIASYNTHETICGUARDRAIL",
"aws_secret_access_key": "synthetic-guardrail-secret",
"aws_bedrock_runtime_endpoint": policy.url,
},
}
]
},
)
def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response:
blocked: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]},
key=key,
)
assert blocked.status_code == 400, blocked.text
return blocked
@pytest.mark.timeout(480)
def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None:
with (
wire_server(_policy) as policy,
wire_server(peer(None)) as wire,
owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned,
owned.gateway.scenario() as scenario,
):
candidate: Final = owned.gateway
model: Final = priced_model(scenario, wire)
meter: Final = scenario.key(models=[model], max_budget=100)
_attack(candidate, model, meter)
(cost,) = failure_spend(candidate, meter)
assert cost > 0
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
chats_sent(wire)
_attack(candidate, model, key)
_attack(candidate, model, key)
after: Final = chat(candidate, model, key)
assert_budget_refusal(after)
assert chats_sent(wire) == 0
@pytest.mark.timeout(480)
def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled(
gateway: Gateway, tmp_path: Path
) -> None:
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config: Final = _proxy_config(
tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}}
)
with (
wire_server(peer(DELIVERED_PIECES)) as wire,
gateway.scenario() as metering,
owned_proxy_process(gateway, tmp_path, {}, config=config) as owned,
owned.gateway.scenario() as scenario,
):
cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire))
candidate: Final = owned.gateway
model: Final = priced_model(scenario, wire)
key: Final = scenario.key(models=[model], max_budget=cost / 2)
broken_stream(candidate, model, key)
after: Final = chat(candidate, model, key)
assert after.status_code == 200, after.text

350
tests/integration/spend/test_partial_stream_live_budget.py Normal file
View file

@ -0,0 +1,350 @@
import json
import threading
import uuid
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from hashlib import sha256
from typing import Final
import httpx
import pytest
from integration._support.client import Gateway, eventually, held
from integration._support.wire import Reply, Request, wire_server
from integration.spend._partial_stream import (
BACKEND,
DELIVERED_PIECES,
JSON_OBJECT,
PIECE,
ROWS,
StatusRow,
assert_budget_refusal,
broken_stream,
chat,
chats_sent,
completion,
failure_spend,
peer,
priced_model,
recovered_cost,
stream,
)
from pydantic import JsonValue
_CONCURRENT_STREAMS: Final = 6
_MESSAGES_BACKEND: Final = "claude-integration-partial-stream"
_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}}
@pytest.mark.timeout(240)
def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
key: Final = scenario.key(models=[model], max_budget=cost / 2)
chats_sent(wire)
broken_stream(gateway, model, key)
after: Final = chat(gateway, model, key)
assert_budget_refusal(after)
assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0))
assert chats_sent(wire) == 1
@pytest.mark.timeout(240)
def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
broken_stream(gateway, model, key)
broken_stream(gateway, model, key)
after: Final = chat(gateway, model, key)
assert_budget_refusal(after)
assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0))
@pytest.mark.timeout(240)
def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None:
with wire_server(peer(0)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
key: Final = scenario.key(models=[model], max_budget=0.05)
refused: Final = stream(gateway, model, key)
after: Final = chat(gateway, model, key)
assert refused.status_code == 500, refused.text
assert after.status_code == 200, after.text
assert failure_spend(gateway, key) == (0.0,)
def _broken_stream_then_refusal(
gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue
) -> httpx.Response:
broken_stream(gateway, model, spender, **fields)
return chat(gateway, model, other, **fields)
@pytest.mark.timeout(240)
def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
team: Final = scenario.team(models=[model], max_budget=cost / 2)
after: Final = _broken_stream_then_refusal(
gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model])
)
assert_budget_refusal(after)
@pytest.mark.timeout(240)
def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
user: Final = scenario.user(max_budget=cost / 2)
after: Final = _broken_stream_then_refusal(
gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model])
)
assert_budget_refusal(after)
@pytest.mark.timeout(240)
def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None:
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2})
scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]})
after: Final = _broken_stream_then_refusal(
gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user
)
assert_budget_refusal(after)
@pytest.mark.timeout(240)
def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None:
tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}"
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2})
scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag})
key: Final = scenario.key(models=[model])
after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]})
untagged: Final = chat(gateway, model, key)
assert_budget_refusal(after)
assert untagged.status_code == 200, untagged.text
@pytest.mark.timeout(240)
def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}])
after: Final = _broken_stream_then_refusal(gateway, model, key, key)
assert_budget_refusal(after)
@pytest.mark.timeout(240)
def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None:
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
cost: Final = recovered_cost(gateway, scenario, model)
key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5))
with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool:
tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS)))
within_budget: Final = chat(gateway, model, key)
broken_stream(gateway, model, key)
over_budget: Final = chat(gateway, model, key)
assert within_budget.status_code == 200, within_budget.text
assert_budget_refusal(over_budget)
assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx(
[0.0] + [cost] * (_CONCURRENT_STREAMS + 1)
)
@pytest.mark.timeout(240)
def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None:
with wire_server(peer(None)) as wire, gateway.scenario() as scenario:
model: Final = priced_model(scenario, wire)
meter: Final = scenario.key(models=[model], max_budget=100)
assert stream(gateway, model, meter, **_UNCACHED).status_code == 200
(metered,) = eventually(
lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)),
lambda found: len(found) == 1,
seconds=70,
)
cost: Final = StatusRow.model_validate(metered).spend
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
first: Final = stream(gateway, model, key, **_UNCACHED)
second: Final = stream(gateway, model, key, **_UNCACHED)
assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text)
assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2
third: Final = stream(gateway, model, key, **_UNCACHED)
assert_budget_refusal(third)
def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method != "POST":
return Reply(status=404)
if JSON_OBJECT.validate_json(request.body).get("model") == hung_model:
return Reply(chunks=(completion(),), gate_before_first=release)
return Reply(body=completion())
return respond
def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float:
timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}])
assert timed_out.status_code in (408, 500), timed_out.text
digest: Final = sha256(key.encode()).hexdigest()
landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70)
rows: Final = tuple(
StatusRow.model_validate(row)
for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2)
)
assert {row.status for row in rows} == {"failure"}, rows
return sum(row.spend for row in rows)
@pytest.mark.timeout(240)
def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None:
hung_backend: Final = f"{BACKEND}-hung"
release: Final = threading.Event()
with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario:
scenario.cleanups.callback(release.set)
hung: Final = scenario.model(
model=f"openai/{hung_backend}",
api_base=f"{wire.url}/v1",
api_key="synthetic-openai-key",
input_cost_per_token=0.001,
output_cost_per_token=0.002,
timeout=1,
num_retries=2,
)
served: Final = priced_model(scenario, wire)
billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100))
assert billed > 0
key: Final = scenario.key(models=[hung, served], max_budget=billed / 2)
assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed)
after: Final = chat(gateway, served, key)
assert_budget_refusal(after)
def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes:
return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode()
def _messages_peer(request: Request) -> Reply:
if request.method != "POST" or request.target != "/v1/messages":
return Reply(status=404)
message: Final[dict[str, JsonValue]] = {
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": _MESSAGES_BACKEND,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 1},
}
if JSON_OBJECT.validate_json(request.body).get("stream") is not True:
return Reply(
body=json.dumps(
{**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"}
).encode()
)
return Reply(
content_type="text/event-stream",
chunks=(
_messages_event("message_start", {"message": {**message, "content": []}}),
_messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}),
*(
_messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}})
for _ in range(DELIVERED_PIECES)
),
_messages_event("content_block_stop", {"index": 0}),
_messages_event(
"message_delta",
{"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}},
),
_messages_event("message_stop", {}),
),
abort_after=DELIVERED_PIECES + 2,
pause_between_chunks=0.02,
)
def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response:
return gateway.request(
"POST",
"/v1/messages",
{
"model": model,
"max_tokens": 5,
"stream": streamed,
"messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}],
},
key=key,
)
def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None:
response: Final = _messages(gateway, model, key, streamed=True)
assert response.status_code == 200, response.text
assert response.text.count(PIECE) == DELIVERED_PIECES, response.text
assert "message_stop" not in response.text, response.text
def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]:
digest: Final = sha256(key.encode()).hexdigest()
eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70)
return tuple(
StatusRow.model_validate(row)
for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2)
)
@pytest.mark.timeout(240)
def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None:
with wire_server(_messages_peer) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"anthropic/{_MESSAGES_BACKEND}",
api_base=wire.url,
api_key="synthetic-anthropic-key",
input_cost_per_token=0.001,
output_cost_per_token=0.002,
)
meter: Final = scenario.key(models=[model], max_budget=100)
_broken_messages_stream(gateway, model, meter)
(metered,) = _rows(gateway, meter, 1)
assert metered.spend > 0, metered
key: Final = scenario.key(models=[model], max_budget=metered.spend / 2)
_broken_messages_stream(gateway, model, key)
after: Final = _messages(gateway, model, key, streamed=False)
assert_budget_refusal(after)
assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0))

225
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -14,6 +14,7 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.dual_cache import DualCache
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth
from litellm.proxy.collector import SpendEventConsumer
@ -36,6 +37,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payloa
from litellm.proxy.utils import (
REJECTED_BEFORE_REQUEST_SETUP,
REJECTED_BEFORE_REQUEST_SETUP_KEY,
BilledFailedAttempt,
ProxyUpdateSpend,
)
from litellm.types.router import Deployment, LiteLLM_Params
@ -2676,6 +2678,229 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend():
assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05
@pytest.fixture()
def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache:
counters: Final = DualCache()
monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache())
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
return counters
def _failure_with_billed_attempts() -> Mapping[str, object]:
usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24)
return {
"combined_usage_object": usage,
"response_cost": 0.5,
"billed_failed_attempts": (
BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25),
BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5),
),
}
_LIVE_COUNTERS: Final = {
"spend:key:hashed-key": 0.2,
"spend:team:team-1": 0.3,
"spend:team_member:user-1:team-1": 0.35,
"spend:user:user-1": 0.4,
"spend:end_user:end-user-1": 0.45,
"spend:org:org-1": 0.5,
"spend:project:project-1": 0.55,
"spend:tag:batch-jobs": 0.6,
"spend:model_access_group:premium": 0.65,
}
async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None:
await logger.async_post_call_failure_hook(
request_data={
"model": "gpt-4.1",
"messages": [{"role": "user", "content": "Hello"}],
"metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]},
"proxy_server_request": {"request_id": "rid"},
**failure,
},
original_exception=Exception("stream broke mid-flight"),
user_api_key_dict=UserAPIKeyAuth(
api_key="hashed-key",
team_id="team-1",
user_id="user-1",
end_user_id="end-user-1",
org_id="org-1",
project_id="project-1",
),
)
async def _counter_values(counters: DualCache) -> dict[str, float]:
return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("failure", "charged"),
[
(
{
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
"response_cost": 0.125,
},
0.125,
),
(_failure_with_billed_attempts(), 1.25),
({}, 0.0),
],
)
async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost(
live_spend_counters: DualCache, failure: Mapping[str, object], charged: float
):
"""#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the
spend log, or an over-budget key keeps passing until the counter is next repaired from the database."""
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
await _fail_request(ProxyDBLogger(), failure)
assert await _counter_values(live_spend_counters) == pytest.approx(
{key: spend + charged for key, spend in _LIVE_COUNTERS.items()}
)
@pytest.mark.asyncio
async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache):
"""A failure whose cost was already finalized settles through the success path's write, which raises the
counters itself."""
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
await _fail_request(
ProxyDBLogger(), {"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}}
)
assert await _counter_values(live_spend_counters) == pytest.approx(
{key: spend + 0.125 for key, spend in _LIVE_COUNTERS.items()}
)
@pytest.mark.asyncio
async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True})
for key, spend in _LIVE_COUNTERS.items():
await live_spend_counters.async_set_cache(key, spend)
await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts())
assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS)
@pytest.mark.asyncio
async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row():
counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable"))
await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts())
assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5]
@pytest.mark.asyncio
@pytest.mark.timeout(30)
async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers():
store_answers: Final = asyncio.Event()
raised_cost: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future()
async def slow_counter_store(**counter_args: object) -> None:
await store_answers.wait()
raised_cost.set_result(counter_args["response_cost"])
logger: Final = ProxyDBLogger(spend_counters=lambda: slow_counter_store, failure_counter_wait_seconds=0)
await _fail_request(
logger,
{
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
"response_cost": 0.125,
},
)
assert not raised_cost.done()
store_answers.set()
assert await raised_cost == 0.125
@pytest.mark.asyncio
async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
):
"""The window row's one-time seed from the spend log stops at the request start, which keeps the failure row
it already logged from being counted twice."""
from litellm.proxy.proxy_server import proxy_logging_obj
window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue
await window_queue.flush_and_get_aggregated_window_spend_transactions()
key_cache: Final = DualCache()
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache)
request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc)
reset_at: Final = request_start + timedelta(days=10)
await key_cache.async_set_cache(
"hashed-key",
UserAPIKeyAuth(
token="hashed-key",
budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}],
),
)
await _fail_request(
ProxyDBLogger(),
{
"combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70),
"response_cost": 0.125,
"start_time": request_start,
},
)
assert [
(row["entity_type"], row["entity_id"], row["spend"], row["started_at"])
for row in await window_queue.flush_and_get_aggregated_window_spend_transactions()
] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")]
@pytest.mark.asyncio
async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group(
live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch
):
"""Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the
final attempt's groups either."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.llm_router",
litellm.Router(
model_list=[
{
"model_name": "gpt-4.1",
"litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"},
"model_info": {"id": "dep-b", "access_groups": ["pool"]},
}
]
),
)
await _fail_request(
ProxyDBLogger(),
{
**_failure_with_billed_attempts(),
"metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}},
},
)
assert {
group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}")
for group in ("premium", "pool")
} == pytest.approx({"premium": 0.25, "pool": 1.25})
@pytest.mark.asyncio
async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata():
"""MCP tool calls may only carry user_api_key; user/team rollups still need user_id."""