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

feat(proxy): total_timeout bounds a whole non-streaming request (#50496)

Browse source 
* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

* fix(proxy): a total_timeout cuts slow pre-call hooks short, holding only the rate limit reservation

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
joshua-berri 2026-11-08 00:43:50 -08:00 • committed by GitHub
parent 7ea4028625
commit 43112547fb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 1170 additions and 162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

197
litellm/litellm_core_utils/total_deadline.py Normal file
View file

@ -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

174
litellm/proxy/common_request_processing.py
View file

@ -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),

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

@ -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,

3
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -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,

2
litellm/proxy/route_llm_request.py
View file

@ -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(

12
litellm/proxy/utils.py
View file

@ -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]:

20
litellm/router_rust/core.py
View file

@ -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))

36
tests/unit/litellm_core_utils/test_total_deadline.py Normal file
View file

@ -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)

30
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -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,

1
tests/unit/proxy/search_endpoints/test_endpoints.py
View file

@ -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

351
tests/unit/proxy/test_common_request_processing.py
View file

@ -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

103
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -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=[

156
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
View file

@ -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

78
tests/unit/proxy/test_response_polling_pre_call_checks.py
View file

@ -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

112
tests/unit/proxy/test_route_llm_request.py
View file

@ -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)

55
tests/unit/test_router/core/test_retries_cooldowns.py
View file

@ -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