scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(proxy): release unclaimed budget reservations at request end (#42304)

Browse source 
* fix(proxy): release unclaimed budget reservations at request end

* fix(proxy): release unclaimed budget reservations of websocket sessions too

* test(proxy): drop the structural middleware inheritance check

* fix(proxy): claim the budget reservation on streaming pass-through before its cost callback

The SSE chunk processor hands its success handler to the logging worker
after the response, so the request-end release freed the reservation
first and left the key unguarded until the worker drained. Claim it at
both end-of-stream hand-offs, the immediate enqueue and the coroutine
parked for deferred dispatch.

Give the xai realtime test double the litellm_params attribute every
real Logging object carries, since the wrapper now reads it.

* test(pass-through): give the vertex streaming test doubles a litellm_params dict

The spec'd Logging mocks in test_vertex_ai_anthropic_streaming_cost_injection.py
lacked the instance attribute the chunk processor now reads to claim the budget
reservation. Also restores main's _lazy_openapi_snapshot.json: the branch's copy
had been regenerated under Python 3.14, which dedents one docstring description
that the CI regeneration on Python 3.12 keeps indented, and the PR adds no lazily
loaded route, so main's file is the correct one.

* fix(pass-through): claim the budget reservation only after its cost callback is enqueued

Every pass-through success hand-off stamped callback_bound before handing the
coroutine to the logging worker. When that enqueue raised, the reservation stayed
claimed with no callback left to reconcile it, so the request-end release skipped it
and the reserved cost stayed pinned on the key's counter. Enqueue first, then claim,
so a failed hand-off leaves the reservation for the request-end release.

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-21 19:51:12 -07:00 • committed by GitHub
parent 537e8ac068
commit e7cd97c6b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 1057 additions and 41 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/litellm_core_utils/core_helpers.py
View file

@ -359,6 +359,46 @@ def get_litellm_metadata_from_kwargs(kwargs: dict):
return {}
def _budget_reservation_on_auth_object(user_api_key_auth: object) -> object:
if isinstance(user_api_key_auth, Mapping):
return user_api_key_auth.get("budget_reservation")
return getattr(user_api_key_auth, "budget_reservation", None)
def budget_reservation_from_metadata(metadata: Mapping[str, object]) -> dict | None:
stamped: Final = metadata.get("user_api_key_budget_reservation")
if isinstance(stamped, dict):
return stamped
on_auth_object: Final = _budget_reservation_on_auth_object(metadata.get("user_api_key_auth"))
return on_auth_object if isinstance(on_auth_object, dict) else None
def _stamp_budget_reservation_callback_bound(litellm_params: Mapping[str, object], callback_bound: bool) -> None:
for metadata_variable_name in ("metadata", "litellm_metadata"):
metadata = litellm_params.get(metadata_variable_name)
if not isinstance(metadata, Mapping):
continue
budget_reservation = budget_reservation_from_metadata(metadata)
if budget_reservation is not None:
budget_reservation["callback_bound"] = callback_bound
def bind_budget_reservation_to_callbacks(litellm_params: Mapping[str, object]) -> None:
"""Mark the request's budget reservation as owned by the success callbacks of this call.
The proxy releases any reservation still unbound when the request ends; one bound here
is left for the cost callback, which may finish after the response has been sent. Bind
only where a success handler is guaranteed to run: a logging object merely existing is
not that, since the proxy builds one for every route before calling anything.
"""
_stamp_budget_reservation_callback_bound(litellm_params, True)
def unbind_budget_reservation_from_callbacks(litellm_params: Mapping[str, object]) -> None:
"""Hand a failed call's reservation back to the request-end release: failure handlers never settle it."""
_stamp_budget_reservation_callback_bound(litellm_params, False)
def reconstruct_model_name(
model_name: str,
custom_llm_provider: str | None,

2
litellm/litellm_core_utils/litellm_logging.py
View file

@ -5493,7 +5493,7 @@ def request_model_access_groups_from_litellm_params(litellm_params: Mapping[str,
"""Access groups the auth layer stamped onto this request, from whichever metadata field carries them.
Detached internal sub-calls only inherit the identity keys, so the auth object is the
fallback there, exactly as _get_budget_reservation_from_metadata does for reservations.
fallback there, exactly as budget_reservation_from_metadata does for reservations.
"""
for metadata_variable_name in ("metadata", "litellm_metadata"):
metadata = litellm_params.get(metadata_variable_name)

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

@ -650,6 +650,7 @@ async def user_api_key_auth_websocket_for_model(websocket: WebSocket, model: str
"type": "http",
"headers": scope_headers,
"path": ws_scope.get("path", ""),
"state": ws_scope.setdefault("state", {}), # mutable-ok: Starlette's socket state, shared with the request
}
for key in ("root_path", "app_root_path"):
if key in ws_scope:
@ -3086,31 +3087,30 @@ async def _reserve_budget_after_common_checks(
request: Request | None = None,
) -> None:
user_api_key_auth_obj.budget_reservation = None
if skip_budget_checks:
return
if general_settings.get("disable_budget_reservation") is True:
return
if not skip_budget_checks and general_settings.get("disable_budget_reservation") is not True:
from litellm.proxy.spend_tracking.budget_reservation import (
reserve_budget_for_request,
)
from litellm.proxy.spend_tracking.budget_reservation import (
reserve_budget_for_request,
)
user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request(
request_body=request_data,
route=route,
llm_router=llm_router,
valid_token=user_api_key_auth_obj,
team_object=team_object,
user_object=user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
raw_body=await read_raw_json_body(request=request),
)
user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request(
request_body=request_data,
route=route,
llm_router=llm_router,
valid_token=user_api_key_auth_obj,
team_object=team_object,
user_object=user_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
apply_user_budget_to_team_keys=general_settings.get("apply_user_budget_to_team_keys") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
raw_body=await read_raw_json_body(request=request),
)
if request is not None:
reservation: Final = user_api_key_auth_obj.budget_reservation
request.state.budget_reservation = reservation # rebind-ok: read by the release middleware
def _should_skip_budget_checks(

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

@ -11,6 +11,7 @@ from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
budget_reservation_from_metadata,
get_litellm_metadata_from_kwargs,
)
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@ -630,17 +631,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]:
def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None:
metadata_budget_reservation: Final = metadata.get("user_api_key_budget_reservation")
if isinstance(metadata_budget_reservation, dict):
return metadata_budget_reservation
user_api_key_auth_obj: Final = metadata.get("user_api_key_auth")
if user_api_key_auth_obj is None:
return None
if isinstance(user_api_key_auth_obj, dict):
budget_reservation: Final = user_api_key_auth_obj.get("budget_reservation")
return budget_reservation if isinstance(budget_reservation, dict) else None
return getattr(user_api_key_auth_obj, "budget_reservation", None)
return budget_reservation_from_metadata(metadata)
def _get_request_tags_for_cost_tracking(

33
litellm/proxy/middleware/budget_reservation_release_middleware.py Normal file
View file

@ -0,0 +1,33 @@
from collections.abc import Awaitable, Callable, Mapping
from typing import Final
from starlette.types import ASGIApp, Receive, Scope, Send
_SCOPES_AUTH_STAMPS: Final = frozenset({"http", "websocket"})
class BudgetReservationReleaseMiddleware:
"""Releases the budget reservation auth made for a request once no callback owns it.
Auth stamps the reservation on the request or socket state; a call that starts
claims it for the cost callbacks, which settle it on success or failure. When the
response has been sent or the socket has closed and the reservation is still
unclaimed, nothing else ever would, so it is released here instead of pinning the
spend counter until its TTL.
"""
def __init__(self, app: ASGIApp, release: Callable[[Mapping[str, object]], Awaitable[None]]) -> None:
self.app = app
self.release = release
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
if scope["type"] not in _SCOPES_AUTH_STAMPS:
await self.app(scope, receive, send)
return
try:
await self.app(scope, receive, send)
finally:
state: Final = scope.get("state")
budget_reservation: Final = state.get("budget_reservation") if isinstance(state, Mapping) else None
if isinstance(budget_reservation, Mapping):
await self.release(budget_reservation)

4
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -47,6 +47,7 @@ from litellm.constants import (
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
bind_budget_reservation_to_callbacks,
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
)
@ -1637,6 +1638,7 @@ async def pass_through_request(
**kwargs,
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
## CUSTOM HEADERS - `x-litellm-*`
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
@ -2561,6 +2563,7 @@ async def websocket_passthrough_request(
**success_kwargs,
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
# Call the proxy logging success hook
if proxy_logging_obj:
@ -2732,6 +2735,7 @@ async def _relay_passthrough_response_bytes(
**success_handler_kwargs,
)
)
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
def _extract_model_from_vertex_ai_setup(setup_response: Mapping[str, object]) -> str | None:

4
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -9,6 +9,7 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.asyncify import asyncify
from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
@ -218,6 +219,7 @@ class PassThroughStreamingHandler:
and response.status_code < 400
):
logging_scheduled = True
bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),)
except Exception as e:
verbose_proxy_logger.error("Error in chunk_processor: %s", e)
@ -250,6 +252,8 @@ class PassThroughStreamingHandler:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine())
except Exception as e:
verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)
else:
bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params)
@staticmethod
async def _route_streaming_logging_to_handler(

9
litellm/proxy/proxy_server.py
View file

@ -654,6 +654,9 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import (
BillableRequestMetricsMiddleware,
BillingRecorder,
)
from litellm.proxy.middleware.budget_reservation_release_middleware import (
BudgetReservationReleaseMiddleware,
)
from litellm.proxy.plugin_routes import (
register_plugins_from_config,
)
@ -729,7 +732,10 @@ from litellm.proxy.shutdown.scheduled_jobs import (
pause_scheduled_jobs,
stop_in_flight_scheduler_jobs,
)
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
from litellm.proxy.spend_tracking.budget_reservation import (
get_budget_window_start,
release_unbound_budget_reservation,
)
from litellm.proxy.spend_tracking.daily_global_spend_rollup import (
run_scheduled_daily_global_spend_reconcile,
)
@ -2358,6 +2364,7 @@ app.add_middleware(
# it sees prisma_client as of the first request rather than import time.
sink_factory=lambda: gateway_request_accumulator if prisma_client is not None else None,
)
app.add_middleware(BudgetReservationReleaseMiddleware, release=release_unbound_budget_reservation)
app.add_middleware(InFlightRequestsMiddleware)
app.add_middleware(SecurityHeadersMiddleware)

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

@ -366,6 +366,7 @@ async def reserve_budget_for_request(
"reserved_cost": reservation_cost,
"entries": applied_entries,
"finalized": False,
"callback_bound": False,
"input_cost": min(float(input_cost or 0.0), reservation_cost),
"input_tokens": max(input_token_counts.values(), default=None),
}
@ -474,6 +475,19 @@ async def release_or_invalidate_budget_reservation(
budget_reservation["finalized"] = True
async def release_unbound_budget_reservation(budget_reservation: Mapping[str, object]) -> None:
"""Release a reservation no logging callback took ownership of, once the request ended.
A handler whose litellm call never builds a logging object (batch cancel, file
content, anything without the client decorator) runs no cost callback, so nothing
else would ever reconcile its reservation. A bound reservation is left alone: its
success or failure handler settles it, possibly after the response has been sent.
"""
if not isinstance(budget_reservation, dict) or budget_reservation.get("callback_bound") is True:
return
await release_or_invalidate_budget_reservation(budget_reservation=budget_reservation)
async def _get_budget_counters(
request_body: dict,
valid_token: UserAPIKeyAuth,

19
litellm/rust_bridge/callbacks_legacy_python.py
View file

@ -59,9 +59,17 @@ def setup(
}
supplied: Final = arguments.get("litellm_logging_obj")
if isinstance(supplied, Logging):
return CallSetup(supplied, arguments)
return _claim_budget_reservation(CallSetup(supplied, arguments), asynchronous)
logger, prepared = function_setup(call_type, Rules(), start_time, *args, is_async_call=asynchronous, **arguments)
return CallSetup(logger, prepared)
return _claim_budget_reservation(CallSetup(logger, prepared), asynchronous)
def _claim_budget_reservation(call_setup: CallSetup, asynchronous: bool) -> CallSetup:
from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks
if asynchronous and not is_internal_call():
bind_budget_reservation_to_callbacks(call_setup.logger.litellm_params)
return call_setup
def check_limits(kwargs: Mapping[str, object]) -> None:
@ -96,6 +104,9 @@ def finalize(
class LoggingSurface(Protocol):
@property
def litellm_params(self) -> Mapping[str, object]: ...
def update_from_kwargs(
self,
kwargs: dict[str, object],
@ -236,8 +247,12 @@ def sync_success_for_async_call(
def failure_handler(
logger: LoggingSurface, error: Exception, start: datetime.datetime, end: datetime.datetime, asynchronous: bool
) -> Coroutine[object, object, None] | None:
from litellm.litellm_core_utils.core_helpers import unbind_budget_reservation_from_callbacks
trace: Final = "".join(traceback.format_exception(error))
if asynchronous:
if not is_internal_call():
unbind_budget_reservation_from_callbacks(logger.litellm_params)
return logger.async_failure_handler(error, trace, start, end)
logger.failure_handler(error, trace, start, end)
return None

9
litellm/utils.py
View file

@ -81,7 +81,11 @@ from litellm.constants import (
PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO,
TOOL_CHOICE_OBJECT_TOKEN_COUNT,
)
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.core_helpers import (
bind_budget_reservation_to_callbacks,
normalize_drop_params,
unbind_budget_reservation_from_callbacks,
)
from litellm.litellm_core_utils.fallback_generalizations import (
match_capability_generalizations,
match_fill_missing_generalizations,
@ -1880,6 +1884,8 @@ def client(original_function):
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
if not _is_litellm_internal_call:
bind_budget_reservation_to_callbacks(logging_obj.litellm_params)
kwargs["litellm_logging_obj"] = logging_obj
modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type)
@ -2081,6 +2087,7 @@ def client(original_function):
# the failure hook ran, so a slow callback doesn't inflate the reported duration.
end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with
if logging_obj and not _is_litellm_internal_call:
unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
try:
logging_obj.failure_handler(
e, traceback_exception, start_time, end_time

4
tests/pass_through_unit_tests/test_vertex_ai_anthropic_streaming_cost_injection.py
View file

@ -53,6 +53,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_enabled():
# Setup logging object with model info
litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
litellm_logging_obj.litellm_params = {}
litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"}
litellm_logging_obj.completion_start_time = None
litellm_logging_obj.async_success_handler = AsyncMock()
@ -132,6 +133,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_disabled():
response.aiter_bytes = mock_aiter_bytes
litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
litellm_logging_obj.litellm_params = {}
litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"}
litellm_logging_obj.completion_start_time = None
litellm_logging_obj.async_success_handler = AsyncMock()
@ -194,6 +196,7 @@ async def test_vertex_ai_anthropic_streaming_cost_injection_no_usage_chunk():
response.aiter_bytes = mock_aiter_bytes
litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
litellm_logging_obj.litellm_params = {}
litellm_logging_obj.model_call_details = {"model": "claude-sonnet-4@20250514"}
litellm_logging_obj.completion_start_time = None
litellm_logging_obj.async_success_handler = AsyncMock()
@ -249,6 +252,7 @@ async def test_vertex_ai_anthropic_streaming_model_extraction():
response.aiter_bytes = mock_aiter_bytes
litellm_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
litellm_logging_obj.litellm_params = {}
litellm_logging_obj.model_call_details = {}
litellm_logging_obj.completion_start_time = None
litellm_logging_obj.async_success_handler = AsyncMock()

55
tests/test_litellm/litellm_core_utils/test_core_helpers.py
View file

@ -6,6 +6,8 @@ import pytest
from litellm.litellm_core_utils.core_helpers import (
_FINISH_REASON_MAP,
bind_budget_reservation_to_callbacks,
budget_reservation_from_metadata,
drop_params_env_flag,
drop_params_flag,
get_or_create_metadata_bucket,
@ -13,7 +15,60 @@ from litellm.litellm_core_utils.core_helpers import (
normalize_drop_params,
reconstruct_model_name,
redact_nested_match_and_regex_keys,
unbind_budget_reservation_from_callbacks,
)
from litellm.proxy._types import UserAPIKeyAuth
class TestBudgetReservationBinding:
"""The request-end release skips a reservation a cost callback has claimed, so the claim
must land on the one dict auth stamped, through whichever metadata field or auth object
carries it, and a failed call must be able to hand it back."""
@staticmethod
def _reservation() -> dict:
return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
@pytest.mark.parametrize("metadata_variable_name", ["metadata", "litellm_metadata"])
def test_reservation_stamped_on_the_metadata_is_bound(self, metadata_variable_name: str):
reservation = self._reservation()
bind_budget_reservation_to_callbacks({metadata_variable_name: {"user_api_key_budget_reservation": reservation}})
assert reservation["callback_bound"] is True
def test_reservation_reachable_only_through_the_auth_object_is_bound(self):
reservation = self._reservation()
user_api_key_auth = UserAPIKeyAuth(token="hashed")
user_api_key_auth.budget_reservation = reservation
bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": user_api_key_auth}})
assert reservation["callback_bound"] is True
def test_reservation_reachable_only_through_a_dumped_auth_object_is_bound(self):
reservation = self._reservation()
bind_budget_reservation_to_callbacks({"metadata": {"user_api_key_auth": {"budget_reservation": reservation}}})
assert reservation["callback_bound"] is True
def test_unbind_hands_a_claimed_reservation_back(self):
reservation = self._reservation()
litellm_params = {"litellm_metadata": {"user_api_key_budget_reservation": reservation}}
bind_budget_reservation_to_callbacks(litellm_params)
unbind_budget_reservation_from_callbacks(litellm_params)
assert reservation["callback_bound"] is False
def test_request_without_a_reservation_binds_nothing(self):
metadata = {"user_api_key_auth": UserAPIKeyAuth(token="hashed")}
bind_budget_reservation_to_callbacks({"metadata": metadata, "litellm_metadata": None})
assert budget_reservation_from_metadata(metadata) is None
assert "user_api_key_budget_reservation" not in metadata
class TestGetOrCreateMetadataBucket:

17
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -26,6 +26,7 @@ from litellm.litellm_core_utils.litellm_logging import (
set_callbacks,
)
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse
from litellm.types.utils import (
CallTypes,
@ -7941,3 +7942,19 @@ def test_response_cost_calculator_prices_terminal_responses_event_from_its_respo
assert event_cost is not None and event_cost > 0
assert event_cost == inner_cost
assert logging_obj.cost_breakdown["input_cost"] is not None and logging_obj.cost_breakdown["input_cost"] > 0
class TestBudgetReservationBinding:
"""The proxy builds a logging object for every route before calling anything, so a
logging object seeing the reservation is no promise that a cost callback will settle
it: the claim belongs to the call wrapper, and this object must leave it unbound."""
def test_update_environment_variables_leaves_the_reservation_unbound(self, logging_obj):
reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
logging_obj.update_environment_variables(
litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={}
)
assert logging_obj.litellm_params["metadata"]["user_api_key_budget_reservation"] is reservation
assert reservation["callback_bound"] is False

3
tests/test_litellm/llms/xai/test_xai_key_fallback.py
View file

@ -12,6 +12,9 @@ from litellm.types.router import GenericLiteLLMParams
class FakeLogging:
def __init__(self) -> None:
self.litellm_params: dict = {}
def update_from_kwargs(self, **kwargs):
pass

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

@ -59,6 +59,7 @@ from litellm.proxy.auth.user_api_key_auth import (
_user_api_key_auth_builder,
get_api_key,
user_api_key_auth,
user_api_key_auth_websocket_for_model,
)
from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
@ -9043,3 +9044,90 @@ async def test_router_settings_model_group_alias_authorizes_target_for_team(monk
await authorize()
assert (await request.json())["model"] == target
assert get_client_requested_model(request) == "AgentX-LLM"
@pytest.mark.asyncio
async def test_reserve_budget_after_common_checks_hands_the_reservation_to_the_request_state():
from fastapi import Request
request = Request(scope={"type": "http"})
user_api_key_auth_obj = UserAPIKeyAuth(token="test_token")
reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
with patch(
"litellm.proxy.spend_tracking.budget_reservation.reserve_budget_for_request",
new=AsyncMock(return_value=reservation),
):
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
request_data={"model": "gpt-4o"},
route="/v1/batches/batch_123/cancel",
llm_router=None,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
skip_budget_checks=False,
general_settings={},
request=request,
)
assert user_api_key_auth_obj.budget_reservation is reservation
assert request.state.budget_reservation is reservation
assert request.scope["state"]["budget_reservation"] is reservation
@pytest.mark.asyncio
async def test_reserve_budget_after_common_checks_clears_the_request_state_when_budget_checks_skip():
from fastapi import Request
request = Request(scope={"type": "http", "state": {"budget_reservation": {"reserved_cost": 0.5}}})
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=UserAPIKeyAuth(token="test_token"),
request_data={"model": "free-model"},
route="/v1/chat/completions",
llm_router=None,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=MagicMock(),
proxy_logging_obj=MagicMock(),
skip_budget_checks=True,
general_settings={},
request=request,
)
assert request.state.budget_reservation is None
@pytest.mark.asyncio
async def test_websocket_auth_hands_the_reservation_to_the_socket_state():
from fastapi import WebSocket
reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
websocket = WebSocket(
scope={
"type": "websocket",
"path": "/v1/realtime",
"headers": [(b"authorization", b"Bearer sk-1234")],
"query_string": b"model=gpt-realtime",
},
receive=AsyncMock(),
send=AsyncMock(),
)
async def auth_that_reserves(request, api_key):
request.state.budget_reservation = reservation
return UserAPIKeyAuth(token="hashed", budget_reservation=reservation)
with patch(
"litellm.proxy.auth.user_api_key_auth.user_api_key_auth",
new=AsyncMock(side_effect=auth_that_reserves),
):
result = await user_api_key_auth_websocket_for_model(websocket, model="gpt-realtime")
assert result.budget_reservation == reservation
assert websocket.state.budget_reservation is reservation
assert websocket.scope["state"]["budget_reservation"] is reservation

349
tests/test_litellm/proxy/middleware/test_budget_reservation_release_middleware.py Normal file
View file

@ -0,0 +1,349 @@
"""
Tests for BudgetReservationReleaseMiddleware.
Auth reserves budget before the handler runs and hands the reservation to the
request or socket state. A litellm call made through the async client wrapper
claims it for the cost callback that runs after the call; anything still unclaimed
when the response is done or the socket has closed would keep the spend counter
pinned until its TTL, so the middleware releases it.
"""
import asyncio
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
from datetime import datetime
from typing import Final
import pytest
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from starlette.types import ASGIApp, Message, Receive, Scope, Send
from starlette.websockets import WebSocket
import litellm
from litellm.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.middleware.budget_reservation_release_middleware import (
BudgetReservationReleaseMiddleware,
)
from litellm.proxy.spend_tracking.budget_reservation import (
reconcile_budget_reservation,
release_unbound_budget_reservation,
reserve_budget_for_request,
)
from litellm.proxy.utils import ProxyLogging
from litellm.utils import Rules, function_setup
KEY_TOKEN: Final = "hashed-release-middleware-key"
COUNTER_KEY: Final = f"spend:key:{KEY_TOKEN}"
CHAT_BODY: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}
@pytest.fixture
def spend_counter_cache(monkeypatch: pytest.MonkeyPatch) -> DualCache:
cache: Final = DualCache()
monkeypatch.setattr(proxy_server, "spend_counter_cache", cache)
monkeypatch.setattr(proxy_server, "prisma_client", None)
return cache
@pytest.fixture
def no_callbacks(monkeypatch: pytest.MonkeyPatch) -> None:
for callback_list_name in (
"callbacks",
"success_callback",
"failure_callback",
"_async_success_callback",
"_async_failure_callback",
):
monkeypatch.setattr(litellm, callback_list_name, [])
async def _reserve() -> dict:
reservation: Final = await reserve_budget_for_request(
request_body=CHAT_BODY,
route="/v1/chat/completions",
llm_router=None,
valid_token=UserAPIKeyAuth(token=KEY_TOKEN, max_budget=1.0, spend=0.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
proxy_logging_obj=ProxyLogging(user_api_key_cache=UserApiKeyCache()),
)
assert reservation is not None
assert reservation["reserved_cost"] > 0
return reservation
async def _chat(reservation: dict, **kwargs: object) -> object:
return await litellm.acompletion(
**CHAT_BODY,
metadata={"user_api_key_budget_reservation": reservation},
**kwargs,
)
def _proxy_pre_call_setup(route_type: str, reservation: dict) -> None:
function_setup(
original_function=route_type,
rules_obj=Rules(),
start_time=datetime.now(),
**CHAT_BODY,
litellm_call_id="proxy-pre-call-setup",
metadata={"user_api_key_budget_reservation": reservation},
)
def _app(
handler: Callable[[Request], Awaitable[Response]],
release: Callable[[Mapping[str, object]], Awaitable[None]] = release_unbound_budget_reservation,
) -> Starlette:
app: Final = Starlette(routes=[Route("/", handler, methods=["POST"])])
app.add_middleware(BudgetReservationReleaseMiddleware, release=release)
return app
async def _post(app: ASGIApp) -> None:
scope: Final = {
"type": "http",
"method": "POST",
"path": "/",
"raw_path": b"/",
"headers": [],
"query_string": b"",
"scheme": "http",
"server": ("testserver", 80),
"client": ("testclient", 1),
}
body_delivered: Final = asyncio.Event()
client_never_disconnects: Final = asyncio.Event()
async def receive() -> Message:
if body_delivered.is_set():
await client_never_disconnects.wait()
body_delivered.set()
return {"type": "http.request", "body": b"", "more_body": False}
async def send(message: Message) -> None:
return None
await app(scope, receive, send)
def _counter(spend_counter_cache: DualCache) -> float | None:
return spend_counter_cache.in_memory_cache.get_cache(key=COUNTER_KEY)
@pytest.mark.asyncio
async def test_unbound_reservation_is_released_after_the_response(spend_counter_cache: DualCache):
reservation: Final = await _reserve()
assert _counter(spend_counter_cache) == pytest.approx(reservation["reserved_cost"])
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
return JSONResponse({"id": "batch_123", "status": "cancelling"})
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_unbound_reservation_is_released_when_the_handler_raises(spend_counter_cache: DualCache):
reservation: Final = await _reserve()
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
raise RuntimeError("upstream refused the cancel")
with pytest.raises(RuntimeError):
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_seen_only_by_the_proxy_pre_call_logging_object_is_released(
spend_counter_cache: DualCache, no_callbacks: None
):
reservation: Final = await _reserve()
async def cancel_batch_without_a_client_wrapper() -> dict:
return {"id": "batch_123", "status": "cancelling"}
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
_proxy_pre_call_setup("acancel_batch", reservation)
return JSONResponse(await cancel_batch_without_a_client_wrapper())
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_of_a_failed_call_is_released_after_the_error_response(
spend_counter_cache: DualCache, no_callbacks: None
):
reservation: Final = await _reserve()
refused: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o")
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
_proxy_pre_call_setup("acompletion", reservation)
try:
await _chat(reservation, mock_response=refused)
except litellm.AuthenticationError:
return JSONResponse({"error": {"message": "bad key"}}, status_code=401)
raise AssertionError("the mocked call must fail")
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_reservation_claimed_by_a_completed_call_is_left_for_the_callback(
spend_counter_cache: DualCache, no_callbacks: None
):
reservation: Final = await _reserve()
reserved_cost: Final = reservation["reserved_cost"]
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
_proxy_pre_call_setup("acompletion", reservation)
response: Final = await _chat(reservation, mock_response="ok")
return JSONResponse(response.model_dump())
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(reserved_cost)
assert reservation["finalized"] is False
@pytest.mark.asyncio
async def test_reservation_claimed_by_a_streaming_call_is_left_for_the_callback_that_finishes_after_the_response(
spend_counter_cache: DualCache, no_callbacks: None
):
reservation: Final = await _reserve()
reserved_cost: Final = reservation["reserved_cost"]
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
_proxy_pre_call_setup("acompletion", reservation)
stream: Final = await _chat(reservation, mock_response="ok", stream=True)
async def sse() -> AsyncIterator[bytes]:
async for chunk in stream:
yield f"data: {chunk.model_dump_json()}\n\n".encode()
yield b"data: [DONE]\n\n"
return StreamingResponse(sse(), media_type="text/event-stream")
await _post(_app(handler))
assert _counter(spend_counter_cache) == pytest.approx(reserved_cost)
assert reservation["finalized"] is False
actual_cost: Final = reserved_cost / 4
await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=actual_cost)
assert _counter(spend_counter_cache) == pytest.approx(actual_cost)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_unbound_reservation_of_a_websocket_session_is_released_when_the_socket_closes(
spend_counter_cache: DualCache,
):
reservation: Final = await _reserve()
async def listen_without_a_provider_key(scope: Scope, receive: Receive, send: Send) -> None:
websocket: Final = WebSocket(scope, receive, send)
websocket.state.budget_reservation = reservation
await websocket.close(code=1011, reason="Required 'DEEPGRAM_API_KEY' in environment")
async def receive() -> Message:
return {"type": "websocket.connect"}
async def send(message: Message) -> None:
return None
middleware: Final = BudgetReservationReleaseMiddleware(
listen_without_a_provider_key, release=release_unbound_budget_reservation
)
await middleware({"type": "websocket", "path": "/deepgram/v1/listen", "headers": []}, receive, send)
assert _counter(spend_counter_cache) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_release_runs_once_per_request_with_the_stamped_reservation():
released: Final = []
reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
async def release(budget_reservation: Mapping[str, object]) -> None:
released.append(budget_reservation)
async def handler(request: Request) -> Response:
request.state.budget_reservation = reservation
return JSONResponse({})
await _post(_app(handler, release=release))
assert released == [reservation]
assert released[0] is reservation
@pytest.mark.asyncio
async def test_request_without_a_reservation_releases_nothing():
released: Final = []
async def release(budget_reservation: Mapping[str, object]) -> None:
released.append(budget_reservation)
async def unauthenticated(request: Request) -> Response:
return JSONResponse({})
async def budget_checks_skipped(request: Request) -> Response:
request.state.budget_reservation = None
return JSONResponse({})
await _post(_app(unauthenticated, release=release))
await _post(_app(budget_checks_skipped, release=release))
assert released == []
@pytest.mark.asyncio
async def test_lifespan_scopes_pass_through():
released: Final = []
seen: Final = []
async def release(budget_reservation: Mapping[str, object]) -> None:
released.append(budget_reservation)
async def inner(scope: Scope, receive: Receive, send: Send) -> None:
seen.append(scope["type"])
async def receive() -> Message:
return {"type": "lifespan.startup"}
async def send(message: Message) -> None:
return None
middleware: Final = BudgetReservationReleaseMiddleware(inner, release=release)
await middleware({"type": "lifespan", "state": {"budget_reservation": {"reserved_cost": 1.0}}}, receive, send)
assert seen == ["lifespan"]
assert released == []

106
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -4256,6 +4256,112 @@ async def test_pass_through_request_non_streaming_success_unchanged():
mock_success_handler.assert_called_once()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"upstream_status_code, claimed_by_the_success_handler",
[(200, True), (500, False)],
ids=["success-claims-the-reservation", "upstream-error-leaves-it-for-the-request-end-release"],
)
async def test_pass_through_request_claims_the_budget_reservation_only_when_its_success_handler_runs(
upstream_status_code: int, claimed_by_the_success_handler: bool
):
reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed")
user_api_key_dict.budget_reservation = reservation
upstream_response: Final = httpx.Response(
status_code=upstream_status_code,
headers={"content-type": "application/json"},
content=b'{"status": "upstream"}',
request=httpx.Request("POST", "http://target-api.com/api/generate"),
)
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processing,
patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker,
):
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_processing.get_custom_headers.return_value = {}
mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=lambda async_coroutine: async_coroutine.close())
async_client = MagicMock()
async_client.build_request = MagicMock(return_value=MagicMock())
async_client.send = AsyncMock(return_value=upstream_response)
mock_get_client.return_value = MagicMock(client=async_client)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/mock-upstream/api/generate"
mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}')
mock_request.headers = Headers({"content-type": "application/json"})
mock_request.query_params = QueryParams({})
response = await pass_through_request(
request=mock_request,
target="http://target-api.com/api/generate",
custom_headers={},
user_api_key_dict=user_api_key_dict,
)
assert response.status_code == upstream_status_code
assert reservation["callback_bound"] is claimed_by_the_success_handler
assert mock_worker.ensure_initialized_and_enqueue.call_count == int(claimed_by_the_success_handler)
@pytest.mark.asyncio
async def test_pass_through_request_leaves_the_budget_reservation_for_the_request_end_release_when_its_success_handler_cannot_be_enqueued():
reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
user_api_key_dict: Final = UserAPIKeyAuth(api_key="hashed")
user_api_key_dict.budget_reservation = reservation
upstream_response: Final = httpx.Response(
status_code=200,
headers={"content-type": "application/json"},
content=b'{"status": "upstream"}',
request=httpx.Request("POST", "http://target-api.com/api/generate"),
)
def refuse_to_enqueue(async_coroutine):
async_coroutine.close()
raise RuntimeError("logging worker is shutting down")
with (
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client") as mock_get_client,
patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.ProxyBaseLLMRequestProcessing"
) as mock_processing,
patch("litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER") as mock_worker,
):
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_failure_hook = AsyncMock()
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value=None)
mock_processing.get_custom_headers.return_value = {}
mock_worker.ensure_initialized_and_enqueue = MagicMock(side_effect=refuse_to_enqueue)
async_client = MagicMock()
async_client.build_request = MagicMock(return_value=MagicMock())
async_client.send = AsyncMock(return_value=upstream_response)
mock_get_client.return_value = MagicMock(client=async_client)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/mock-upstream/api/generate"
mock_request.body = AsyncMock(return_value=b'{"prompt": "hi"}')
mock_request.headers = Headers({"content-type": "application/json"})
mock_request.query_params = QueryParams({})
with pytest.raises(ProxyException):
await pass_through_request(
request=mock_request,
target="http://target-api.com/api/generate",
custom_headers={},
user_api_key_dict=user_api_key_dict,
)
assert reservation["callback_bound"] is False
@pytest.mark.asyncio
async def test_pass_through_request_internal_failure_still_raises_proxy_exception():
"""

69
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py
View file

@ -838,3 +838,72 @@ async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exceptio
assert failure_payload["completion_tokens"] == 12
assert failure_payload["response_cost"] > 12 * 3.75e-06
assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"deferred_dispatch_armed",
[False, True],
ids=["enqueued-at-end-of-stream", "parked-for-deferred-dispatch"],
)
async def test_chunk_processor_claims_the_budget_reservation_before_handing_it_to_the_cost_callback(
deferred_dispatch_armed: bool,
):
reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
response = _make_streaming_response([b"event-1", b"event-2"])
logging_obj = _unarmed_logging_obj()
logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}}
if deferred_dispatch_armed:
logging_obj._on_deferred_stream_complete = AsyncMock()
claimed_when_the_callback_ran = []
async def cost_callback(**kwargs):
claimed_when_the_callback_ran.append(reservation["callback_bound"])
async for _ in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/bedrock/model/claude/invoke-with-response-stream",
route_streaming_logging=cost_callback,
):
pass
if deferred_dispatch_armed:
(parked_cost_callback,) = logging_obj._deferred_stream_complete_args
await parked_cost_callback
else:
await GLOBAL_LOGGING_WORKER.flush()
assert reservation["callback_bound"] is True
assert claimed_when_the_callback_ran == [True]
@pytest.mark.asyncio
async def test_chunk_processor_leaves_the_budget_reservation_for_the_request_end_release_when_the_cost_callback_cannot_be_enqueued():
reservation = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
response = _make_streaming_response([b"event-1", b"event-2"])
logging_obj = _unarmed_logging_obj()
logging_obj.litellm_params = {"metadata": {"user_api_key_budget_reservation": reservation}}
def refuse_to_enqueue(async_coroutine):
async_coroutine.close()
raise RuntimeError("logging worker is shutting down")
with patch.object(GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=refuse_to_enqueue):
async for _ in PassThroughStreamingHandler.chunk_processor(
response=response,
request_body={"model": "claude-3-haiku"},
litellm_logging_obj=logging_obj,
endpoint_type=EndpointType.GENERIC,
start_time=datetime.now(),
passthrough_success_handler_obj=MagicMock(),
url_route="/bedrock/model/claude/invoke-with-response-stream",
route_streaming_logging=AsyncMock(),
):
pass
assert reservation["callback_bound"] is False

39
tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py
View file

@ -26,6 +26,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
_get_team_member_budget_counter,
count_request_input_tokens,
estimate_request_max_cost,
release_unbound_budget_reservation,
reserve_budget_for_request,
)
from litellm.proxy.utils import ProxyLogging
@ -546,3 +547,41 @@ async def test_team_member_reservation_counter_adds_temp_increase_to_live_team_d
assert counter is not None
assert counter.max_budget == expected_max_budget
assert counter.fallback_spend == 0.5
@pytest.mark.asyncio
async def test_reservation_starts_unbound_to_any_callback():
reservation: Final = await _reserve("/v1/responses")
assert reservation is not None
assert reservation["callback_bound"] is False
@pytest.mark.asyncio
async def test_release_unbound_budget_reservation_frees_the_counter(spend_counter_cache: DualCache):
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
reservation: Final = await _reserve_for_tiny_budget_key(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}
)
assert reservation is not None
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"])
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(0.0)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_release_unbound_budget_reservation_leaves_a_bound_one_to_its_callback(spend_counter_cache: DualCache):
counter_key: Final = f"spend:key:{TINY_BUDGET_KEY_TOKEN}"
reservation: Final = await _reserve_for_tiny_budget_key(
"/v1/chat/completions", {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]}
)
assert reservation is not None
reservation["callback_bound"] = True
await release_unbound_budget_reservation(reservation)
assert spend_counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(reservation["reserved_cost"])
assert reservation["finalized"] is False

79
tests/test_litellm/rust_bridge/test_callbacks_legacy_python.py
View file

@ -9,9 +9,10 @@ import pytest
from pydantic import TypeAdapter
import litellm
from litellm._internal_context import is_internal_call
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.rust_bridge import callbacks_legacy_python as legacy
from litellm.rust_bridge.callbacks_legacy_python import check_limits, setup
from litellm.rust_bridge.callbacks_legacy_python import check_limits, failure_handler, setup
_OCR_KWARGS: Final = MappingProxyType(
{
@ -81,6 +82,82 @@ def test_setup_builds_a_logger_when_none_is_supplied(call_type: str, kwargs: Map
assert result.logger.litellm_call_id == result.kwargs["litellm_call_id"]
def _budget_reservation() -> dict:
return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False}
def _kwargs_with_a_budget_reservation(reservation: dict) -> dict[str, object]:
return {**_OCR_KWARGS, "metadata": {"user_api_key_budget_reservation": reservation}}
def test_setup_claims_the_budget_reservation_for_an_async_call() -> None:
reservation: Final = _budget_reservation()
setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True)
assert reservation["callback_bound"] is True
def test_setup_claims_the_budget_reservation_a_supplied_logger_already_saw() -> None:
reservation: Final = _budget_reservation()
supplied: Final = _supplied_logger()
supplied.update_environment_variables(
litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}}, optional_params={}
)
assert reservation["callback_bound"] is False
setup("aocr", (), {**_OCR_KWARGS, "litellm_logging_obj": supplied}, datetime.datetime.now(), asynchronous=True)
assert reservation["callback_bound"] is True
def test_setup_leaves_the_budget_reservation_alone_for_a_sync_call() -> None:
reservation: Final = _budget_reservation()
setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=False)
assert reservation["callback_bound"] is False
def test_setup_leaves_the_budget_reservation_alone_for_an_internal_call() -> None:
reservation: Final = _budget_reservation()
token: Final = is_internal_call.set(True)
try:
setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), datetime.datetime.now(), asynchronous=True)
finally:
is_internal_call.reset(token)
assert reservation["callback_bound"] is False
def test_failure_handler_hands_the_budget_reservation_back_for_an_async_call() -> None:
reservation: Final = _budget_reservation()
now: Final = datetime.datetime.now()
result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True)
assert reservation["callback_bound"] is True
pending: Final = failure_handler(result.logger, RuntimeError("upstream refused"), now, now, asynchronous=True)
assert reservation["callback_bound"] is False
assert pending is not None
pending.close()
def test_failure_handler_of_an_internal_call_leaves_the_outer_budget_reservation_claim_in_place() -> None:
reservation: Final = _budget_reservation()
now: Final = datetime.datetime.now()
result: Final = setup("aocr", (), _kwargs_with_a_budget_reservation(reservation), now, asynchronous=True)
token: Final = is_internal_call.set(True)
try:
pending: Final = failure_handler(result.logger, RuntimeError("inner step failed"), now, now, asynchronous=True)
finally:
is_internal_call.reset(token)
assert reservation["callback_bound"] is True
assert pending is not None
pending.close()
CONTRACT_PATH: Final = (
Path(__file__).parents[3] / "litellm-rust/crates/callbacks-legacy-python/python_contract.json"
)

94
tests/test_litellm/test_utils.py
View file

@ -4984,6 +4984,100 @@ async def test_wrapper_async_fires_post_call_failure_deployment_hook_on_internal
assert isinstance(recorder.calls[0][1], litellm.AuthenticationError)
def _budget_reservation(callback_bound: bool = False) -> dict:
return {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": callback_bound}
_BUDGET_RESERVATION_CALL_KWARGS: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
_BUDGET_RESERVATION_REFUSAL: Final = litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o")
@pytest.mark.asyncio
async def test_wrapper_async_claims_the_budget_reservation_for_the_cost_callback() -> None:
reservation = _budget_reservation()
await litellm.acompletion(
**_BUDGET_RESERVATION_CALL_KWARGS,
mock_response="ok",
metadata={"user_api_key_budget_reservation": reservation},
)
assert reservation["callback_bound"] is True
@pytest.mark.asyncio
async def test_wrapper_async_claims_the_budget_reservation_before_the_stream_is_consumed() -> None:
reservation = _budget_reservation()
stream = await litellm.acompletion(
**_BUDGET_RESERVATION_CALL_KWARGS,
mock_response="ok",
stream=True,
metadata={"user_api_key_budget_reservation": reservation},
)
assert reservation["callback_bound"] is True
async for _ in stream:
pass
@pytest.mark.asyncio
async def test_wrapper_async_claims_the_budget_reservation_a_supplied_logging_object_already_saw() -> None:
reservation = _budget_reservation()
logging_obj, kwargs = litellm.utils.function_setup(
original_function="acompletion",
rules_obj=litellm.utils.Rules(),
start_time=datetime.now(),
**_BUDGET_RESERVATION_CALL_KWARGS,
litellm_call_id="proxy-pre-call-setup",
metadata={"user_api_key_budget_reservation": reservation},
)
assert reservation["callback_bound"] is False
await litellm.acompletion(**kwargs, litellm_logging_obj=logging_obj, mock_response="ok")
assert reservation["callback_bound"] is True
@pytest.mark.asyncio
async def test_wrapper_async_hands_the_budget_reservation_back_when_the_call_fails() -> None:
reservation = _budget_reservation()
with pytest.raises(litellm.AuthenticationError):
await litellm.acompletion(
**_BUDGET_RESERVATION_CALL_KWARGS,
mock_response=_BUDGET_RESERVATION_REFUSAL,
metadata={"user_api_key_budget_reservation": reservation},
)
assert reservation["callback_bound"] is False
@pytest.mark.asyncio
async def test_wrapper_async_leaves_the_budget_reservation_alone_on_internal_calls() -> None:
claimed_by_the_outer_call = _budget_reservation(callback_bound=True)
never_claimed = _budget_reservation()
token = is_internal_call.set(True)
try:
await litellm.acompletion(
**_BUDGET_RESERVATION_CALL_KWARGS,
mock_response="ok",
metadata={"user_api_key_budget_reservation": never_claimed},
)
with pytest.raises(litellm.AuthenticationError):
await litellm.acompletion(
**_BUDGET_RESERVATION_CALL_KWARGS,
mock_response=_BUDGET_RESERVATION_REFUSAL,
metadata={"user_api_key_budget_reservation": claimed_by_the_outer_call},
)
finally:
is_internal_call.reset(token)
assert never_claimed["callback_bound"] is False
assert claimed_by_the_outer_call["callback_bound"] is True
@pytest.mark.asyncio
async def test_wrapper_async_does_not_fire_failure_hook_for_pre_call_budget_error(
monkeypatch: pytest.MonkeyPatch,