From 43112547fb9c0a8d482f4fe031e09a6c2cc234e1 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sun, 8 Nov 2026 00:43:50 -0800 Subject: [PATCH] feat(proxy): total_timeout bounds a whole non-streaming request (#50496) * feat(proxy): total_timeout bounds a whole non-streaming request A `total_timeout` body field or `x-litellm-total-timeout` header (the header wins) gives a non-streaming /v1/chat/completions, /v1/messages or /v1/responses request one budget for preparation, retries, moderation and post-call processing. Past it the proxy answers 408 and cancels the upstream call, and a retry whose backoff cannot fit ends early with the error it would have retried. A body value on a streaming or background request is refused with a 400, a header value leaves it unbounded. A timed out request is billed like a client disconnect. Co-Authored-By: Claude Opus 5.5 * fix(proxy): ignore any total_timeout header on requests it cannot bound, and never interrupt pre-call admission A globally configured x-litellm-total-timeout that is malformed no longer fails the streaming and background requests it is ignored on. The deadline holds while pre-call admission runs, since admission reserves shared quota before it records it, and an expired deadline then ends the request before dispatch, where the failure path hands the reservation back. Co-Authored-By: Claude Opus 5.5 * fix(proxy): a total_timeout cuts slow pre-call hooks short, holding only the rate limit reservation Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- litellm/litellm_core_utils/total_deadline.py | 197 ++++++++++ litellm/proxy/common_request_processing.py | 174 +++++++-- .../hooks/parallel_request_limiter_v3.py | 2 + .../proxy/response_api_endpoints/endpoints.py | 3 + litellm/proxy/route_llm_request.py | 2 + litellm/proxy/utils.py | 12 +- litellm/router_rust/core.py | 20 + .../litellm_core_utils/test_total_deadline.py | 36 ++ .../hooks/test_parallel_request_limiter_v3.py | 30 ++ .../proxy/search_endpoints/test_endpoints.py | 1 + .../proxy/test_common_request_processing.py | 351 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 103 +++-- ..._utils_model_creation_and_error_logging.py | 156 +++++--- .../test_response_polling_pre_call_checks.py | 78 +++- tests/unit/proxy/test_route_llm_request.py | 112 ++++-- .../core/test_retries_cooldowns.py | 55 +++ 16 files changed, 1170 insertions(+), 162 deletions(-) create mode 100644 litellm/litellm_core_utils/total_deadline.py create mode 100644 tests/unit/litellm_core_utils/test_total_deadline.py diff --git a/litellm/litellm_core_utils/total_deadline.py b/litellm/litellm_core_utils/total_deadline.py new file mode 100644 index 00000000000..f44eb97719d --- /dev/null +++ b/litellm/litellm_core_utils/total_deadline.py @@ -0,0 +1,197 @@ +""" +One wall-clock budget for a whole non-streaming proxy request (``total_timeout``): the proxy opens it, +and the router reads it to refuse a retry backoff that cannot finish in time. A request that turns out +to stream or run in the background is not bounded: a body ``total_timeout`` refuses it, a header one +(often set globally on a client) lets it go unbounded. +""" + +import asyncio +import contextlib +import math +import sys +from collections.abc import AsyncGenerator, Callable, Generator, Mapping +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Final + +from litellm._logging import verbose_router_logger + +if sys.version_info >= (3, 11): + from asyncio import Timeout as _Scope + from asyncio import timeout_at as _timeout_at + + _ScopeTimeout: Final = TimeoutError + + def _expired(scope: _Scope) -> bool: + return scope.expired() + + def _disarm(scope: _Scope) -> None: + scope.reschedule(None) + + def _rearm(scope: _Scope, when: float) -> None: + scope.reschedule(when) + +else: + from async_timeout import Timeout as _Scope + from async_timeout import timeout_at as _timeout_at + + _ScopeTimeout: Final = asyncio.TimeoutError + + def _expired(scope: _Scope) -> bool: + return scope.expired + + def _disarm(scope: _Scope) -> None: + scope.reject() + + def _rearm(scope: _Scope, when: float) -> None: + scope.update(when) + + +TOTAL_TIMEOUT_HEADER: Final = "x-litellm-total-timeout" +TOTAL_TIMEOUT_FIELD: Final = "total_timeout" + + +@dataclass(frozen=True, slots=True) +class TotalTimeout: + seconds: float + from_body: bool + + +@dataclass(frozen=True, slots=True) +class InvalidTotalTimeout: + message: str + + +def _positive_finite_seconds(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float | str): + return None + try: + seconds: Final = float(value) + except (ValueError, OverflowError): + return None + return seconds if math.isfinite(seconds) and seconds > 0 else None + + +def _total_timeout(value: object, *, from_body: bool) -> TotalTimeout | InvalidTotalTimeout: + seconds: Final = _positive_finite_seconds(value) + if seconds is None: + name: Final = TOTAL_TIMEOUT_FIELD if from_body else TOTAL_TIMEOUT_HEADER + return InvalidTotalTimeout(f"{name} must be a positive, finite number of seconds, got {value!r}") + return TotalTimeout(seconds, from_body) + + +def total_timeout_from_request( + header: str | None, body: Mapping[str, object] +) -> TotalTimeout | InvalidTotalTimeout | None: + """The request's ``total_timeout``: the header wins over the body field.""" + if header is not None: + return _total_timeout(header, from_body=False) + if TOTAL_TIMEOUT_FIELD not in body: + return None + return _total_timeout(body[TOTAL_TIMEOUT_FIELD], from_body=True) + + +def bounded_by_total_timeout(call: Mapping[str, object]) -> bool: + """Only a request that neither streams nor runs in the background can be bounded.""" + return call.get("stream") is not True and call.get("background") is not True + + +class TotalTimeoutNotApplicable(Exception): + """A request bounded by a body ``total_timeout`` streams or runs in the background.""" + + +class TotalTimeoutExpired(TimeoutError): + """The request's ``total_timeout`` ran out.""" + + +class TotalDeadline: + __slots__ = ("_at", "_clock", "_released", "_scope", "_scope_at", "from_body") + + def __init__( + self, at: float, clock: Callable[[], float], scope: _Scope, scope_at: float, *, from_body: bool + ) -> None: + self._at: Final = at + self._clock: Final = clock + self._scope: Final = scope + self._scope_at: Final = scope_at + self._released: bool = False + self.from_body: Final = from_body + + def expired(self) -> bool: + return _expired(self._scope) + + def allows_wait(self, seconds: float) -> bool: + return self._released or self._clock() + seconds < self._at + + @contextlib.contextmanager + def held(self) -> Generator[None, None, None]: + """Run the block without interrupting it; a deadline that passed meanwhile fires right after it.""" + if self._released: + yield + return + _disarm(self._scope) + try: + yield + finally: + if not self._released: + _rearm(self._scope, self._scope_at) + + def release(self) -> None: + self._released = True + if not self.expired(): + _disarm(self._scope) + + +_CURRENT: Final[ContextVar[TotalDeadline | None]] = ContextVar("total_deadline", default=None) + + +def current_total_deadline() -> TotalDeadline | None: + return _CURRENT.get() + + +@contextlib.asynccontextmanager +async def total_timeout_held() -> AsyncGenerator[None, None]: + deadline: Final = _CURRENT.get() + if deadline is None: + yield + return + with deadline.held(): + yield + + +def enforce_total_timeout_scope(call: Mapping[str, object]) -> None: + """Raise ``TotalTimeoutNotApplicable`` for a body ``total_timeout`` on a streaming or background call, + and stop bounding it for a header one.""" + deadline: Final = _CURRENT.get() + if deadline is None or deadline._released or bounded_by_total_timeout(call): # pyright: ignore[reportPrivateUsage] # same module + return + if deadline.from_body: + raise TotalTimeoutNotApplicable( + "total_timeout applies only to non-streaming, non-background requests; drop it or the stream/background flag" + ) + verbose_router_logger.debug("x-litellm-total-timeout ignored: the request streams or runs in the background") + deadline.release() + + +@contextlib.asynccontextmanager +async def total_deadline( + seconds: float, *, from_body: bool, clock: Callable[[], float] | None = None +) -> AsyncGenerator[TotalDeadline, None]: + """Bound the block by ``seconds``: past it the block is cancelled and ``TotalTimeoutExpired`` raised, + while a ``TimeoutError`` raised inside the block passes through as it was.""" + loop: Final = asyncio.get_running_loop() + now: Final = clock if clock is not None else loop.time + scope_at: Final = loop.time() + seconds + scope: Final = _timeout_at(scope_at) + deadline: Final = TotalDeadline(now() + seconds, now, scope, scope_at, from_body=from_body) + try: + async with scope: + token: Final = _CURRENT.set(deadline) + try: + yield deadline + finally: + _CURRENT.reset(token) + except _ScopeTimeout as error: + if not deadline.expired(): + raise + raise TotalTimeoutExpired(f"the request did not finish within its total_timeout of {seconds}s") from error diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 887b8f98c2b..8f73d47d96b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -174,6 +174,7 @@ _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") _ResponseT: Final = TypeVar("_ResponseT") _ChunkT: Final = TypeVar("_ChunkT") +_ProcessedT: Final = TypeVar("_ProcessedT") KNOWN_PROXY_ROUTES: Final = frozenset( route for member in LiteLLMRoutes for route in member.value if route.startswith("/") @@ -288,6 +289,18 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.litellm_core_utils.total_deadline import ( + TOTAL_TIMEOUT_FIELD, + TOTAL_TIMEOUT_HEADER, + InvalidTotalTimeout, + TotalTimeout, + TotalTimeoutExpired, + TotalTimeoutNotApplicable, + bounded_by_total_timeout, + current_total_deadline, + total_deadline, + total_timeout_from_request, +) from litellm.proxy.anthropic_endpoints.answered_status import with_answered_status from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( AnthropicStreamModelRestamper, @@ -583,6 +596,31 @@ async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, par return True +_TOTAL_TIMEOUT_ROUTE_TYPES: Final = frozenset({"acompletion", "aresponses", "anthropic_messages"}) + + +def _total_timeout_refused(message: str) -> ProxyException: + return ProxyException(message=message, type="invalid_request_error", param=TOTAL_TIMEOUT_FIELD, code=400) + + +def _total_timeout_expired() -> bool: + deadline: Final = current_total_deadline() + return deadline is not None and deadline.expired() + + +async def _cost_a_call_answered_past_its_total_timeout( + llm_call: "asyncio.Task[object]", logging_obj: LiteLLMLoggingObj +) -> None: + """The provider answered but the request ran out of time after it (a slower moderation check): the + failure it is logged as still carries what the provider charged.""" + if not llm_call.done() or llm_call.cancelled() or llm_call.exception() is not None: + return + response: Final = llm_call.result() + prepare_response_cost(logging_obj, response, datetime.now(timezone.utc)) + with anyio.CancelScope(shield=True): + await finalize_proxy_response_cost(logging_obj, response) + + async def _cancel_pending_gather_tasks(tasks: Sequence["asyncio.Task[object]"]) -> None: pending_tasks: Final = [task for task in tasks if not task.done()] for task in pending_tasks: @@ -3128,33 +3166,88 @@ class ProxyBaseLLMRequestProcessing: ) return await open_sse_before_first_byte( - self._process_llm_request( + self._within_total_timeout( request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, route_type=route_type, - proxy_logging_obj=proxy_logging_obj, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - llm_router=llm_router, - model=model, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, is_streaming_request=is_streaming_request, - contents=contents, - skip_pre_call_logic=skip_pre_call_logic, - first_chunk_framing=sse_framing.first_chunk, + processing=self._process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type=route_type, + proxy_logging_obj=proxy_logging_obj, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + llm_router=llm_router, + model=model, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + is_streaming_request=is_streaming_request, + contents=contents, + skip_pre_call_logic=skip_pre_call_logic, + first_chunk_framing=sse_framing.first_chunk, + ), ), ping_interval_seconds=ttft_keepalive_interval(self.data, llm_router), on_late_failure=_audit_late_failure, late_framing=ANTHROPIC_LATE_SSE_FRAMING if answers_as_anthropic else sse_framing.late, ) + def total_timeout_for( + self, request: Request, route_type: ProxyRouteType, is_streaming_request: bool | None + ) -> TotalTimeout | None: + """The request's ``total_timeout`` when it bounds this request: a body one on a request it cannot + bound is refused, a header one is ignored.""" + body: Final = _REQUEST_FIELDS.validate_python(self.data) # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict + header: Final = request.headers.get(TOTAL_TIMEOUT_HEADER) + self.data.pop(TOTAL_TIMEOUT_FIELD, None) # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict + bounded: Final = ( + route_type in _TOTAL_TIMEOUT_ROUTE_TYPES + and is_streaming_request is not True + and bounded_by_total_timeout(body) + ) + if header is not None and not bounded: + verbose_proxy_logger.debug("%s ignored for this %s request", TOTAL_TIMEOUT_HEADER, route_type) + return None + parsed: Final = total_timeout_from_request(header, body) + if parsed is None: + return None + if isinstance(parsed, InvalidTotalTimeout): + raise _total_timeout_refused(parsed.message) + if not bounded: + raise _total_timeout_refused( + f"{TOTAL_TIMEOUT_FIELD} applies only to non-streaming, non-background " + "/v1/chat/completions, /v1/messages and /v1/responses requests" + ) + return parsed + + async def _within_total_timeout( + self, + request: Request, + route_type: ProxyRouteType, + is_streaming_request: bool | None, + processing: Coroutine[object, object, _ProcessedT], + ) -> _ProcessedT: + try: + total_timeout: Final = self.total_timeout_for(request, route_type, is_streaming_request) + except ProxyException: + processing.close() + raise + if total_timeout is None: + return await processing + try: + async with total_deadline(total_timeout.seconds, from_body=total_timeout.from_body): + return await processing + except TotalTimeoutExpired as expired: + raise ProxyException(message=str(expired), type="timeout", param=TOTAL_TIMEOUT_FIELD, code=408) from expired + except TotalTimeoutNotApplicable as refused: + raise _total_timeout_refused(str(refused)) from refused + async def _process_llm_request( self, request: Request, @@ -3238,25 +3331,32 @@ class ProxyBaseLLMRequestProcessing: if contents: self.data["contents"] = contents - ### ROUTE THE REQUEST ### - # Do not change this - it should be a constant time fetch - ALWAYS - llm_call: Final = await route_request( - data=self.data, - route_type=route_type, - llm_router=llm_router, - user_model=user_model, - user_api_key_dict=user_api_key_dict, - ) - llm_call_task: Final = asyncio.create_task(with_proxy_compaction_executor(llm_call, request)) - tasks.append(llm_call_task) - - llm_responses: Final = asyncio.gather(*tasks) # run the moderation check in parallel to the actual llm api call - try: - if general_settings.get("cancel_on_disconnect", False): - responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) - else: - responses = await llm_responses + ### ROUTE THE REQUEST ### + # Do not change this - it should be a constant time fetch - ALWAYS + llm_call: Final = await route_request( + data=self.data, + route_type=route_type, + llm_router=llm_router, + user_model=user_model, + user_api_key_dict=user_api_key_dict, + ) + llm_call_task: Final = asyncio.create_task(with_proxy_compaction_executor(llm_call, request)) + tasks.append(llm_call_task) + + llm_responses: Final = asyncio.gather( + *tasks + ) # run the moderation check in parallel to the actual llm api call + + try: + if general_settings.get("cancel_on_disconnect", False): + responses = await _await_llm_call_cancelling_on_disconnect(request, llm_responses) + else: + responses = await llm_responses + except asyncio.CancelledError: + if _total_timeout_expired(): + await _cost_a_call_answered_past_its_total_timeout(llm_call_task, logging_obj) + raise finally: await _cancel_pending_gather_tasks(tasks) @@ -3526,7 +3626,7 @@ class ProxyBaseLLMRequestProcessing: _guardrail_entries_for_headers(get_or_create_metadata_bucket(self.data)[1], logging_obj) ), ) - if isinstance(exc, asyncio.CancelledError): + if isinstance(exc, asyncio.CancelledError) and not _total_timeout_expired(): await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=HTTPException(status_code=499, detail=_CLIENT_DISCONNECT_DETAIL), diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6b36a6ac8bc..9fb026a64d9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -65,6 +65,7 @@ from litellm.litellm_core_utils.token_counter import ( offload_token_count, split_embedding_input_tokens, ) +from litellm.litellm_core_utils.total_deadline import total_timeout_held from litellm.llms.azure.ptu_capacity import normalized_tokens from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( @@ -3801,6 +3802,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raise cancellation @with_service_target(RATE_LIMITS_TARGET) + @total_timeout_held() async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 50d1c7c5b59..574a677061d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -398,6 +398,9 @@ async def responses_api( # instead of a polling ID that immediately fails in the background task. processor = ProxyBaseLLMRequestProcessing(data=data) try: + # A background request runs after its polling id is returned, so a total_timeout it + # cannot apply is refused (or a header one dropped) before the id is handed out. + _ = processor.total_timeout_for(request, "aresponses", is_streaming_request=None) data, _logging_obj = await processor.common_processing_pre_call_logic( request=request, general_settings=general_settings, diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 08942851af9..ffc0e31408d 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -17,6 +17,7 @@ import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.routing_common_utils import is_proxy_admin_request +from litellm.litellm_core_utils.total_deadline import enforce_total_timeout_scope from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router import router_arguments from litellm.router_api import router_init_args @@ -761,6 +762,7 @@ async def route_request( """ Common helper to route the request """ + enforce_total_timeout_scope(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body router_only_settings_from_override: Final = _merge_router_settings_override(data) # pyright: ignore[reportUnknownArgumentType] # route_request takes an untyped dict body try: return await _route_request_single_attempt( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 503dfb69793..8b9189f143f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -188,6 +188,7 @@ from litellm.litellm_core_utils.served_output_texts import ( served_stream_output_texts, ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model +from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired from litellm.llms import load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -1737,8 +1738,17 @@ def _client_disconnected(exception: Exception) -> bool: return isinstance(exception, HTTPException) and exception.status_code == LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED +def _ran_out_of_total_timeout(exception: Exception) -> bool: + """The request's ``total_timeout`` ran out while the provider was working on it, like a disconnect.""" + return any(isinstance(cause, TotalTimeoutExpired) for cause in _exception_chain(exception)) + + def _provider_billed_the_failure(exception: Exception) -> bool: - return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) + return ( + _timed_out_reading_from_provider(exception) + or _client_disconnected(exception) + or _ran_out_of_total_timeout(exception) + ) def _as_str_mapping(value: object) -> Mapping[str, object]: diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index b7c9a96c9ff..55155a96fa3 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -69,6 +69,11 @@ from litellm.litellm_core_utils.routing_common_utils import ( ) from litellm.litellm_core_utils.routing_groups import validate_routing_strategy from litellm.litellm_core_utils.secret_redaction import redact_string +from litellm.litellm_core_utils.total_deadline import ( + TotalTimeoutNotApplicable, + current_total_deadline, + enforce_total_timeout_scope, +) from litellm.redis.stores.connection import connect from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, RouterCooldownStore from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage @@ -2915,6 +2920,9 @@ class RustCore( case "attempt": outcome = await driver.attempt(step) case "sleep": + if (outlasted := _error_outlasting_total_timeout(driver, _sleep_seconds(step))) is not None: + driver.close_route_phase() + raise outlasted await self._host.sleep(_sleep_seconds(step)) outcome = None case "allow_fallback": @@ -3873,6 +3881,15 @@ def _origin_level(origin: Mapping[str, object]) -> int: return level if isinstance(level, int) else 0 +def _error_outlasting_total_timeout(driver: "_Driver", seconds: float) -> BaseException | None: + """The error a retry would have answered when its backoff outlasts the request's total_timeout: the + call ends with it now, rather than with a timeout after waiting out the rest of the budget.""" + deadline: Final = current_total_deadline() + if deadline is None or deadline.allows_wait(seconds): + return None + return driver.last_error + + def _sleep_seconds(step: Mapping[str, object]) -> float: """A retry's sleep: ``_time_to_sleep_before_retry`` returns a plain ``0`` when another deployment is healthy, else a backoff in float seconds.""" @@ -5198,6 +5215,7 @@ class _Driver: return {**self._failed(step, self._default_copy(step), None, refused_call), "recorded": True} try: deployment, call = self._call(step) # rebind-ok: rebound as the Python router rebinds it + enforce_total_timeout_scope(call.kwargs) if call.refused is not None: await log_refused_call(call, call.refused) raise call.refused @@ -5257,6 +5275,8 @@ class _Driver: ATTEMPT_SETTLEMENT_GRACE_SECONDS, ) raise + except TotalTimeoutNotApplicable: + raise except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop if call is not None: await released(self._release(call, attempt, refund_io_tokens=True)) diff --git a/tests/unit/litellm_core_utils/test_total_deadline.py b/tests/unit/litellm_core_utils/test_total_deadline.py new file mode 100644 index 00000000000..f3bc99354ea --- /dev/null +++ b/tests/unit/litellm_core_utils/test_total_deadline.py @@ -0,0 +1,36 @@ +from collections.abc import Mapping + +import pytest + +from litellm.litellm_core_utils.total_deadline import InvalidTotalTimeout, TotalTimeout, total_timeout_from_request + + +@pytest.mark.parametrize( + ("header", "body", "expected"), + [ + (None, {}, None), + (None, {"total_timeout": 10}, TotalTimeout(10.0, from_body=True)), + (None, {"total_timeout": "2.5"}, TotalTimeout(2.5, from_body=True)), + ("4", {}, TotalTimeout(4.0, from_body=False)), + ("4", {"total_timeout": 10}, TotalTimeout(4.0, from_body=False)), + ("4", {"total_timeout": "bad"}, TotalTimeout(4.0, from_body=False)), + ], + ids=["absent", "body", "body_numeric_string", "header", "header_wins", "header_wins_over_a_bad_body"], +) +def test_total_timeout_is_read_from_the_header_first_then_the_body( + header: str | None, body: Mapping[str, object], expected: TotalTimeout | None +) -> None: + assert total_timeout_from_request(header, body) == expected + + +@pytest.mark.parametrize( + "value", + [None, True, False, 0, -1, "0", "-2", "nan", "inf", "-inf", float("nan"), 10**400, "ten", "", [10], {"s": 10}], +) +def test_a_total_timeout_that_is_not_a_positive_finite_number_of_seconds_is_invalid(value: object) -> None: + assert isinstance(total_timeout_from_request(None, {"total_timeout": value}), InvalidTotalTimeout) + + +@pytest.mark.parametrize("value", ["0", "-1", "nan", "inf", "ten", ""]) +def test_an_invalid_total_timeout_header_is_invalid(value: str) -> None: + assert isinstance(total_timeout_from_request(value, {}), InvalidTotalTimeout) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 64e7cde4034..741b6cb74b5 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias from litellm.litellm_core_utils.ptu_shares import PTUTeamCeiling +from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired, total_deadline from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError @@ -66,6 +67,7 @@ from litellm.redis.stores.rate_limits import ( RateLimitStore, ) from litellm.responses.streaming_iterator import _record_usage_for_failure +from tests._support.clock_loops import AutojumpClockLoop from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits, scope_key from litellm.responses.utils import ResponseAPILoggingUtils from litellm.types.extraction import ExtractionResponse @@ -7951,6 +7953,34 @@ async def _admit(handler, auth, data=None): ) +def test_a_total_timeout_never_interrupts_a_token_reservation_so_the_failure_path_can_refund_it() -> None: + async def reserve_outlasting_the_budget(keys: list[str], args: list[object]) -> list[object]: + await asyncio.sleep(120) + return [0, *[value for index in range(0, len(args), 4) for value in (0, args[index + 1])]] + + async def admit_under_a_total_timeout() -> tuple[bool, int]: + redis: Final = FakeRateLimits(check_and_increment_by_n_script=reserve_outlasting_the_budget) + handler: Final = _handler_with_redis(redis) + auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-total-timeout-reservation"), tpm_limit=1000) + stash: Final = get_or_create_request_stash() + try: + async with total_deadline(60, from_body=True): + await _admit(handler, auth) + await asyncio.sleep(3600) + except TotalTimeoutExpired: + return True, stash.reserved_tokens + return False, stash.reserved_tokens + + loop: Final = AutojumpClockLoop() + try: + expired, reserved = loop.run_until_complete(admit_under_a_total_timeout()) + finally: + loop.close() + + assert expired + assert reserved > 0 + + async def _read_only_check(handler, auth): descriptors = handler.create_rate_limit_descriptors( user_api_key_dict=auth, diff --git a/tests/unit/proxy/search_endpoints/test_endpoints.py b/tests/unit/proxy/search_endpoints/test_endpoints.py index 7bb4df37314..92e9076678f 100644 --- a/tests/unit/proxy/search_endpoints/test_endpoints.py +++ b/tests/unit/proxy/search_endpoints/test_endpoints.py @@ -142,6 +142,7 @@ def test_direct_search_body_tool_name_is_denied_under_search_tool_deny_by_defaul def _json_request(body: dict[str, object]) -> MagicMock: request = MagicMock() request.body = AsyncMock(return_value=orjson.dumps(body)) + request.headers = {} return request diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a519da41f61..901dbb6ea97 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4,9 +4,9 @@ import copy import datetime import gzip import json -from collections.abc import AsyncIterable, Awaitable, Mapping +from collections.abc import AsyncIterable, Awaitable, Coroutine, Mapping from types import MappingProxyType, SimpleNamespace -from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence +from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeVar from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch @@ -82,6 +82,7 @@ from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException +from litellm.litellm_core_utils.total_deadline import current_total_deadline from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.hooks.model_max_budget_limiter import ( @@ -11213,7 +11214,7 @@ async def test_base_process_llm_request_pings_while_the_upstream_call_is_still_r with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_inner): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="acompletion", @@ -11391,7 +11392,7 @@ async def test_base_process_llm_request_audits_a_failure_that_lands_after_its_ke with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_failure): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=user_api_key_dict, route_type="acompletion", @@ -11437,7 +11438,7 @@ async def test_base_process_llm_request_honours_a_deployment_hard_disable(deploy with patch.object(litellm, "sse_keepalive_ping_interval_seconds", 0.05): with patch.object(ProxyBaseLLMRequestProcessing, "_process_llm_request", slow_inner): response = await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=Response(), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth), route_type="acompletion", @@ -13321,3 +13322,343 @@ async def test_a_chained_proxy_redacts_secrets_and_internal_details_from_relayed assert {key: answered[key] for key in _DOWNSTREAM_GUARDRAIL_FIELDS} == dict(_DOWNSTREAM_GUARDRAIL_FIELDS) assert answered["scanner"]["attempts"] == 2 assert [value for value in _DOWNSTREAM_LEAKABLE_VALUES if value in json.dumps(answered)] == [] + + +_ScenarioT = TypeVar("_ScenarioT") + + +def _run_on_autojump_clock(scenario: Coroutine[object, object, _ScenarioT]) -> _ScenarioT: + loop: Final = AutojumpClockLoop() + try: + return loop.run_until_complete(scenario) + finally: + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.close() + + +class TestTotalTimeout: + """``total_timeout`` / ``x-litellm-total-timeout`` bound a whole non-streaming request.""" + + def _request(self, headers: Mapping[str, str]) -> Request: + async def receive(): + await asyncio.Event().wait() + + return Request( + scope={"type": "http", "headers": [(k.encode(), v.encode()) for k, v in headers.items()]}, + receive=receive, + ) + + async def _process( + self, + monkeypatch: pytest.MonkeyPatch, + llm_call: Callable[[Mapping[str, object]], object], + *, + routing: Callable[[], Awaitable[None]] | None = None, + body: Mapping[str, object] = MappingProxyType({}), + headers: Mapping[str, str] = MappingProxyType({}), + moderation: Callable[..., object] | None = None, + post_call: Callable[..., object] | None = None, + route_type: str = "acompletion", + proxy_logging_obj: MagicMock | None = None, + preparation: Callable[..., Awaitable[tuple[dict[str, object], MagicMock]]] | None = None, + ) -> tuple[object, MagicMock]: + from litellm.proxy._types import UserAPIKeyAuth + + logging_obj: Final = MagicMock() + logging_obj.litellm_call_id = "total-timeout-call" + logging_obj.defer_async_logging = False + logging_obj._on_deferred_stream_complete = None + logging_obj.cost_breakdown = None + logging_obj.finalized_cost = None + logging_obj.prepared_cost = None + + def prepared(response: object, now: object) -> object: + logging_obj.prepared_cost = response + return response + + logging_obj.prepare_response_cost.side_effect = prepared + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": "fake-model", "litellm_logging_obj": logging_obj, **body} + ) + hooks: Final = proxy_logging_obj if proxy_logging_obj is not None else MagicMock(spec=ProxyLogging) + hooks.during_call_hook = AsyncMock(side_effect=moderation, return_value=None) + hooks.update_request_status = AsyncMock(return_value=None) + hooks.post_call_success_hook = AsyncMock( + side_effect=post_call or (lambda data, user_api_key_dict, response: response) + ) + hooks.post_call_response_headers_hook = AsyncMock(return_value=None) + hooks.post_call_failure_hook = AsyncMock(return_value=None) + + async def fake_route_request(**kwargs): + if routing is not None: + await routing() + return llm_call(kwargs["data"]) + + monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request) + if preparation is not None: + monkeypatch.setattr(processor, "_pre_call_with_fallbacks", preparation) + result: Final = await processor.base_process_llm_request( + request=self._request(headers), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + route_type=route_type, + proxy_logging_obj=hooks, + general_settings={}, + proxy_config=MagicMock(spec=ProxyConfig), + skip_pre_call_logic=preparation is None, + ) + return result, logging_obj + + @pytest.mark.parametrize( + ("body", "headers"), + [({"total_timeout": 60}, {}), ({}, {"x-litellm-total-timeout": "60"})], + ids=["body", "header"], + ) + def test_a_request_past_its_total_timeout_is_answered_with_a_408_and_its_call_is_cancelled( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], headers: Mapping[str, str] + ) -> None: + async def provider_hangs_past_the_budget() -> tuple[ProxyException, bool]: + provider_called: Final = asyncio.Event() + upstream_cancelled: Final = asyncio.Event() + + async def hanging(_: Mapping[str, object]) -> object: + provider_called.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + upstream_cancelled.set() + raise + raise AssertionError + + processing: Final = asyncio.create_task(self._process(monkeypatch, hanging, body=body, headers=headers)) + await provider_called.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, upstream_cancelled.is_set() + + expired, upstream_cancelled = _run_on_autojump_clock(provider_hangs_past_the_budget()) + + assert expired.code == "408" + assert expired.param == "total_timeout" + assert upstream_cancelled + + def test_a_moderation_check_outlasting_the_budget_after_the_provider_answered_prepares_its_cost( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + answer: Final = litellm.ModelResponse() + + async def moderation_outlasts_the_budget() -> tuple[ProxyException, MagicMock]: + answered: Final = asyncio.Event() + moderating: Final = asyncio.Event() + logging_objs: Final[list[MagicMock]] = [] + + async def provider(data: Mapping[str, object]) -> object: + logging_objs.append(data["litellm_logging_obj"]) # pyright: ignore[reportArgumentType] # the harness's logging double + answered.set() + return answer + + async def slow_moderation(**_: object) -> None: + moderating.set() + await asyncio.Event().wait() + + processing: Final = asyncio.create_task( + self._process(monkeypatch, provider, body={"total_timeout": 60}, moderation=slow_moderation) + ) + await answered.wait() + await moderating.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, logging_objs[0] + + expired, logging_obj = _run_on_autojump_clock(moderation_outlasts_the_budget()) + + costed: Final = [call.args[0] for call in logging_obj.prepare_response_cost.call_args_list] + assert expired.code == "408" + assert costed + assert all(response is answer for response in costed) + + async def test_a_provider_timeout_inside_the_budget_is_not_turned_into_a_408( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def provider_timed_out(_: Mapping[str, object]) -> object: + raise TimeoutError("provider read timeout") + + with pytest.raises(TimeoutError, match="provider read timeout"): + await self._process(monkeypatch, provider_timed_out, body={"total_timeout": 60}) + + async def test_cancelling_a_request_inside_its_budget_still_cancels_it( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + provider_called: Final = asyncio.Event() + + async def hanging(_: Mapping[str, object]) -> object: + provider_called.set() + await asyncio.Event().wait() + raise AssertionError + + processing: Final = asyncio.create_task(self._process(monkeypatch, hanging, body={"total_timeout": 60})) + await provider_called.wait() + processing.cancel() + + with pytest.raises(asyncio.CancelledError): + await processing + + async def test_a_bounded_request_reaches_the_router_bounded_and_without_its_total_timeout_field( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + seen: Final[list[tuple[bool, bool]]] = [] + + async def answered(data: Mapping[str, object]) -> object: + seen.append(("total_timeout" in data, current_total_deadline() is not None)) + return litellm.ModelResponse() + + await self._process(monkeypatch, answered, body={"total_timeout": 60}) + + assert seen == [(False, True)] + + @pytest.mark.parametrize( + ("body", "route_type"), + [ + ({"total_timeout": 60, "stream": True}, "acompletion"), + ({"total_timeout": 60, "background": True}, "aresponses"), + ({"total_timeout": 60}, "aembedding"), + ({"total_timeout": 0}, "acompletion"), + ({"total_timeout": None}, "acompletion"), + ], + ids=["streaming", "background_responses", "other_route", "zero", "null"], + ) + async def test_a_body_total_timeout_it_cannot_apply_is_refused_before_the_call( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], route_type: str + ) -> None: + called: Final[list[object]] = [] + + async def answered(data: Mapping[str, object]) -> object: + called.append(data) + return litellm.ModelResponse() + + with pytest.raises(ProxyException) as raised: + await self._process(monkeypatch, answered, body=body, route_type=route_type) + + assert raised.value.code == "400" + assert raised.value.param == "total_timeout" + assert called == [] + + @pytest.mark.parametrize("header", ["60", "10s", "", "0"], ids=["valid", "with_a_unit", "empty", "zero"]) + @pytest.mark.parametrize( + ("body", "route_type"), + [({"stream": True}, "acompletion"), ({"background": True}, "aresponses"), ({}, "aembedding")], + ids=["streaming", "background_responses", "other_route"], + ) + async def test_a_header_total_timeout_leaves_a_request_it_cannot_apply_unbounded_whatever_its_value( + self, monkeypatch: pytest.MonkeyPatch, body: Mapping[str, object], route_type: str, header: str + ) -> None: + bounded: Final[list[bool]] = [] + + async def answered(_: Mapping[str, object]) -> object: + bounded.append(current_total_deadline() is not None) + return litellm.ModelResponse() + + await self._process( + monkeypatch, answered, body=body, headers={"x-litellm-total-timeout": header}, route_type=route_type + ) + + assert bounded == [False] + + def test_a_post_call_hook_outlasting_the_budget_is_a_408_not_a_client_disconnect( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + hooks: Final = MagicMock(spec=ProxyLogging) + + async def guardrail_outlasts_the_budget() -> ProxyException: + guarding: Final = asyncio.Event() + + async def answered(_: Mapping[str, object]) -> object: + return litellm.ModelResponse() + + async def slow_guardrail(**_: object) -> object: + guarding.set() + await asyncio.Event().wait() + raise AssertionError + + processing: Final = asyncio.create_task( + self._process( + monkeypatch, answered, body={"total_timeout": 60}, post_call=slow_guardrail, proxy_logging_obj=hooks + ) + ) + await guarding.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value + + expired: Final = _run_on_autojump_clock(guardrail_outlasts_the_budget()) + + assert expired.code == "408" + hooks.post_call_failure_hook.assert_not_awaited() + + def test_a_budget_that_runs_out_while_routing_also_stops_the_moderation_check( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def routing_outlasts_the_budget() -> tuple[ProxyException, bool]: + moderating: Final = asyncio.Event() + routing: Final = asyncio.Event() + moderation_stopped: Final = asyncio.Event() + + async def moderation(**_: object) -> None: + moderating.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + moderation_stopped.set() + raise + + async def stuck_routing() -> None: + routing.set() + await asyncio.Event().wait() + + async def never_called(_: Mapping[str, object]) -> object: + raise AssertionError + + processing: Final = asyncio.create_task( + self._process( + monkeypatch, never_called, routing=stuck_routing, body={"total_timeout": 60}, moderation=moderation + ) + ) + await moderating.wait() + await routing.wait() + with pytest.raises(ProxyException) as raised: + await processing + return raised.value, moderation_stopped.is_set() + + expired, moderation_stopped = _run_on_autojump_clock(routing_outlasts_the_budget()) + + assert expired.code == "408" + assert moderation_stopped + + def test_a_budget_that_runs_out_while_preparing_cuts_preparation_short_and_never_calls_the_provider( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def preparation_outlasts_the_budget() -> tuple[ProxyException, bool, bool, float]: + loop: Final = asyncio.get_running_loop() + started: Final = loop.time() + prepared: Final = asyncio.Event() + called: Final = asyncio.Event() + + async def slow_guardrail(**_: object) -> tuple[dict[str, object], MagicMock]: + await asyncio.sleep(120) + prepared.set() + return {"model": "fake-model"}, MagicMock() + + async def provider(_: Mapping[str, object]) -> object: + called.set() + return litellm.ModelResponse() + + with pytest.raises(ProxyException) as raised: + await self._process(monkeypatch, provider, body={"total_timeout": 60}, preparation=slow_guardrail) + return raised.value, prepared.is_set(), called.is_set(), loop.time() - started + + expired, prepared, called, elapsed = _run_on_autojump_clock(preparation_outlasts_the_budget()) + + assert expired.code == "408" + assert not prepared + assert not called + assert elapsed < 120 diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 97c21fcc777..af463fa19d4 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,8 +17,6 @@ from pydantic import ValidationError as PydanticValidationError from starlette.datastructures import Headers import litellm -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs from litellm.constants import ( CALLER_SESSION_ID_METADATA_KEY, COMPLEXITY_ROUTER_CHILD_SESSION_ID_METADATA_KEY, @@ -40,6 +38,7 @@ from litellm.litellm_core_utils.redact_messages import _get_turn_off_message_log from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.proxy._types import AddTeamCallback, ProxyException, TeamCallbackMetadata, UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.litellm_pre_call_utils import ( KeyAndTeamLoggingSettings, LiteLLMProxyRequestSetup, @@ -64,6 +63,7 @@ from litellm.proxy.litellm_pre_call_utils import ( from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES from litellm.types.utils import CredentialItem +from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs def test_check_if_token_is_service_account(): @@ -112,7 +112,8 @@ class TestGetMetadataVariableName: assert get_metadata_variable_name(request) == "litellm_metadata" @pytest.mark.parametrize( - "path", ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"] + "path", + ["/v1beta/interactions", "/interactions", "/v1beta/interactions/v1_abc123", "/interactions/v1_abc123/cancel"], ) def test_returns_litellm_metadata_for_interactions_routes(self, path): request = self._make_request(path) @@ -808,7 +809,9 @@ async def test_interactions_spend_logs_metadata_survives_the_router_metadata_buc # The router then adds its deployment fields to the same bucket, and the spend log # resolves that bucket. - data["litellm_metadata"].update({"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"}) + data["litellm_metadata"].update( + {"model_group": "gemini-3.1-flash-lite", "deployment": "gemini/gemini-3.1-flash-lite"} + ) resolved = get_litellm_metadata_from_kwargs( {"litellm_params": {"metadata": data.get("metadata") or {}, "litellm_metadata": data["litellm_metadata"]}} ) @@ -3759,53 +3762,87 @@ async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypa "surface,headers,body,child_id,trace_id", [ ( - "chat/completions", {"X-Session-Id": "child-01"}, {"metadata": {"session_id": "body-001"}}, - "child-01", "parent-01", + "chat/completions", + {"X-Session-Id": "child-01"}, + {"metadata": {"session_id": "body-001"}}, + "child-01", + "parent-01", ), ("responses", {"x-vendor-session-id": "child-01"}, {}, "child-01", "parent-01"), ("messages", {"session-id": "child-01", "user-agent": "codex_cli_rs/1.0"}, {}, "child-01", "parent-01"), ("chat/completions", {"x-litellm-session-id": "short", "x-session-id": "child-01"}, {}, "short", "short"), ("responses", {"x-litellm-trace-id": "trace-01", "x-session-id": "child-01"}, {}, "child-01", "trace-01"), ( - "messages", {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, - {}, "child-01", "trace-01", + "messages", + {"x-litellm-trace-id": "trace-01", "x-litellm-session-id": "child-01"}, + {}, + "child-01", + "trace-01", ), ("responses", {}, {"metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("chat/completions", {}, {"litellm_metadata": {"session_id": "body-001"}}, "body-001", "parent-01"), ("messages", {"baggage": "session.id=child-01"}, {}, "child-01", "parent-01"), ( - "responses", {"baggage": "session.id=child-01"}, {"metadata": {"session_id": "body-001"}}, - "body-001", "parent-01", + "responses", + {"baggage": "session.id=child-01"}, + {"metadata": {"session_id": "body-001"}}, + "body-001", + "parent-01", ), ( - "chat/completions", {}, {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, - "root-001", "parent-01", + "chat/completions", + {}, + {"litellm_session_id": "root-001", "metadata": {"session_id": "body-001"}}, + "root-001", + "parent-01", ), ( - "responses", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, - "active-1", "parent-01", + "responses", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": "active-1"}}, + "active-1", + "parent-01", ), ( - "chat/completions", {}, + "chat/completions", + {}, {"metadata": {"session_id": "active-1"}, "litellm_metadata": {"session_id": "body-001"}}, - "active-1", "parent-01", + "active-1", + "parent-01", ), ( - "chat/completions", {}, {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, - None, "parent-01", + "chat/completions", + {}, + {"metadata": {"session_id": ""}, "litellm_metadata": {"session_id": "body-001"}}, + None, + "parent-01", ), ( - "messages", {}, {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, - None, "parent-01", + "messages", + {}, + {"metadata": {"session_id": "body-001"}, "litellm_metadata": {"session_id": ""}}, + None, + "parent-01", ), ( - "messages", {}, {"litellm_session_id": "root-001", "litellm_metadata": {"session_id": "active-1"}, - "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}}, - "anthro-1", "parent-01", + "messages", + {}, + { + "litellm_session_id": "root-001", + "litellm_metadata": {"session_id": "active-1"}, + "metadata": {"session_id": "body-001", "user_id": "user_abc_account__session_anthro-1"}, + }, + "anthro-1", + "parent-01", + ), + ( + "messages", + {"x-session-id": "child-01"}, + {"metadata": {"user_id": {"session_id": "body-001"}}}, + "child-01", + "parent-01", ), - ("messages", {"x-session-id": "child-01"}, {"metadata": {"user_id": {"session_id": "body-001"}}}, - "child-01", "parent-01"), ("chat/completions", {}, {}, None, "parent-01"), ("responses", {"x-session-id": "bad"}, {}, None, "parent-01"), ("messages", {}, {"metadata": {"session_id": 123}}, None, "parent-01"), @@ -4347,9 +4384,7 @@ def _bedrock_passthrough_data( ), ], ) -def test_bedrock_request_metadata_thread_id_becomes_session_id( - headers: dict[str, str], client_body: dict[str, object] -): +def test_bedrock_request_metadata_thread_id_becomes_session_id(headers: dict[str, str], client_body: dict[str, object]): sent_body: Final = copy.deepcopy(client_body) data: Final = _bedrock_passthrough_data(client_body) @@ -4429,7 +4464,9 @@ def test_bedrock_request_metadata_without_usable_session_id_sets_no_session( @pytest.mark.parametrize( ("body_tags", "header_tags", "expected"), [ - pytest.param({"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key"), + pytest.param( + {"team": "search"}, {"threadId": "thread-abc-123"}, "thread-abc-123", id="body_without_session_key" + ), pytest.param({"sessionId": "body-session"}, {"sessionId": "header-session"}, "body-session", id="body_wins"), ], ) @@ -11226,9 +11263,7 @@ async def test_header_mapped_user_id_cannot_load_another_users_connection(monkey from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ @@ -11281,9 +11316,7 @@ async def test_per_user_token_resolution_uses_key_identity_without_header_mappin from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request _per_user_env(monkeypatch) - load = AsyncMock( - return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}} - ) + load = AsyncMock(return_value={"copilot-cred": {"access_token": "gho_user_a_token", "github_host": "github.com"}}) monkeypatch.setattr(upc, "aget_user_provider_credentials", load) router = Router( model_list=[ diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index f74ff9b1fca..867259a9b3b 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -561,7 +561,11 @@ class TestPostCallFailureHookDoesNotInventAProviderHandoff: litellm_call_id="pre-handoff-auth-failure", function_id="pre-handoff", ) - request_data: Final[dict[str, object]] = {"litellm_logging_obj": logging_obj, "model": "gpt-4.1", "metadata": {}} + request_data: Final[dict[str, object]] = { + "litellm_logging_obj": logging_obj, + "model": "gpt-4.1", + "metadata": {}, + } proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging_obj.alert_types = [] with ( @@ -1116,6 +1120,44 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: assert expected_input_cost > 0 assert request_data["response_cost"] == pytest.approx(expected_input_cost) + @pytest.mark.asyncio + async def test_a_total_timeout_that_ran_out_after_dispatch_is_priced_at_input_rate(self): + from datetime import datetime + + from litellm.litellm_core_utils.total_deadline import TotalTimeoutExpired + from litellm.proxy._types import ProxyException + from litellm.types.utils import Usage + + def ran_out() -> ProxyException: + try: + raise TotalTimeoutExpired("the request did not finish within its total_timeout of 2.0s") + except TotalTimeoutExpired as expired: + try: + raise ProxyException( + message=str(expired), type="timeout", param="total_timeout", code=408 + ) from expired + except ProxyException as answered: + return answered + + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "count these input tokens please"}], + "call_type": "acompletion", + } + ), + "metadata": {}, + } + await self._run(request_data, original_exception=ran_out()) + + estimated = request_data["combined_usage_object"] + assert isinstance(estimated, Usage) + expected_input_cost, _ = litellm.cost_per_token(model="gpt-3.5-turbo", prompt_tokens=estimated.prompt_tokens) + assert expected_input_cost > 0 + assert request_data["response_cost"] == pytest.approx(expected_input_cost) + @pytest.mark.asyncio async def test_upstream_499_error_is_a_plain_failure_not_a_priced_disconnect(self): """A provider answering HTTP 499 surfaces as litellm.APIError with status_code 499. Only the @@ -2693,7 +2735,12 @@ def test_an_earlier_attempt_on_another_model_is_priced_at_that_models_rate_not_t monkeypatch.setitem( litellm.model_cost, pricing_model, - {"litellm_provider": provider, "mode": "chat", "input_cost_per_token": input_rate, "output_cost_per_token": 1.0}, + { + "litellm_provider": provider, + "mode": "chat", + "input_cost_per_token": input_rate, + "output_cost_per_token": 1.0, + }, ) usage = litellm.Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) @@ -2926,9 +2973,7 @@ async def test_dropped_stream_failure_attribution_reaches_terminal_spend_row() - request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3101,9 +3146,7 @@ async def test_latest_dropped_attempt_publishes_its_payload_after_earlier_dispat request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3158,9 +3201,7 @@ async def test_dropped_attempt_attribution_is_published_while_failure_callback_i request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3237,9 +3278,7 @@ async def test_failed_attempt_publishes_baseline_reason_after_invalidation( request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3254,9 +3293,7 @@ async def test_failed_attempt_publishes_baseline_reason_after_invalidation( assert row["prompt_tokens"] == 4 assert row["completion_tokens"] == 1 lifted_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(lifted["standard_logging_object"]) - callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( - recorder.standard_logging_object - ) + callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(recorder.standard_logging_object) def baseline_reason(payload: StandardLoggingPayload) -> str: encoded: Final = payload["autorouter_baseline_observation"] @@ -3335,9 +3372,7 @@ async def test_queued_stale_attempt_dispatch_does_not_publish_over_current_attem request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3402,9 +3437,7 @@ async def test_dropped_stream_failure_does_not_overwrite_newer_terminal_payload( request_data: Final = {"litellm_logging_obj": logging_obj} lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) terminal_kwargs: Final = { - key: value - for key, value in {**request_data, **lifted}.items() - if key != "litellm_logging_obj" + key: value for key, value in {**request_data, **lifted}.items() if key != "litellm_logging_obj" } row: Final = get_logging_payload( kwargs=terminal_kwargs, @@ -3658,7 +3691,6 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) - COMPACTION_INSTRUCTIONS = {"role": "system", "content": "answer briefly"} COMPACTED_HISTORY = [ {"role": "user", "content": "a long turn " * 200}, @@ -3735,8 +3767,6 @@ def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) - - @pytest.mark.parametrize( "current_attempt_api_call_start_time", [None, 1700000001.0], @@ -4076,9 +4106,7 @@ def test_create_model_info_response_resolves_mode_through_deployment_model(): ] ) - response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) finally: litellm.model_cost.clear() litellm.model_cost.update(saved_model_cost) @@ -4108,12 +4136,8 @@ def test_create_model_info_response_resolves_model_group_alias_to_target(model_g model_group_alias=model_group_alias, ) - alias_response = create_model_info_response( - model_id="team-embeddings", provider="openai", llm_router=router - ) - target_response = create_model_info_response( - model_id="my-embeddings", provider="openai", llm_router=router - ) + alias_response = create_model_info_response(model_id="team-embeddings", provider="openai", llm_router=router) + target_response = create_model_info_response(model_id="my-embeddings", provider="openai", llm_router=router) assert alias_response["id"] == "team-embeddings" for field in ("mode", "max_input_tokens", "max_output_tokens"): @@ -4196,7 +4220,9 @@ async def test_post_call_failure_hook_redacts_traceback_before_callbacks(monkeyp with patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()): await proxy_logging_obj.post_call_failure_hook( request_data={"metadata": {}}, - original_exception=HTTPException(status_code=400, detail="Upstream passthrough request failed with status 400"), + original_exception=HTTPException( + status_code=400, detail="Upstream passthrough request failed with status 400" + ), user_api_key_dict=UserAPIKeyAuth(), traceback_str=upstream_traceback, ) @@ -4260,9 +4286,13 @@ def test_mcp_conversion_preserves_request_policy_and_isolates_guardrail_data(buc parent = { "model": "parent-model", bucket: { - "guardrails": ["policy-rule"], "guardrail_config": {"language": "en"}, - "applied_policies": ["parent-policy"], "policy_sources": {"parent-policy": "model"}, - "_guardrail_pipelines": [], "_pipeline_managed_guardrails": ["pipeline-rule"], "tags": ["review"], + "guardrails": ["policy-rule"], + "guardrail_config": {"language": "en"}, + "applied_policies": ["parent-policy"], + "policy_sources": {"parent-policy": "model"}, + "_guardrail_pipelines": [], + "_pipeline_managed_guardrails": ["pipeline-rule"], + "tags": ["review"], }, "guardrails": [{"request-rule": {"extra_body": {"threshold": 0.9}}}], "guardrail_config": {"entities": ["EMAIL_ADDRESS"]}, @@ -4292,13 +4322,26 @@ def test_mcp_conversion_honors_only_authenticated_global_guardrail_opt_outs(opt_ from litellm.responses.mcp.request_context import MCPRequestContext auth = UserAPIKeyAuth(metadata={"opted_out_global_guardrails": ["global-rule"] if opt_out else []}) - context = MCPRequestContext.resolve(kwargs={"metadata": { - "user_api_key_auth": auth, "disable_global_guardrails": True, - "user_api_key_metadata": {"disable_global_guardrails": True}, - }}, tools=None) + context = MCPRequestContext.resolve( + kwargs={ + "metadata": { + "user_api_key_auth": auth, + "disable_global_guardrails": True, + "user_api_key_metadata": {"disable_global_guardrails": True}, + } + }, + tools=None, + ) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) - kwargs = {"name": "execute", "arguments": {}, "user_api_key_auth": auth, "guardrail_context": context.guardrail_context} - synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs) + kwargs = { + "name": "execute", + "arguments": {}, + "user_api_key_auth": auth, + "guardrail_context": context.guardrail_context, + } + synthetic = proxy_logging._convert_mcp_to_llm_format( + proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs + ) guardrail = CustomGuardrail(guardrail_name="global-rule", event_hook="pre_mcp_call", default_on=True) assert guardrail.should_run_guardrail(synthetic, GuardrailEventHooks.pre_mcp_call) is (not opt_out) synthetic["metadata"]["user_api_key_metadata"]["opted_out_global_guardrails"].append("unrelated") @@ -4312,19 +4355,26 @@ def test_mcp_auth_policy_uses_original_request_model(monkeypatch, model, expecte from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails registry = policy_registry.PolicyRegistry() - registry._policies = {"model-policy": Policy( - condition=PolicyCondition(model="parent-model"), guardrails=PolicyGuardrails(add=["model-rule"]) - )} + registry._policies = { + "model-policy": Policy( + condition=PolicyCondition(model="parent-model"), guardrails=PolicyGuardrails(add=["model-rule"]) + ) + } registry._initialized = True monkeypatch.setattr(policy_registry, "_policy_registry", registry) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) proxy_logging = ProxyLogging(user_api_key_cache=DualCache()) kwargs = { - "name": "execute", "arguments": {}, + "name": "execute", + "arguments": {}, "user_api_key_auth": UserAPIKeyAuth(metadata={"policies": ["model-policy"]}), - "guardrail_context": MCPRequestContext.resolve_guardrail_context({"model": model, "guardrails": ["request-rule"]}), + "guardrail_context": MCPRequestContext.resolve_guardrail_context( + {"model": model, "guardrails": ["request-rule"]} + ), } - synthetic = proxy_logging._convert_mcp_to_llm_format(proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs) + synthetic = proxy_logging._convert_mcp_to_llm_format( + proxy_logging._create_mcp_request_object_from_kwargs(kwargs), kwargs + ) assert ("model-rule" in synthetic["metadata"]["guardrails"]) is expected assert "request-rule" in synthetic["metadata"]["guardrails"] @@ -4365,7 +4415,5 @@ def test_is_client_error_exception_existing_behaviour_unchanged() -> None: assert _is_client_error_exception(HTTPException(status_code=400, detail="x")) is True assert _is_client_error_exception(HTTPException(status_code=500, detail="x")) is False - assert _is_client_error_exception( - ProxyException(message="x", type="x", param=None, code="429") - ) is True + assert _is_client_error_exception(ProxyException(message="x", type="x", param=None, code="429")) is True assert _is_client_error_exception(ValueError("x")) is False diff --git a/tests/unit/proxy/test_response_polling_pre_call_checks.py b/tests/unit/proxy/test_response_polling_pre_call_checks.py index c6e4301a501..4228c1ff8f8 100644 --- a/tests/unit/proxy/test_response_polling_pre_call_checks.py +++ b/tests/unit/proxy/test_response_polling_pre_call_checks.py @@ -6,6 +6,8 @@ BEFORE a polling ID is created, so rate-limited requests get a synchronous error instead of a polling ID that immediately fails. """ +from collections.abc import Mapping +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -83,7 +85,7 @@ class TestSkipPreCallLogic: ): try: await processor.base_process_llm_request( - request=MagicMock(spec=Request), + request=MagicMock(spec=Request, headers={}), fastapi_response=MagicMock(spec=Response), user_api_key_dict=UserAPIKeyAuth(), route_type="aresponses", @@ -176,3 +178,77 @@ class TestPollingEndpointPreCallGuard: assert exc_info.value.status_code == 429 generate_polling_id_mock.assert_not_called() + + +async def _polled_background_response( + body: Mapping[str, object], headers: Mapping[str, str] +) -> tuple[MagicMock, Mapping[str, object]]: + from litellm.proxy.response_api_endpoints.endpoints import responses_api + from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler + + generate_polling_id: Final = MagicMock(return_value="litellm_poll_test") + pre_call_saw: Final[list[Mapping[str, object]]] = [] + + async def pre_call(processor: ProxyBaseLLMRequestProcessing, **_: object) -> tuple[Mapping[str, object], MagicMock]: + pre_call_saw.append(dict(processor.data)) + return processor.data, MagicMock() + + request: Final = MagicMock(spec=Request) + request.headers = headers + proxy_server_patches: Final[Mapping[str, object]] = { + "read_request_body": AsyncMock(return_value=dict(body)), + "general_settings": {}, + "llm_router": MagicMock(), + "native_background_mode": None, + "polling_cache_ttl": 3600, + "polling_via_cache_enabled": True, + "proxy_config": MagicMock(), + "proxy_logging_obj": AsyncMock(), + "redis_usage_cache": AsyncMock(), + "user_api_base": None, + "user_max_tokens": None, + "user_model": None, + "user_request_timeout": None, + "user_temperature": None, + "version": "1.0.0", + } + with ( + patch.multiple("litellm.proxy.proxy_server", **proxy_server_patches), + patch("litellm.proxy.response_polling.polling_handler.should_use_polling_for_request", return_value=True), + patch.object(ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", pre_call), + patch.object( + ProxyBaseLLMRequestProcessing, + "handle_llm_api_exception", + new_callable=AsyncMock, + side_effect=lambda e, **_: e, + ), + patch("litellm.proxy.response_api_endpoints.endpoints.raise_if_required_body_param_missing"), + patch.object(ResponsePollingHandler, "generate_polling_id", generate_polling_id), + patch("asyncio.create_task"), + patch.object(ResponsePollingHandler, "create_initial_state", new_callable=AsyncMock, return_value=MagicMock()), + ): + await responses_api( + request=request, fastapi_response=MagicMock(spec=Response), user_api_key_dict=MagicMock(spec=UserAPIKeyAuth) + ) + return generate_polling_id, pre_call_saw[0] if pre_call_saw else {} + + +@pytest.mark.asyncio +async def test_a_body_total_timeout_on_a_polled_background_response_is_refused_before_a_polling_id_is_issued() -> None: + from litellm.proxy._types import ProxyException + + with pytest.raises(ProxyException) as refused: + _ = await _polled_background_response({"model": "gpt-4", "background": True, "total_timeout": 10}, {}) + + assert refused.value.code == "400" + assert refused.value.param == "total_timeout" + + +@pytest.mark.asyncio +async def test_a_header_total_timeout_on_a_polled_background_response_is_dropped_and_polling_proceeds() -> None: + generate_polling_id, pre_call_saw = await _polled_background_response( + {"model": "gpt-4", "background": True, "total_timeout": "ignored"}, {"x-litellm-total-timeout": "10"} + ) + + generate_polling_id.assert_called_once() + assert "total_timeout" not in pre_call_saw diff --git a/tests/unit/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 25c62034f56..7f384d8e6bf 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,29 +1,24 @@ import asyncio import contextlib import gc +import json import threading import weakref -import httpx -import json -import litellm -import pytest -import redis -import respx - - from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping from datetime import datetime from types import SimpleNamespace -from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from uuid import uuid4 -from litellm.router_api import RouterSettings -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs - +import httpx +import pytest +import redis +import respx from fakeredis import TcpFakeServer from fastapi import HTTPException +import litellm from litellm.caching.caching import DualCache from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -37,11 +32,13 @@ from litellm.proxy.route_llm_request import ( user_config_routers, ) from litellm.proxy.utils import ProxyLogging +from litellm.router_api import RouterSettings from litellm.router_strategy.adaptive_router.config import ( ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, ADAPTIVE_ROUTER_RESPONSE_HEADER, ) from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError +from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs @pytest.mark.parametrize( @@ -79,7 +76,6 @@ async def test_route_request_dynamic_credentials(route_type, required_body_param @pytest.mark.asyncio async def test_route_request_proxy_admin_can_call_all_team_scoped_deployments_without_team_id(): import litellm - from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth router = litellm.Router( @@ -671,7 +667,6 @@ async def test_route_request_forwards_mock_params_when_opted_in(monkeypatch): """End-to-end through ``route_request``: with the opt-in set the param reaches the router, which is what makes a fallback drill possible.""" import litellm.proxy.proxy_server as proxy_server - from litellm.proxy.route_llm_request import MOCK_TESTING_CONFIG_KEY monkeypatch.setattr( @@ -776,10 +771,12 @@ async def test_route_request_realtime_wildcard_model_resolves_credentials( POST /realtime/client_secrets with a request model like openai/gpt-realtime must match an openai/* deployment and forward its api_key upstream. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -814,10 +811,12 @@ async def test_route_request_realtime_team_scoped_model_resolves_credentials( Team-scoped deployments (team_public_model_name) must be selected when user_api_key_team_id is present, same as /chat/completions. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -860,10 +859,12 @@ async def test_route_request_realtime_litellm_credential_name_resolves_api_key( litellm_credential_name on a wildcard deployment must resolve to the stored api_key when routing acreate_realtime_client_secret through the router. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -895,9 +896,9 @@ async def test_route_request_realtime_unresolvable_model_raises_not_found( An unknown model must not silently fall through to litellm with an empty OPENAI_API_KEY env var. """ - import litellm from unittest.mock import AsyncMock, patch + import litellm from litellm.proxy.route_llm_request import ProxyModelNotFoundError monkeypatch.delenv("OPENAI_API_KEY", raising=False) @@ -927,10 +928,12 @@ async def test_route_request_realtime_calls_resolves_api_base(monkeypatch): non-default (self-hosted / proxied) OpenAI endpoint is honored, instead of defaulting to https://api.openai.com. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) monkeypatch.delenv("OPENAI_API_BASE", raising=False) @@ -970,10 +973,12 @@ async def test_route_request_realtime_transcription_session_resolves_credentials /realtime/transcription_sessions must resolve credentials through the router (wildcard deployment) rather than falling back to an empty OPENAI_API_KEY. """ - import httpx - import litellm from unittest.mock import AsyncMock, patch + import httpx + + import litellm + monkeypatch.delenv("OPENAI_API_KEY", raising=False) router = litellm.Router( @@ -1400,7 +1405,9 @@ async def test_route_request_applies_override_fallbacks_when_request_carries_cre @pytest.mark.parametrize( "router_only_setting", [ - pytest.param({"model_group_retry_policy": {"gpt-4": {"TimeoutErrorRetries": 1}}}, id="model-group-retry-policy"), + pytest.param( + {"model_group_retry_policy": {"gpt-4": {"TimeoutErrorRetries": 1}}}, id="model-group-retry-policy" + ), pytest.param({"context_window_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="context-window-fallbacks"), pytest.param({"content_policy_fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}]}, id="content-policy-fallbacks"), ], @@ -2103,11 +2110,12 @@ async def test_route_request_read_through_disabled_without_store_model_in_db(mon @pytest.mark.asyncio async def test_route_request_read_through_supplies_db_model_default_for_missing_param(monkeypatch): - import litellm - import litellm.proxy.proxy_server as proxy_server from types import SimpleNamespace from unittest.mock import AsyncMock, patch + import litellm + import litellm.proxy.proxy_server as proxy_server + model_name = "e2e-db-only-max-tokens-default" router = litellm.Router( model_list=[{"model_name": "some-other-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}] @@ -2662,7 +2670,9 @@ async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() }, { **_mock_deployment("smart"), - "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + "model_info": { + "adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]} + }, }, ] }, @@ -3125,3 +3135,47 @@ def test_fallback_probe_skips_only_a_loaded_model_that_refuses_the_route() -> No assert skipped("embed-only") is True assert skipped("chat-fallback") is False assert skipped("added-on-another-worker") is False + + +def _router_answering_hi() -> "litellm.Router": + import litellm + + return litellm.Router( + model_list=[{"model_name": "g", "litellm_params": {"model": "openai/gpt-5.4-nano", "mock_response": "hi"}}] + ) + + +@pytest.mark.parametrize("flag", ["stream", "background"]) +@pytest.mark.asyncio +async def test_a_request_a_hook_made_stream_or_run_in_the_background_is_refused_under_a_body_total_timeout( + flag: str, +) -> None: + from litellm.litellm_core_utils.total_deadline import TotalTimeoutNotApplicable, total_deadline + + async with total_deadline(60, from_body=True): + with pytest.raises(TotalTimeoutNotApplicable): + await route_request( + data={"model": "g", "messages": [{"role": "user", "content": "hi"}], flag: True}, + llm_router=_router_answering_hi(), + user_model=None, + route_type="acompletion", + ) + + +@pytest.mark.asyncio +async def test_a_request_a_hook_turned_streaming_runs_unbounded_under_a_header_total_timeout() -> None: + from litellm.litellm_core_utils.total_deadline import current_total_deadline, total_deadline + + async with total_deadline(60, from_body=False): + llm_call: Final = await route_request( + data={"model": "g", "messages": [{"role": "user", "content": "hi"}], "stream": True}, + llm_router=_router_answering_hi(), + user_model=None, + route_type="acompletion", + ) + stream: Final = await llm_call + chunks: Final = [chunk async for chunk in stream] + deadline: Final = current_total_deadline() + + assert chunks + assert deadline is not None and deadline.allows_wait(10**9) diff --git a/tests/unit/test_router/core/test_retries_cooldowns.py b/tests/unit/test_router/core/test_retries_cooldowns.py index 004b1762240..d5ffa8491b1 100644 --- a/tests/unit/test_router/core/test_retries_cooldowns.py +++ b/tests/unit/test_router/core/test_retries_cooldowns.py @@ -13,6 +13,7 @@ import pytest import litellm from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, INITIAL_RETRY_DELAY, JITTER +from litellm.litellm_core_utils.total_deadline import TotalTimeoutNotApplicable, total_deadline from litellm.router_api import ChangeSet from litellm.router_core import RouterCore from litellm.types.router import AllowedFailsPolicy, RetryPolicy @@ -430,6 +431,60 @@ async def test_lone_deployment_waits_the_retry_after_header(make_core: MakeCore, assert all(7 <= slept < 7 + JITTER for slept in host.sleeps[:2]) +@pytest.mark.parametrize( + ("budget", "expected_calls"), + [(10, 2), (100, 4)], + ids=["second_backoff_outlasts_the_budget", "every_backoff_fits_the_budget"], +) +@pytest.mark.asyncio +async def test_a_retry_runs_only_while_its_backoff_fits_the_total_timeout( + make_core: MakeCore, host: ScriptedHost, budget: float, expected_calls: int +) -> None: + core: Final = make_core(model_list=[deployment("g", "a")], num_retries=3) + host.answer("openai/a", lambda **_: rate_limited(retry_after="6")) + + async with total_deadline(budget, from_body=True, clock=lambda: sum(host.sleeps)): + failure: Final = await attempt(core) + + assert isinstance(failure, litellm.RateLimitError) + assert len(host.calls_to("acompletion")) == expected_calls + assert sum(host.sleeps) < budget + + +@pytest.mark.asyncio +async def test_a_streaming_attempt_under_a_body_total_timeout_is_refused_before_the_provider_is_called( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + num_retries=3, + ) + host.answer("openai/a", lambda **_: litellm.ModelResponse()) + + async with total_deadline(10, from_body=True): + refused: Final = await attempt(core, stream=True) + + assert isinstance(refused, TotalTimeoutNotApplicable) + assert host.calls_to("acompletion") == [] + + +@pytest.mark.asyncio +async def test_a_streaming_attempt_under_a_header_total_timeout_runs_unbounded( + make_core: MakeCore, host: ScriptedHost +) -> None: + core: Final = make_core( + model_list=[deployment("g", "a")], + num_retries=3, + ) + host.answer("openai/a", lambda **_: rate_limited(retry_after="6")) + + async with total_deadline(10, from_body=False, clock=lambda: sum(host.sleeps)): + failure: Final = await attempt(core, stream=True) + + assert isinstance(failure, litellm.RateLimitError) + assert len(host.calls_to("acompletion")) == 4 + + @pytest.mark.asyncio async def test_lone_deployment_ignores_a_retry_after_header_over_a_minute( make_core: MakeCore, host: ScriptedHost