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