diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79edcde9f66..e215d55ebdf 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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, diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 6b39511d89b..7b1c13e2c09 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -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( diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 0aa2528d46b..6116b618c24 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -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: """ diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 40d80877056..27d6bc9389a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 05d5002b52c..7aa93c253d4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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, ) diff --git a/litellm/proxy/common_utils/request_tag_mappings.py b/litellm/proxy/common_utils/request_tag_mappings.py index dca269dca16..301544b698f 100644 --- a/litellm/proxy/common_utils/request_tag_mappings.py +++ b/litellm/proxy/common_utils/request_tag_mappings.py @@ -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) ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 91554b0c759..3afe66754e9 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -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) ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 2690c65c108..86f16411724 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 2ded84562fc..e5bb5011b15 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 194a7ba4f1b..0bde66b5cf8 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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. diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f3d7a9499dc..5d02078fdf9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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): diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index af68d7b242b..ecfc6b13edb 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -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: diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index b191f7835a8..92b29232271 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -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) diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index cc3e77f9c91..c39e8be7358 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -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( diff --git a/tests/unit/proxy/auth/test_handle_jwt.py b/tests/unit/proxy/auth/test_handle_jwt.py index 36ba8a734b8..418a719c1d8 100644 --- a/tests/unit/proxy/auth/test_handle_jwt.py +++ b/tests/unit/proxy/auth/test_handle_jwt.py @@ -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, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index 0b27b9a5f84..6a41c9d0601 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -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): diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 0f4b24a9ab9..b5dc6042474 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -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) diff --git a/tests/unit/proxy/common_utils/test_request_tag_mappings.py b/tests/unit/proxy/common_utils/test_request_tag_mappings.py index e8e7d881e1d..349d670bc79 100644 --- a/tests/unit/proxy/common_utils/test_request_tag_mappings.py +++ b/tests/unit/proxy/common_utils/test_request_tag_mappings.py @@ -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) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 1e5ad21fc5d..c574c8cded4 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -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), ], ) diff --git a/tests/unit/proxy/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index e11e30dfa76..1f197510dfe 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -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 diff --git a/tests/unit/proxy/proxy_server/test_streaming_helpers.py b/tests/unit/proxy/proxy_server/test_streaming_helpers.py index 4136a4acb16..dd61fba8526 100644 --- a/tests/unit/proxy/proxy_server/test_streaming_helpers.py +++ b/tests/unit/proxy/proxy_server/test_streaming_helpers.py @@ -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}, diff --git a/tests/unit/proxy/spend_tracking/test_budget_reservation.py b/tests/unit/proxy/spend_tracking/test_budget_reservation.py index a12e99d0ecd..b08b97ddccc 100644 --- a/tests/unit/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/unit/proxy/spend_tracking/test_budget_reservation.py @@ -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 diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index b3d83831578..8aa1742a2e7 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -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, diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 4386d5aec94..7a8c8b0e319 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -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, }