scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(proxy): run baggage tag mappings and budget checks on /queue/chat/completions (#51753)

Browse source 
* fix(proxy): run baggage tag mappings and budget checks on /queue/chat/completions

* fix(proxy): check the key's max_budget on /queue/chat/completions

* fix(proxy): run budget, tag, and request policy checks on queued chat calls

* fix(proxy): bill a dropped queued stream and run the trace ID and failure logging checks on queued calls

A queued stream now hands its logging object to the stream generator, so a
client that leaves mid-stream is billed for the reply it already got, and a
stream left before the first model chunk releases its budget reservation.
A team's require_trace_id and the failure logging for rejected calls now
use the same route check as the other inference checks.

* fix(proxy): settle a dropped queued stream's budget hold through the first-chunk step chat uses

* fix(proxy): keep the trace ID, call ID, unpriced fallback filter, and JWT admin team header on queued chat calls

* fix(proxy): send the call ID on a queued call the provider rejects before streaming

* fix(proxy): release a dropped stream's budget hold only when nothing bills it

A guardrail that buffers the stream can pull model chunks the client never
saw. The early-drop release settled the hold at the input cost first, then
the partial bill for those chunks added its full cost on top. The release
now runs after partial billing and only when no bill fired

* fix(proxy): fold the caller's litellm_metadata into a queued call's metadata

A queued call that sent a litellm_metadata bucket, even an empty one, made tag
routing and logging read that bucket, so the key and team tags the handler
merged into metadata no longer narrowed deployments. The queue handler now
folds the bucket into metadata and drops it, as chat does, after stripping the
proxy-owned user_api_key_* and internal control fields from both buckets, so a
caller can no longer set the queued spend row's organization ID either

* fix(proxy): release a queued call's budget hold when the call is cancelled before it answers

* fix(proxy): refund a cancelled queued call's hold in full until the model call claims it

* fix(proxy): keep a cancelled queued call's hold release running when the call is cancelled again

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
norkelbyte 2026-11-14 05:54:54 -08:00 • committed by GitHub
parent 27c796c412
commit ac1cc5111c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 822 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -904,7 +904,7 @@ def _enforce_user_param_check(general_settings: dict, request: Request, request_
http_method: Final = request.method if hasattr(request, "method") else None
is_post_method: Final = http_method and http_method.upper() == "POST"
is_openai_route: Final = RouteChecks.is_llm_api_route(route=route)
is_openai_route: Final = RouteChecks.runs_inference_checks(route=route)
is_mcp_route: Final = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
@ -931,7 +931,7 @@ def _reject_clientside_metadata_tags_check(
and "tags" in request_body["metadata"]
)
)
if RouteChecks.is_llm_api_route(route=route) and has_client_tags:
if RouteChecks.runs_inference_checks(route=route) and has_client_tags:
raise ProxyException(
message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via key or team metadata, or by configured request_tag_mappings.",
type=ProxyErrorTypes.bad_request_error,
@ -945,7 +945,7 @@ def _global_proxy_budget_check(global_proxy_spend: float | None, skip_budget_che
litellm.max_budget > 0
and not skip_budget_checks
and global_proxy_spend is not None
and RouteChecks.is_llm_api_route(route=route)
and RouteChecks.runs_inference_checks(route=route)
and route not in MODEL_LIST_ROUTES
):
if math.isfinite(litellm.max_budget) and global_proxy_spend > litellm.max_budget:
@ -1129,7 +1129,7 @@ def route_skips_budget_checks(route: str) -> bool:
return False
if route in MODEL_DISCOVERY_ROUTES:
return True
if RouteChecks.is_llm_api_route(route=route):
if RouteChecks.runs_inference_checks(route=route):
return False
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( # noqa: PLC0415 # pass-through endpoints import the auth stack
InitPassThroughEndpointHelpers,
@ -1230,7 +1230,7 @@ async def common_checks(
unpriced_models: Final = (
_unpriced_models_in_request(model=_model, llm_router=llm_router)
if litellm.block_requests_for_models_without_pricing and RouteChecks.is_llm_api_route(route=route)
if litellm.block_requests_for_models_without_pricing and RouteChecks.runs_inference_checks(route=route)
else ()
)
if unpriced_models:
@ -1322,7 +1322,7 @@ async def common_checks(
team_object is not None
and _team_metadata_requires_trace_id(team_object)
and request.method not in ("GET", "HEAD", "OPTIONS")
and (RouteChecks.is_llm_api_route(route=route) or pass_through_route)
and (RouteChecks.runs_inference_checks(route) or pass_through_route)
and not RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes.trace_telemetry_routes.value,

9
litellm/proxy/auth/handle_jwt.py
View file

@ -2277,12 +2277,13 @@ class JWTAuthManager:
) -> None:
"""Attach team context from x-litellm-team-id to an admin result.
Only applies on LLM API routes so team TPM/RPM limits and attribution
are enforced when admins act on behalf of a team. Admin management
routes ignore the header to preserve pre-existing bypass behavior.
Only applies on inference routes (LLM API routes and the queued chat
route) so team TPM/RPM limits and attribution are enforced when admins
act on behalf of a team. Admin management routes ignore the header to
preserve pre-existing bypass behavior.
"""
header_team_id: Final = request_headers.get("x-litellm-team-id") if request_headers else None
if not header_team_id or not RouteChecks.is_llm_api_route(route=route):
if not header_team_id or not RouteChecks.runs_inference_checks(route=route):
return
try:
team_object: Final = await get_team_object(

6
litellm/proxy/auth/route_checks.py
View file

@ -95,6 +95,8 @@ _SELF_SCOPED_CREDENTIAL_ROUTES: Final = (
)
_QUEUED_INFERENCE_ROUTES: Final = frozenset({"/queue/chat/completions"})
_UNBILLED_ROUTES: Final[frozenset[str]] = frozenset(
{
*MODEL_LIST_ROUTES,
@ -457,6 +459,10 @@ class RouteChecks:
resource, _, action = route.rsplit("/", 1)[-1].partition(":")
return route in _UNBILLED_ROUTES or resource in _TOKEN_COUNTING_SEGMENTS or action == _TOKEN_COUNTING_ACTION
@staticmethod
def runs_inference_checks(route: str) -> bool:
return route in _QUEUED_INFERENCE_ROUTES or RouteChecks.is_llm_api_route(route=route)
@staticmethod
def is_llm_api_route(route: str) -> bool:
"""

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

@ -1086,7 +1086,7 @@ async def _try_team_member_model_budget_fallback(
try:
if (
litellm.block_requests_for_models_without_pricing
and RouteChecks.is_llm_api_route(route=route)
and RouteChecks.runs_inference_checks(route=route)
and _unpriced_models_in_request(model=served_model, llm_router=llm_router)
):
continue
@ -3479,7 +3479,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
)
# Check 5. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
if RouteChecks.runs_inference_checks(route=route):
await virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,

27
litellm/proxy/common_request_processing.py
View file

@ -189,6 +189,7 @@ from litellm.proxy.route_llm_request import (
endpoint_unsupported_by_model,
route_request,
)
from litellm.proxy.spend_tracking.budget_reservation import release_budget_reservation_on_cancel
from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost
from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id
from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
@ -385,7 +386,7 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma
}
def _withheld_provider_output(response: object) -> bool:
def withheld_provider_output(response: object) -> bool:
return getattr(response, "has_buffered_provider_output", False) is True
@ -1331,18 +1332,18 @@ class _UpstreamClosingStreamingResponse(StreamingResponse):
def __init__(
self,
content: AsyncGenerator[str, None],
content: AsyncGenerator[str | bytes, None],
*,
media_type: str | None = None,
headers: Mapping[str, str] | None = None,
status_code: int = status.HTTP_200_OK,
upstream_generator: AsyncGenerator[str, None] | None = None,
upstream_generator: AsyncGenerator[str | bytes, None] | None = None,
) -> None:
super().__init__(content, status_code=status_code, headers=headers, media_type=media_type)
self._upstream_generator = upstream_generator
@property
def upstream_generator(self) -> AsyncGenerator[str, None] | None:
def upstream_generator(self) -> AsyncGenerator[str | bytes, None] | None:
"""The upstream LLM stream, for a caller that has to run this response's cleanup itself."""
return self._upstream_generator
@ -1613,7 +1614,7 @@ _NO_GENERAL_SETTINGS: Final[Mapping[str, object]] = MappingProxyType({})
async def create_response(
generator: AsyncGenerator[str, None] | Coroutine[object, object, AsyncGenerator[str, None]],
generator: AsyncGenerator[str | bytes, None] | Coroutine[object, object, AsyncGenerator[str | bytes, None]],
media_type: str,
headers: Mapping[str, str],
default_status_code: int = status.HTTP_200_OK,
@ -1631,7 +1632,7 @@ async def create_response(
callers whose headers can only be known then. ``general_settings`` decides whether
the first-chunk error body also carries the ``x-litellm-call-id`` header's value.
"""
first_chunk_value: str | None = None
first_chunk_value: str | bytes | None = None
final_status_code = default_status_code
try:
@ -1724,7 +1725,7 @@ async def create_response(
status_code=error_status,
)
async def combined_generator() -> AsyncGenerator[str, None]:
async def combined_generator() -> AsyncGenerator[str | bytes, None]:
if not _DD_STREAMING_TRACE_ENABLED:
# Fast path: no per-chunk span object / context-manager overhead.
if first_chunk_value is not None:
@ -4736,7 +4737,7 @@ class ProxyBaseLLMRequestProcessing:
# No spend row settles a stream that returned nothing it could log, so release what the request
# holds, as chat completions does.
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation_on_cancel, # pyright: ignore[reportUnknownVariableType] # takes the untyped reservation dict
release_budget_reservation_on_cancel,
)
await release_budget_reservation_on_cancel(getattr(user_api_key_dict, "budget_reservation", None))
@ -5398,6 +5399,8 @@ class ProxyBaseLLMRequestProcessing:
client_disconnected: bool = False,
user_api_key_dict: UserAPIKeyAuth | None = None,
proxy_logging_obj: ProxyLogging | None = None,
*,
release_reservation_if_unbilled: bool = False,
) -> None:
with anyio.CancelScope(shield=True):
if stream_completed and proxy_logging_obj is not None and user_api_key_dict is not None:
@ -5429,6 +5432,12 @@ class ProxyBaseLLMRequestProcessing:
and user_api_key_dict is not None
):
await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict)
if (
release_reservation_if_unbilled
and not success_event_owns_slot_release
and user_api_key_dict is not None
):
await release_budget_reservation_on_cancel(user_api_key_dict.budget_reservation)
if has_attributes(response, _ClosesAsync, "aclose"):
try:
@ -5580,7 +5589,7 @@ class ProxyBaseLLMRequestProcessing:
# only sees GeneratorExit on GC) cannot own the refund.
client_disconnected = not stream_completed
await close_guarded_stream(guarded_stream)
if not delivered_chunk and not _withheld_provider_output(response):
if not delivered_chunk and not withheld_provider_output(response):
from litellm.proxy.spend_tracking.budget_reservation import (
release_budget_reservation_on_cancel,
)

5
litellm/proxy/common_utils/request_tag_mappings.py
View file

@ -19,7 +19,10 @@ def is_baggage_mapped_route(route: str) -> bool:
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import InitPassThroughEndpointHelpers
return (
(RouteChecks.is_llm_api_route(route) or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route))
(
RouteChecks.runs_inference_checks(route)
or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route)
)
and not RouteChecks.is_unbilled_route(route)
and not RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value)
)

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

@ -497,7 +497,7 @@ class _ProxyDBLogger(CustomLogger):
route_is_trackable: Final = (
request_route is None
or RouteChecks.is_llm_api_route(route=request_route)
or RouteChecks.runs_inference_checks(request_route)
or RouteChecks.is_info_route(route=request_route)
or InitPassThroughEndpointHelpers.has_configured_route_ignoring_trailing_slash(route=request_route)
)

43
litellm/proxy/litellm_pre_call_utils.py
View file

@ -592,6 +592,11 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
def _is_proxy_owned_metadata_key(key: str) -> bool:
return key.startswith("user_api_key_") or key in _UNTRUSTED_METADATA_CONTROL_FIELDS
UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
{
"litellm-disable-message-redaction",
@ -2252,6 +2257,40 @@ class LiteLLMProxyRequestSetup:
project_tags: Final = (user_api_key_dict.project_metadata or {}).get("tags") or ()
return tuple(dict.fromkeys((*key_tags, *team_tags, *project_tags)))
@staticmethod
def caller_metadata(value: object) -> dict[str, object]:
if not isinstance(value, dict):
return {}
return {
key: field
for key, field in _REQUEST_FIELDS.validate_python(value).items()
if not _is_proxy_owned_metadata_key(key)
}
@staticmethod
def merged_request_tags(
request: Request,
metadata_buckets: tuple[Mapping[str, object], ...],
body_tags: object,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[str, ...]:
header: Final = request.headers.get("x-litellm-tags") or ""
body_string_tags: Final = _metadata_string_tags(({"tags": body_tags},))
caller_tags: Final = (
body_string_tags
if isinstance(body_tags, list)
else tuple(tag.strip() for tag in header.split(",") if tag.strip())
)
return tuple(
dict.fromkeys(
(
*_metadata_string_tags(metadata_buckets),
*LiteLLMProxyRequestSetup.get_inherited_tags(user_api_key_dict),
*caller_tags,
)
)
)
@staticmethod
def add_user_api_key_auth_to_request_metadata(
data: dict,
@ -3030,9 +3069,7 @@ async def add_litellm_data_to_request(
_user_meta.get("headers"),
allow_client_message_redaction_opt_out=(_allow_client_message_redaction_opt_out),
)
for _k in [
k for k in _user_meta if k.startswith("user_api_key_") or k in _UNTRUSTED_METADATA_CONTROL_FIELDS
]:
for _k in [k for k in _user_meta if isinstance(k, str) and _is_proxy_owned_metadata_key(k)]:
_user_meta.pop(_k, None)
# Strip pricing overrides AFTER the litellm_metadata string-to-dict parse

81
litellm/proxy/proxy_server.py
View file

@ -529,6 +529,7 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l
should_return_raw_model_name,
sse_error_payload,
ttft_keepalive_interval,
withheld_provider_output,
)
from litellm.proxy.common_utils.anthropic_error_payload import (
anthropic_error_json_response,
@ -1015,6 +1016,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
BudgetCounter,
customer_window_counters,
get_budget_window_start,
release_cancelled_budget_reservation,
release_unbound_budget_reservation,
stamp_budget_reservation_actual_cost,
)
@ -11020,6 +11022,7 @@ async def async_data_generator(
)
stream_source: AsyncIterator[object] | None = None # rebind-ok: bound once the keepalive policy resolves
stream_started: bool = False # rebind-ok: set once a ping or a hook-cleared chunk is sent, then errors go in-band
provider_output_started: bool = False # rebind-ok: set on the first provider chunk, which partial billing settles
try:
error_message: str | None = None
requested_model_from_client: Final = _get_client_requested_model_for_streaming(request_data=request_data)
@ -11083,6 +11086,7 @@ async def async_data_generator(
yield ": ping\n\n"
continue
chunk = item
provider_output_started = True
chunk_received_at = time.perf_counter()
if needs_per_chunk_hook:
### CALL HOOKS ### - modify outgoing data
@ -11274,6 +11278,9 @@ async def async_data_generator(
client_disconnected=client_disconnected,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
release_reservation_if_unbilled=client_disconnected
and not provider_output_started
and not withheld_provider_output(_UNTYPED_VALUE.validate_python(response)),
)
@ -19768,6 +19775,7 @@ async def async_queue_request(
Now using a FastAPI background task + /chat/completions compatible endpoint
"""
litellm_call_id: Final = resolve_litellm_call_id(request.headers.get(LITELLM_CALL_ID_HEADER))
data = {}
try:
data = await request.json()
@ -19797,10 +19805,14 @@ async def async_queue_request(
# if users are using user_api_key_auth, set `user` in `data`
data["user"] = user_api_key_dict.user_id
if not isinstance(data.get("metadata"), dict):
# Covers both missing and JSON-string metadata (multipart /
# extra_body); see above for the same guard upstream.
data["metadata"] = {}
data["litellm_call_id"] = litellm_call_id
caller_litellm_metadata: Final = LiteLLMProxyRequestSetup.caller_metadata(data.pop("litellm_metadata", None))
queued_metadata: Final = {
**caller_litellm_metadata,
**LiteLLMProxyRequestSetup.caller_metadata(data.get("metadata")),
}
data["metadata"] = queued_metadata
logged_api_key: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict)
data["metadata"]["user_api_key"] = logged_api_key
data["metadata"]["user_api_key_hash"] = logged_api_key
@ -19816,6 +19828,22 @@ async def async_queue_request(
user_api_key_dict, "team_object_permission_id", None
)
data["metadata"]["endpoint"] = str(request.url)
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=_headers, data=data, _metadata_variable_name="metadata"
)
queued_tags: Final = LiteLLMProxyRequestSetup.merged_request_tags(
request=request,
metadata_buckets=(queued_metadata, caller_litellm_metadata),
body_tags=data.get("tags"),
user_api_key_dict=user_api_key_dict,
)
if queued_tags:
data["metadata"]["tags"] = list(queued_tags)
data["metadata"]["inherited_tags"] = LiteLLMProxyRequestSetup.get_inherited_tags(user_api_key_dict)
data["metadata"]["user_api_key_budget_reservation"] = user_api_key_dict.budget_reservation
LiteLLMProxyRequestSetup.apply_authenticated_baggage_tags(
request=request, request_data=data, metadata_key="metadata"
)
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
# override with user settings, these are params passed via cli
@ -19838,18 +19866,31 @@ async def async_queue_request(
)
)
queued_priority, _ = router.health().split_priority(data["model"], data)
queued_headers: Final = {"x-litellm-priority": str(queued_priority), LITELLM_CALL_ID_HEADER: litellm_call_id}
async def schedule_queued_call() -> object:
try:
return _UNTYPED_VALUE.validate_python(await router.schedule_acompletion(**data))
except asyncio.CancelledError:
await release_cancelled_budget_reservation(user_api_key_dict.budget_reservation)
raise
if "stream" in data and data["stream"] is True: # use generate_responses to stream responses
async def produce_queue_stream() -> StreamingResponse:
return StreamingResponse(
async_data_generator(
async def produce_queue_stream() -> StreamingResponse | JSONResponse:
queued_stream: Final = await schedule_queued_call()
data["litellm_logging_obj"] = getattr(queued_stream, "logging_obj", None)
return await create_response(
generator=async_data_generator(
user_api_key_dict=user_api_key_dict,
response=_UNTYPED_VALUE.validate_python(await router.schedule_acompletion(**data)),
response=queued_stream,
request_data=data,
request=request,
),
media_type="text/event-stream",
headers={"x-litellm-priority": str(queued_priority)},
headers=queued_headers,
request=request,
general_settings=general_settings_view(),
)
async def audit_late_failure(exc: Exception) -> HTTPException | None:
@ -19863,20 +19904,26 @@ async def async_queue_request(
on_late_failure=audit_late_failure,
)
response: Final = _UNTYPED_VALUE.validate_python(await router.schedule_acompletion(**data))
fastapi_response.headers.update({"x-litellm-priority": str(queued_priority)})
response: Final = await schedule_queued_call()
fastapi_response.headers.update(queued_headers)
return response
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
)
if isinstance(e, litellm.BadRequestError):
raise ProxyException(
message=e.message, type="invalid_request_error", param=getattr(e, "param", None), code=e.status_code
) from e
if isinstance(e, ProxyException):
failure: Final = _queued_call_failure(e)
failure.headers = {**failure.headers, LITELLM_CALL_ID_HEADER: litellm_call_id}
if failure is e:
raise
raise proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST) from e
raise failure from e
def _queued_call_failure(e: Exception) -> ProxyException:
if isinstance(e, litellm.BadRequestError):
return ProxyException(
message=e.message, type="invalid_request_error", param=getattr(e, "param", None), code=e.status_code
)
return proxy_exception_for(e, default_status_code=status.HTTP_400_BAD_REQUEST)
@app.get("/fallback/login", tags=["experimental"], include_in_schema=False)

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

@ -334,7 +334,7 @@ async def reserve_budget_for_request(
raw_body: bytes | None = None,
path_params: Mapping[str, object] | None = None,
) -> dict | None:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
if valid_token is None or not RouteChecks.runs_inference_checks(route=route):
return None
if RouteChecks.is_unbilled_route(route):
return None
@ -472,7 +472,7 @@ async def release_budget_reservation(budget_reservation: dict | None) -> None:
async def release_budget_reservation_on_cancel(
budget_reservation: dict | None,
budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation
) -> None:
"""Reconcile a still-open reservation when the request is cancelled mid-flight.
@ -495,7 +495,8 @@ async def release_budget_reservation_on_cancel(
"""
if not budget_reservation or budget_reservation.get("finalized") is True:
return
incurred_cost: Final = float(budget_reservation.get("input_cost") or 0.0)
input_cost: Final = budget_reservation.get("input_cost")
incurred_cost: Final = float(input_cost) if isinstance(input_cost, (int, float)) else 0.0
try:
await asyncio.shield(
reconcile_budget_reservation(budget_reservation=budget_reservation, actual_cost=incurred_cost)
@ -543,6 +544,15 @@ async def release_or_invalidate_budget_reservation(
budget_reservation["finalized"] = True
async def release_cancelled_budget_reservation(
budget_reservation: dict[str, object] | None, # mutable-ok: stamps finalized on the caller's reservation
) -> None:
if budget_reservation is not None and budget_reservation.get("callback_bound") is True:
await release_budget_reservation_on_cancel(budget_reservation)
return
await asyncio.shield(release_or_invalidate_budget_reservation(budget_reservation))
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.

2
litellm/proxy/utils.py
View file

@ -4530,7 +4530,7 @@ class ProxyLogging:
#########################################################
if route is None:
return False
if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)):
if not (RouteChecks.runs_inference_checks(route) or RouteChecks.is_info_route(route)):
return False
# Dashboard session keys poll info routes in the background; their failures are not user requests.
if RouteChecks.is_dashboard_session_info_route(route, team_id):

50
tests/integration/spend/test_baggage_tag_mappings.py
View file

@ -21,6 +21,7 @@ from integration._support.client import (
gateway_from_environment,
list_value,
object_value,
string_value,
)
from integration._support.database import read_rows, scratch_database, write_rows
from integration._support.database_relay import dropped_connection_relay, plaintext_database_url
@ -563,6 +564,55 @@ def test_require_budgeted_tag_rejects_an_unbudgeted_mapped_tag_until_the_tag_has
]
def _queue_chat(rig: Rig, baggage: str) -> httpx.Response:
return rig.proxy.request(
"POST",
"/queue/chat/completions",
{
"model": rig.chat_model,
"messages": [{"role": "user", "content": f"queued {uuid.uuid4().hex}"}],
"priority": 0,
},
headers={"baggage": baggage},
)
def test_queued_chat_runs_the_baggage_mapper_like_chat_completions(mapped: Rig) -> None:
value: Final = uuid.uuid4().hex
mapped.provider.drain()
missing_required: Final = _queue_chat(mapped, "entrypoint=a")
assert missing_required.status_code == 400, missing_required.text
unbudgeted: Final = _queue_chat(mapped, f"budgeted={value},tenant={value}")
assert unbudgeted.status_code == 400, unbudgeted.text
assert mapped.provider.drain() == ()
admitted: Final = _queue_chat(mapped, f"entrypoint=a,tenant={value}")
assert admitted.status_code == 200, admitted.text
request_id: Final = string_value(JSON_OBJECT.validate_json(admitted.content)["id"])
rows: Final = eventually(
lambda: mapped.proxy.rows('SELECT request_tags FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)),
lambda values: len(values) == 1,
seconds=70,
)
assert _caller_tags(rows[0]["request_tags"]) == {mapped.exact_tag, mapped.tenant_tag(value)}
def test_spent_mapped_tag_budget_rejects_the_next_queued_request_before_the_provider(mapped: Rig) -> None:
entry: Final = uuid.uuid4().hex
tag: Final = mapped.usecase_tag(entry)
baggage: Final = f"entrypoint={entry},tenant={entry}"
with mapped.proxy.scenario() as scenario:
mapped.proxy.post("/tag/new", {"name": tag, "max_budget": 0.0001})
scenario.cleanups.callback(lambda: mapped.proxy.request("POST", "/tag/delete", {"name": tag}))
first: Final = _queue_chat(mapped, baggage)
assert first.status_code == 200, first.text
eventually(lambda: _queue_chat(mapped, baggage).status_code, lambda status: status != 200, seconds=70)
mapped.provider.drain()
blocked: Final = _queue_chat(mapped, baggage)
assert blocked.status_code == 422, blocked.text
assert "budget" in blocked.text.lower(), blocked.text
assert mapped.provider.drain() == ()
def _routing_caller_tag(rig: Rig, caller_tag: CallerTag) -> tuple[Mapping[str, JsonValue], Mapping[str, str]]:
json_metadata: Final[Mapping[str, JsonValue]] = {"metadata": json.dumps({"tags": [rig.exact_tag]})}
match caller_tag:

24
tests/unit/proxy/auth/test_auth_checks.py
View file

@ -109,7 +109,7 @@ async def _run_team_trace_id_check(
@pytest.mark.parametrize(
"route",
["/chat/completions", "/v1/messages", "/mcp/", "/a2a/my-agent"],
["/chat/completions", "/queue/chat/completions", "/v1/messages", "/mcp/", "/a2a/my-agent"],
)
@pytest.mark.asyncio
async def test_team_require_trace_id_rejects_llm_mcp_and_agent_requests(route: str) -> None:
@ -2469,3 +2469,25 @@ async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged(
assert usage_redis.finished
assert {json.loads(message)["cache_key"] for message in client.published} == set(hashed_tokens)
@pytest.mark.parametrize("route", ("/v1/chat/completions", "/queue/chat/completions"))
def test_chat_routes_require_the_user_param_when_enforced(route: str) -> None:
from starlette.requests import Request
from litellm.proxy.auth.auth_checks import _enforce_user_param_check
request: Final = Request({"type": "http", "method": "POST", "path": route})
with pytest.raises(Exception, match="'user' param not passed"):
_enforce_user_param_check({"enforce_user_param": True}, request, {}, route)
_enforce_user_param_check({"enforce_user_param": True}, request, {"user": "u"}, route)
@pytest.mark.parametrize("route", ("/v1/chat/completions", "/queue/chat/completions"))
def test_spent_proxy_budget_refuses_chat_routes(monkeypatch: pytest.MonkeyPatch, route: str) -> None:
from litellm.proxy.auth.auth_checks import _global_proxy_budget_check
monkeypatch.setattr(litellm, "max_budget", 0.5)
_global_proxy_budget_check(global_proxy_spend=0.4, skip_budget_checks=False, route=route)
with pytest.raises(litellm.BudgetExceededError):
_global_proxy_budget_check(global_proxy_spend=1.0, skip_budget_checks=False, route=route)

11
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -4306,7 +4306,8 @@ async def test_get_team_object_distinguishes_absent_team_from_unreadable_row():
@pytest.mark.asyncio
async def test_reject_clientside_metadata_tags_enabled_with_tags():
@pytest.mark.parametrize("route", ["/chat/completions", "/queue/chat/completions"])
async def test_reject_clientside_metadata_tags_enabled_with_tags(route):
"""Test that common_checks rejects request when reject_clientside_metadata_tags is True and metadata.tags is present."""
from fastapi import Request
@ -4334,7 +4335,7 @@ async def test_reject_clientside_metadata_tags_enabled_with_tags():
end_user_object=None,
global_proxy_spend=None,
general_settings=general_settings,
route="/chat/completions",
route=route,
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=valid_token,
@ -10544,12 +10545,13 @@ async def test_can_key_call_resolved_model_skips_customer_lookup_without_custome
@pytest.mark.asyncio
async def test_common_checks_blocks_unpriced_model_when_enabled(monkeypatch):
@pytest.mark.parametrize("route", ["/chat/completions", "/queue/chat/completions"])
async def test_common_checks_blocks_unpriced_model_when_enabled(monkeypatch, route):
monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True)
router = _router_with_priced_and_unpriced_models()
with pytest.raises(ProxyException) as exc_info:
await _run_common_checks(model="unpriced-group", llm_router=router)
await _run_common_checks(model="unpriced-group", llm_router=router, route=route)
assert exc_info.value.code == "403"
assert exc_info.value.type == ProxyErrorTypes.model_cost_map_missing
@ -11240,6 +11242,7 @@ def test_route_skips_budget_checks_marks_only_spend_free_routes() -> None:
assert RouteChecks.is_llm_api_route(route="/utils/model_info") is True
assert route_skips_budget_checks(route="/health") is False
assert route_skips_budget_checks(route="/v1/chat/completions") is False
assert route_skips_budget_checks(route="/queue/chat/completions") is False
@pytest.mark.parametrize(

6
tests/unit/proxy/auth/test_handle_jwt.py
View file

@ -2814,7 +2814,8 @@ async def test_auth_builder_rbac_team_denies_passthrough_without_allowlist():
@pytest.mark.asyncio
async def test_auth_builder_admin_on_llm_route_honors_team_header():
@pytest.mark.parametrize("route", ["/chat/completions", "/queue/chat/completions"])
async def test_auth_builder_admin_on_llm_route_honors_team_header(route: str):
"""JWT proxy_admin + x-litellm-team-id on an LLM API route -> team context is
attached to the admin result so team TPM/RPM limits and attribution apply."""
from litellm.caching import DualCache
@ -2832,6 +2833,7 @@ async def test_auth_builder_admin_on_llm_route_honors_team_header():
"management_routes",
"info_routes",
"openai_routes",
"/queue/chat/completions",
],
),
)
@ -2858,7 +2860,7 @@ async def test_auth_builder_admin_on_llm_route_honors_team_header():
jwt_handler=jwt_handler,
request_data={"model": "gpt-4"},
general_settings={},
route="/chat/completions",
route=route,
prisma_client=None,
user_api_key_cache=user_api_key_cache,
parent_otel_span=None,

40
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -541,6 +541,46 @@ async def test_aaauser_personal_budgets(key_ownership):
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
@pytest.mark.asyncio
@pytest.mark.parametrize("spend,expect_refusal", [(5.0, True), (0.5, False)], ids=["over_budget", "under_budget"])
async def test_virtual_key_max_budget_decides_a_queued_chat_call(monkeypatch, spend, expect_refusal):
import time
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.proxy_server import hash_token, user_api_key_cache
route: Final = "/queue/chat/completions"
key: Final = "sk-queued-chat-budget"
user_api_key_cache.set_cache(
key=hash_token(key),
value=UserAPIKeyAuth(
token=hash_token(key),
last_refreshed_at=time.time(),
allowed_routes=[route],
max_budget=1.0,
spend=spend,
),
)
monkeypatch.setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache)
monkeypatch.setattr(litellm.proxy.proxy_server, "master_key", MASTER_KEY)
monkeypatch.setattr(litellm.proxy.proxy_server, "prisma_client", _NoMembershipRowPrisma())
request: Final = Request(scope={"type": "http", "method": "POST", "path": route, "headers": []})
request._url = URL(url=route)
if not expect_refusal:
served: Final = await user_api_key_auth(request=request, api_key="Bearer " + key)
assert served.token == hash_token(key)
return
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer " + key)
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
@pytest.mark.asyncio
@pytest.mark.parametrize("prohibited_param", ["api_base", "base_url"])
async def test_user_api_key_auth_fails_with_prohibited_params(prohibited_param):

5
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -11262,7 +11262,8 @@ class TestCheckTeamMemberModelBudgetWithFallback:
assert request_data["model"] == "gpt-4o"
@pytest.mark.asyncio
async def test_fallback_skips_candidates_without_pricing(self, monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize("route", ["/v1/chat/completions", "/queue/chat/completions"])
async def test_fallback_skips_candidates_without_pricing(self, monkeypatch: pytest.MonkeyPatch, route: str):
valid_token: Final = UserAPIKeyAuth(
user_id="user-1",
team_id="team-1",
@ -11289,7 +11290,7 @@ class TestCheckTeamMemberModelBudgetWithFallback:
request=request,
llm_model_list=None,
llm_router=router,
route="/v1/chat/completions",
route=route,
)
_assert_budget_exceeded_values(exc_info.value)

9
tests/unit/proxy/common_utils/test_request_tag_mappings.py
View file

@ -269,7 +269,14 @@ async def test_budgeted_mapper_admits_a_tag_with_a_max_budget_and_unbudgeted_rul
@pytest.mark.parametrize(
"route",
["/v1/chat/completions", "/chat/completions", "/v1/messages", "/v1/responses", "/a2a/qa-agent/message/send"],
[
"/v1/chat/completions",
"/chat/completions",
"/queue/chat/completions",
"/v1/messages",
"/v1/responses",
"/a2a/qa-agent/message/send",
],
)
def test_inference_routes_take_baggage_mappings(route):
assert is_baggage_mapped_route(route)

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

@ -926,6 +926,7 @@ async def test_async_post_call_failure_hook_non_llm_route():
("/unregistered-search", CallTypes.pass_through.value, False),
("/azure-search/", CallTypes.pass_through.value, True),
("/configured-search/indexes/docs", CallTypes.pass_through.value, True),
("/queue/chat/completions", CallTypes.acompletion.value, True),
(None, CallTypes.acompletion.value, True),
],
)

341
tests/unit/proxy/proxy_server/test_routes_queue.py
View file

@ -6,7 +6,9 @@ Pins (PR2):
from __future__ import annotations
from collections.abc import Callable
import asyncio
import json
from collections.abc import AsyncGenerator, Awaitable, Callable
from contextlib import AbstractContextManager
from datetime import datetime, timezone
from queue import SimpleQueue
@ -15,12 +17,18 @@ from unittest.mock import AsyncMock, MagicMock, patch
import litellm
import pytest
from fastapi import Request, Response
from fastapi.responses import StreamingResponse
from fastapi.testclient import TestClient
from starlette.requests import ClientDisconnect
from starlette.types import Message
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy import proxy_server
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import ProxyLogging
from .conftest import normalize # type: ignore[import-not-found]
@ -307,7 +315,12 @@ def test_streamed_queue_responses_report_the_priority_the_router_queued_at(
client, auth_as, monkeypatch: pytest.MonkeyPatch
) -> None:
_queue_router(monkeypatch, default_priority=3)
payload: Final = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "priority": "1", "stream": True}
payload: Final = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"priority": "1",
"stream": True,
}
with auth_as():
response: Final = client.post("/queue/chat/completions", json=payload)
assert response.status_code == 200, response.text
@ -326,3 +339,327 @@ def test_a_provider_bad_request_on_the_queue_route_keeps_its_own_message(
assert error["type"] == "invalid_request_error"
assert "context window" in error["message"].lower()
assert "Authentication Error" not in error["message"]
@pytest.mark.parametrize(
"body_tags,caller_tag", [(None, "header-tag"), (["body-tag"], "body-tag")], ids=["header_tags", "body_tags"]
)
def test_queued_chat_carries_key_team_and_caller_tags_and_the_budget_reservation(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
patched_queue: MagicMock,
body_tags: list[str] | None,
caller_tag: str,
) -> None:
reservation: Final = {"reserved_cost": 0.01, "callback_bound": False}
payload: Final = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"metadata": {"tags": ["sent"]},
**({"tags": body_tags} if body_tags is not None else {}),
}
with auth_as(metadata={"tags": ["key-tag"]}, team_metadata={"tags": ["team-tag"]}, budget_reservation=reservation):
response: Final = client.post("/queue/chat/completions", json=payload, headers={"x-litellm-tags": "header-tag"})
assert response.status_code == 200, response.text
metadata: Final = patched_queue.schedule_acompletion.call_args.kwargs["metadata"]
assert metadata["tags"] == ["sent", "key-tag", "team-tag", caller_tag]
assert tuple(metadata["inherited_tags"]) == ("key-tag", "team-tag")
assert metadata["user_api_key_budget_reservation"] == reservation
@pytest.mark.parametrize(
("key_metadata", "litellm_metadata"),
[({"tags": ["key-tag"]}, {}), ({}, {"tags": ["key-tag"]})],
ids=["key_tags_with_empty_bucket", "bucket_tags"],
)
def test_a_queued_call_with_a_litellm_metadata_bucket_still_routes_by_its_tags(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
key_metadata: dict[str, object],
litellm_metadata: dict[str, object],
) -> None:
deployments: Final = [
{
"model_name": "gpt-4",
"litellm_params": {
"model": "openai/gpt-4",
"api_key": "sk-queue-test",
"mock_response": reply,
"tags": [reply],
},
}
for reply in ("key-tag", "default")
]
monkeypatch.setattr(
proxy_server, "llm_router", Router(model_list=deployments, default_priority=0, enable_tag_filtering=True)
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock()))
with auth_as(metadata=key_metadata):
response: Final = client.post(
"/queue/chat/completions",
json={
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"litellm_metadata": litellm_metadata,
},
)
assert response.status_code == 200, response.text
assert response.json()["choices"][0]["message"]["content"] == "key-tag"
def test_a_queued_spend_row_ignores_key_fields_forged_in_the_litellm_metadata_bucket(
client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], monkeypatch: pytest.MonkeyPatch
) -> None:
_queue_router(monkeypatch, default_priority=0)
logger: Final = _QueueRequestLogger()
monkeypatch.setattr(litellm, "callbacks", [logger])
with auth_as(team_id="real-team"):
response: Final = client.post(
"/queue/chat/completions",
json={
"model": "gpt-4",
"messages": [{"role": "user", "content": "hi"}],
"litellm_metadata": {"user_api_key_team_id": "forged-team", "user_api_key_org_id": "forged-org"},
},
)
assert response.status_code == 200, response.text
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
payload: Final = get_logging_payload(
kwargs=logger.requests.get_nowait(),
response_obj=litellm.ModelResponse.model_validate(response.json()),
start_time=now,
end_time=now,
)
assert payload["team_id"] == "real-team"
assert payload["organization_id"] != "forged-org"
def test_a_queued_spend_row_stores_the_trace_id_the_caller_sent(
client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], monkeypatch: pytest.MonkeyPatch
) -> None:
_queue_router(monkeypatch, default_priority=0)
logger: Final = _QueueRequestLogger()
monkeypatch.setattr(litellm, "callbacks", [logger])
with auth_as():
response: Final = client.post(
"/queue/chat/completions",
json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]},
headers={"x-litellm-trace-id": "queue-trace-1"},
)
assert response.status_code == 200, response.text
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
payload: Final = get_logging_payload(
kwargs=logger.requests.get_nowait(),
response_obj=litellm.ModelResponse.model_validate(response.json()),
start_time=now,
end_time=now,
)
assert payload["session_id"] == "queue-trace-1"
def test_a_queued_stream_error_carries_the_call_id_in_its_header_and_body(
client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], monkeypatch: pytest.MonkeyPatch
) -> None:
async def provider_rate_limited() -> AsyncGenerator[litellm.ModelResponseStream, None]:
raise litellm.RateLimitError(message="slow down", llm_provider="openai", model="gpt-4")
yield litellm.ModelResponseStream()
router: Final = MagicMock()
router.health.return_value.split_priority.return_value = (0, {})
router.health.return_value.priority_header_fields.return_value = {}
router.registry.return_value.deployments.return_value = []
router.schedule_acompletion = AsyncMock(return_value=provider_rate_limited())
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock()))
monkeypatch.setattr(proxy_server, "general_settings", {"include_call_id_in_error_body": True})
with auth_as():
response: Final = client.post(
"/queue/chat/completions",
json={"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "stream": True},
headers={"x-litellm-call-id": "queue-call-1"},
)
assert response.status_code == 429, response.text
assert response.headers.get("x-litellm-call-id") == "queue-call-1"
assert response.json()["error"]["litellm_call_id"] == "queue-call-1"
@pytest.mark.parametrize("stream", [True, False])
def test_a_queued_call_the_provider_rejects_carries_the_call_id_in_its_header_and_body(
client: TestClient,
auth_as: Callable[..., AbstractContextManager[None]],
monkeypatch: pytest.MonkeyPatch,
stream: bool,
) -> None:
router: Final = MagicMock()
router.health.return_value.split_priority.return_value = (0, {})
router.health.return_value.priority_header_fields.return_value = {}
router.registry.return_value.deployments.return_value = []
router.schedule_acompletion = AsyncMock(
side_effect=litellm.NotFoundError(message="no such model", llm_provider="gemini", model="gemini-missing")
)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock()))
monkeypatch.setattr(proxy_server, "general_settings", {"include_call_id_in_error_body": True})
with auth_as():
response: Final = client.post(
"/queue/chat/completions",
json={"model": "gemini-missing", "messages": [{"role": "user", "content": "hi"}], "stream": stream},
headers={"x-litellm-call-id": "queue-call-2"},
)
assert response.status_code == 404, response.text
assert response.headers.get("x-litellm-call-id") == "queue-call-2"
assert response.json()["error"]["litellm_call_id"] == "queue-call-2"
class _QueueSuccessLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.responses: Final = SimpleQueue[object]()
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.responses.put(response_obj)
_QUEUED_STREAM_SCOPE: Final = {
"type": "http",
"asgi": {"version": "3.0", "spec_version": "2.4"},
"http_version": "1.1",
"method": "POST",
"scheme": "http",
"server": ("testserver", 80),
"path": "/queue/chat/completions",
"root_path": "",
"headers": [(b"content-type", b"application/json")],
"query_string": b"",
}
def _queued_stream_request(after_body: Callable[[], Awaitable[Message]]) -> Request:
body: Final = json.dumps({"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "stream": True})
body_read: Final = asyncio.Event()
async def receive() -> Message:
if body_read.is_set():
return await after_body()
body_read.set()
return {"type": "http.request", "body": body.encode(), "more_body": False}
return Request(_QUEUED_STREAM_SCOPE, receive)
async def _still_connected() -> Message:
await asyncio.Event().wait()
return {"type": "http.disconnect"}
async def _already_gone() -> Message:
return {"type": "http.disconnect"}
async def test_a_queued_stream_the_client_abandons_bills_the_reply_it_already_streamed(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_queue_router(monkeypatch, default_priority=0, mock_response="queued partial reply")
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock()))
logger: Final = _QueueSuccessLogger()
monkeypatch.setattr(litellm, "callbacks", [logger])
request: Final = _queued_stream_request(_still_connected)
response: Final = await proxy_server.async_queue_request(
request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-queue")
)
assert isinstance(response, StreamingResponse)
delivered: Final = SimpleQueue[bytes]()
async def send_until_the_client_hangs_up(message: Message) -> None:
if message["type"] != "http.response.body":
return
if delivered.qsize() == 2:
raise OSError("client hung up")
delivered.put(bytes(message["body"]))
with pytest.raises(ClientDisconnect):
await response(_QUEUED_STREAM_SCOPE, request.receive, send_until_the_client_hangs_up)
billed: Final = logger.responses.get_nowait()
assert isinstance(billed, litellm.ModelResponse)
billed_reply: Final = billed.choices[0].message.content
assert isinstance(billed_reply, str)
assert billed_reply
assert "queued partial reply".startswith(billed_reply)
assert billed_reply != "queued partial reply"
async def test_a_queued_stream_the_client_leaves_before_its_first_chunk_gives_back_its_budget_hold(
monkeypatch: pytest.MonkeyPatch,
) -> None:
provider_closed: Final = asyncio.Event()
async def provider_still_thinking() -> AsyncGenerator[litellm.ModelResponseStream, None]:
try:
await asyncio.Event().wait()
yield litellm.ModelResponseStream()
finally:
provider_closed.set()
router: Final = MagicMock()
router.health.return_value.split_priority.return_value = (0, {})
router.health.return_value.priority_header_fields.return_value = {}
router.registry.return_value.deployments.return_value = []
router.schedule_acompletion = AsyncMock(return_value=provider_still_thinking())
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock()))
auth: Final = UserAPIKeyAuth(
api_key="sk-queue", budget_reservation={"reserved_cost": 0.01, "input_cost": 0.0, "entries": []}
)
response: Final = await proxy_server.async_queue_request(
request=_queued_stream_request(_already_gone), fastapi_response=Response(), user_api_key_dict=auth
)
assert auth.budget_reservation is not None
assert auth.budget_reservation.get("finalized") is True
assert provider_closed.is_set()
assert not isinstance(response, StreamingResponse)
@pytest.mark.parametrize("stream", [False, True])
async def test_a_queued_call_cancelled_before_it_answers_gives_back_its_budget_hold(
monkeypatch: pytest.MonkeyPatch, stream: bool
) -> None:
waiting: Final = asyncio.Event()
async def wait_for_a_worker(**_: object) -> object:
waiting.set()
await asyncio.Event().wait()
return dict(HAPPY_RESPONSE)
router: Final = MagicMock()
router.health.return_value.split_priority.return_value = (0, {})
router.health.return_value.priority_header_fields.return_value = {}
router.registry.return_value.deployments.return_value = []
router.schedule_acompletion = wait_for_a_worker
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=MagicMock()))
auth: Final = UserAPIKeyAuth(
api_key="sk-queue", budget_reservation={"reserved_cost": 0.01, "input_cost": 0.0, "entries": []}
)
body: Final = json.dumps({"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}], "stream": stream})
async def receive() -> Message:
return {"type": "http.request", "body": body.encode(), "more_body": False}
queued: Final = asyncio.ensure_future(
proxy_server.async_queue_request(
request=Request(_QUEUED_STREAM_SCOPE, receive), fastapi_response=Response(), user_api_key_dict=auth
)
)
await waiting.wait()
queued.cancel()
with pytest.raises(asyncio.CancelledError):
await queued
assert auth.budget_reservation is not None
assert auth.budget_reservation.get("finalized") is True

87
tests/unit/proxy/proxy_server/test_streaming_helpers.py
View file

@ -1148,6 +1148,93 @@ async def test_async_data_generator_yields_sse_chunks_and_done(monkeypatch):
}
@pytest.mark.asyncio
@pytest.mark.parametrize(("provider_chunks", "released"), [(0, True), (1, False)], ids=["before_output", "after_output"])
async def test_async_data_generator_refunds_the_reservation_only_when_the_client_leaves_before_provider_output(
monkeypatch: pytest.MonkeyPatch, provider_chunks: int, released: bool
) -> None:
_patch_logging_flags(monkeypatch)
waiting_on_provider: Final = asyncio.Event()
async def provider() -> AsyncIterator[ModelResponseStream]:
for _ in range(provider_chunks):
yield _simple_chunk(content="partial")
waiting_on_provider.set()
await asyncio.Event().wait()
yield _simple_chunk(content="never sent")
auth: Final = UserAPIKeyAuth(
api_key="sk-test-key", budget_reservation={"reserved_cost": 0.01, "input_cost": 0.0, "entries": []}
)
generator: Final = async_data_generator(response=provider(), user_api_key_dict=auth, request_data={"model": "gpt-4"})
for _ in range(provider_chunks):
await generator.__anext__()
next_frame: Final = asyncio.ensure_future(generator.__anext__())
await waiting_on_provider.wait()
next_frame.cancel()
with pytest.raises(asyncio.CancelledError):
await next_frame
assert auth.budget_reservation is not None
assert (auth.budget_reservation.get("finalized") is True) is released
class _PartialBillingRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.billed_texts: Final[list[str]] = [] # mutable-ok: collects the billing events the stream dispatches
self.billed: Final = asyncio.Event()
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.billed_texts.append(litellm.get_response_string(response_obj))
self.billed.set()
@pytest.mark.asyncio
async def test_a_stream_a_guardrail_buffered_leaves_its_reservation_to_the_partial_billing_when_the_client_leaves(
monkeypatch: pytest.MonkeyPatch,
) -> None:
_patch_logging_flags(monkeypatch, needs_wrap=True)
recorder: Final = _PartialBillingRecorder()
provider_stream: Final = await litellm.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "count"}],
mock_response="one two three four five six seven eight",
stream=True,
api_key="test-key",
)
provider_stream.logging_obj.dynamic_async_success_callbacks = [recorder]
guardrail_holding_chunks: Final = asyncio.Event()
async def buffering_guardrail(user_api_key_dict, response, request_data) -> AsyncIterator[object]:
held: Final = [await response.__anext__() for _ in range(3)]
guardrail_holding_chunks.set()
await asyncio.Event().wait()
for chunk in held:
yield chunk
monkeypatch.setattr(ps.proxy_logging_obj, "async_post_call_streaming_iterator_hook", buffering_guardrail)
auth: Final = UserAPIKeyAuth(
api_key="sk-test-key", budget_reservation={"reserved_cost": 0.01, "input_cost": 0.004, "entries": []}
)
generator: Final = async_data_generator(
response=provider_stream,
user_api_key_dict=auth,
request_data={"model": "gpt-4o-mini", "litellm_logging_obj": provider_stream.logging_obj},
)
first_frame: Final = asyncio.ensure_future(generator.__anext__())
await guardrail_holding_chunks.wait()
first_frame.cancel()
with pytest.raises(asyncio.CancelledError):
await first_frame
await asyncio.wait_for(recorder.billed.wait(), timeout=5)
assert recorder.billed_texts == ["".join(chunk.choices[0].delta.content or "" for chunk in provider_stream.chunks)]
assert recorder.billed_texts[0]
assert auth.budget_reservation is not None
assert auth.budget_reservation.get("finalized") is not True
_AZURE_CONTENT_FILTER_RESULTS: Final = {
"hate": {"filtered": False, "severity": "safe"},
"protected_material_text": {"detected": False, "filtered": False},

5
tests/unit/proxy/spend_tracking/test_budget_reservation.py
View file

@ -111,8 +111,9 @@ async def test_token_counting_routes_are_exempt_from_budget_reservation(route):
@pytest.mark.asyncio
async def test_non_exempt_llm_route_still_reserves_budget():
reservation: Final = await _reserve("/v1/responses")
@pytest.mark.parametrize("route", ["/v1/responses", "/queue/chat/completions"])
async def test_non_exempt_llm_route_still_reserves_budget(route):
reservation: Final = await _reserve(route)
assert reservation is not None
assert reservation["reserved_cost"] > 0

91
tests/unit/proxy/test_budget_reservation.py
View file

@ -45,6 +45,7 @@ from litellm.proxy.spend_tracking.budget_reservation import (
invalidate_budget_reservation_counters,
release_budget_reservation,
release_budget_reservation_on_cancel,
release_cancelled_budget_reservation,
reserve_budget_for_request,
)
from litellm.proxy.spend_tracking.input_tokens import (
@ -2658,6 +2659,96 @@ async def test_release_budget_reservation_on_cancel_gives_back_counter(
assert counter_cache.get_local("spend:key:key-cancel-give-back") == pytest.approx(0.5)
@pytest.mark.asyncio
@pytest.mark.parametrize(("callback_bound", "kept_cost"), [(False, 0.0), (True, 0.5)])
async def test_a_cancelled_call_keeps_its_input_cost_only_once_the_model_call_claimed_its_hold(
spend_counter_state, callback_bound: bool, kept_cost: float
):
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
token = f"key-cancelled-bound-{callback_bound}"
valid_token = UserAPIKeyAuth(token=token, spend=0.0, max_budget=10.0)
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=3.0,
),
patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_input_cost",
return_value=0.5,
),
):
reservation = await reserve_budget_for_request(
request_body=_request_body(),
route="/queue/chat/completions",
llm_router=None,
valid_token=valid_token,
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert reservation is not None
assert counter_cache.get_local(f"spend:key:{token}") == pytest.approx(3.0)
reservation["callback_bound"] = callback_bound
await release_cancelled_budget_reservation(reservation)
assert counter_cache.get_local(f"spend:key:{token}") == pytest.approx(kept_cost)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_an_unclaimed_hold_is_dropped_when_its_release_fails_after_the_call_is_cancelled_again(
spend_counter_state,
):
counter_cache, key_cache = spend_counter_state
token: Final = "key-cancelled-twice"
with patch(
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
return_value=3.0,
):
reservation: Final = await reserve_budget_for_request(
request_body=_request_body(),
route="/queue/chat/completions",
llm_router=None,
valid_token=UserAPIKeyAuth(token=token, spend=0.0, max_budget=10.0),
team_object=None,
user_object=None,
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=ProxyLogging(user_api_key_cache=key_cache),
)
assert reservation is not None
assert counter_cache.get_local(f"spend:key:{token}") == pytest.approx(3.0)
release_started: Final = asyncio.Event()
store_dropped: Final = asyncio.Event()
async def release_until_the_store_drops(**_: object) -> tuple[()]:
release_started.set()
await store_dropped.wait()
raise ConnectionError("counter store unreachable")
async def settled() -> None:
while reservation.get("finalized") is not True:
await asyncio.sleep(0)
with patch(
"litellm.proxy.spend_tracking.budget_reservation._set_reserved_entries_actual_cost",
side_effect=release_until_the_store_drops,
):
release: Final = asyncio.ensure_future(release_cancelled_budget_reservation(reservation))
await release_started.wait()
release.cancel()
await asyncio.wait([release])
store_dropped.set()
await asyncio.wait_for(settled(), timeout=1)
assert counter_cache.get_local(f"spend:key:{token}") is None
@pytest.mark.asyncio
async def test_release_budget_reservation_on_cancel_noop_when_finalized(
spend_counter_state,

6
tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py
View file

@ -54,6 +54,11 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging):
error_type=ProxyErrorTypes.auth_error,
route="/chat/completions",
),
"auth_short_circuit_on_queued_route": proxy_logging._is_proxy_only_llm_api_error(
original_exception=Exception("auth"),
error_type=ProxyErrorTypes.auth_error,
route="/queue/chat/completions",
),
"guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error(
original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"),
route="/chat/completions",
@ -64,6 +69,7 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging):
"non_llm_route": False,
"http_on_llm_route": True,
"auth_short_circuit": True,
"auth_short_circuit_on_queued_route": True,
"guardrail_raised_on_llm_route": True,
}