scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_add-claude-sonnet-5-pricing

Browse source
This commit is contained in:
mateo 2026-09-05 19:19:05 +00:00
commit 38106cc0f2
103 changed files with 4207 additions and 536 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -433,9 +433,9 @@ _default_detect_secrets_config = {
"name": "ZendeskSecretKeyDetector",
"path": _custom_plugins_path + "/zendesk_secret_key.py",
},
{"name": "Base64HighEntropyString", "limit": 3.0},
{"name": "Base64HighEntropyString", "limit": 4.5},
{"name": "HexHighEntropyString", "limit": 3.0},
]
],
}
@ -466,16 +466,19 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
os.remove(temp_file.name)
detected_secrets = []
for file in secrets.files:
for found_secret in secrets[file]:
if found_secret.secret_value is None:
continue
detected_secrets.append(
{"type": found_secret.type, "value": found_secret.secret_value}
)
return detected_secrets
return [
{"type": found_secret.type, "value": found_secret.secret_value}
for file in sorted(secrets.files)
for found_secret in sorted(
secrets[file],
key=lambda secret: (
-len(secret.secret_value or ""),
secret.type,
secret.secret_value or "",
),
)
if found_secret.secret_value is not None
]
def redact_text(self, text: str, source: str = "message") -> str:
"""Replace every detected secret in ``text`` with ``[REDACTED]`` and

15
enterprise/litellm_enterprise/enterprise_callbacks/secrets_plugins/openai_api_key.py
View file

@ -3,6 +3,7 @@ This plugin searches for OpenAI API Keys.
"""
import re
from collections.abc import Generator
from detect_secrets.plugins.base import RegexBasedDetector
@ -16,4 +17,16 @@ class OpenAIApiKeyDetector(RegexBasedDetector):
@property
def denylist(self) -> list[re.Pattern]:
return [re.compile(r"""(sk-[a-zA-Z0-9]{5,})""")]
return [
re.compile(
r"((?:(?<![a-zA-Z0-9])|(?<=%[0-9A-Fa-f]{2}))"
r"sk[-_]"
r"[a-zA-Z0-9_-]{5,}"
r"(?![a-zA-Z0-9_-]))"
)
]
def analyze_string(self, string: str) -> Generator[str, None, None]:
# the digit check lives outside the regex: a lookahead re-scans the token
# from every `sk` inside it, which is quadratic on `-sk-sk-sk-...` input
yield from (match for match in super().analyze_string(string) if re.search(r"[0-9]", match))

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.64"
version = "0.1.65"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.64"
version = "0.1.65"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.93"
version = "0.4.94"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.93"
version = "0.4.94"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

6
litellm/constants.py
View file

@ -40,7 +40,7 @@ ROUTER_SETTINGS_MANAGED_OUTSIDE_CONFIG: Final[frozenset[str]] = frozenset(
"router_general_settings",
"ignore_invalid_deployments",
"fallback_access_check",
"heuristic_v2_router_limit",
"auto_router_capability_limit",
}
)
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
@ -89,6 +89,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS: Final = 256 * 1024
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
# in-memory stand-in handed to provider converters for redacted arguments; never stored
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
@ -215,6 +216,9 @@ MAX_CALLBACKS: Final = get_env_int("LITELLM_MAX_CALLBACKS", 100)
# so the deployment-level hook does not re-run them for the same request
PRE_CALL_EXECUTED_GUARDRAILS_KEY: Final = "_pre_call_executed_guardrails"
# Attribute stamped on log_guardrail_information wrappers so __init_subclass__ does not wrap them again
LOGS_GUARDRAIL_INFORMATION_MARKER: Final = "_litellm_logs_guardrail_information"
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET: Final = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)

6
litellm/integrations/cloudzero/cz_stream_api.py
View file

@ -97,7 +97,11 @@ class CloudZeroStreamer:
continue
# Convert lists back to DataFrames
return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records}
return {
date_key: pl.DataFrame(records, infer_schema_length=None)
for date_key, records in daily_batches.items()
if records
}
def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime:
"""Parse timestamp string and convert to UTC."""

2
litellm/integrations/cloudzero/transform.py
View file

@ -95,7 +95,7 @@ class CBFTransformer:
if len(cbf_data) > 0:
console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]")
return pl.DataFrame(cbf_data)
return pl.DataFrame(cbf_data, infer_schema_length=None)
def _create_cbf_record(self, row: dict[str, object]) -> CBFRecord:
"""Create a single CBF record from LiteLLM daily spend row."""

9
litellm/integrations/custom_guardrail.py
View file

@ -46,6 +46,7 @@ dc: Final = DualCache()
from litellm.constants import (
GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS,
LOGS_GUARDRAIL_INFORMATION_MARKER,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
from litellm.exceptions import (
@ -151,6 +152,13 @@ class CustomGuardrail(CustomLogger):
records_own_guardrail_information: ClassVar[bool] = False
def __init_subclass__(cls, **kwargs: object) -> None: # kwargs-ok: forwarded to cooperative __init_subclass__ hooks
super().__init_subclass__(**kwargs)
own_apply_guardrail: Final = cls.__dict__.get("apply_guardrail")
if own_apply_guardrail is None or LOGS_GUARDRAIL_INFORMATION_MARKER in vars(own_apply_guardrail):
return
cls.apply_guardrail = log_guardrail_information(own_apply_guardrail)
def __init__(
self,
guardrail_name: str | None = None,
@ -1559,4 +1567,5 @@ def log_guardrail_information(func):
return async_wrapper(*args, **kwargs)
return sync_wrapper(*args, **kwargs)
vars(wrapper)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the wrapper this call just built
return wrapper

24
litellm/integrations/shadow_eval_logger.py
View file

@ -61,9 +61,9 @@ _MAX_CONCURRENT_SHADOW_TASKS: Final = 16
_MAX_JUDGE_RESPONSE_CHARS: Final = 8_000
_MAX_JUDGE_PROMPT_CHARS: Final = 24_000
# The judge answers with a small JSON object; a tighter budget truncates the JSON
# mid-object and the attempt is lost to an error row.
JUDGE_MAX_OUTPUT_TOKENS: Final = 1500
# Covers the judge's reasoning tokens as well as its small JSON answer: a judge deployment
# carrying an elevated reasoning_effort spends a tight cap before it ever answers.
JUDGE_MAX_OUTPUT_TOKENS: Final = 4096
_MAX_ERROR_CHARS: Final = 500
@ -419,6 +419,20 @@ def _failure_detail(e: BaseException) -> str:
return f"{type(e).__name__}{location}: {e}"
def _judge_reply_shape(response: object) -> str:
"""How an unparseable judge reply was shaped. The parser's own message cannot separate a
judge that answered with nothing from one truncated mid-object, and those want opposite
fixes. Shape only, never the reply text: the judge quotes the sampled turns it compares,
and no attempt row carries sampled content today."""
read: Final = _chat_message_reader(response)
if read is None:
return "unreadable judge reply"
content: Final = read("content")
served: Final = str(_field_reader(response)("model") or "unknown")
body: Final = f"{len(str(content))} chars" if content else "no content"
return f"finish_reason={_chat_finish_reason(response)}, content={body}, model={served}"
def _call_cost(response: object) -> float:
"""Price one eval-arm call with the figure the spend pipeline bills: the router client
stamps _hidden_params.response_cost from the deployment's own pricing, which the public
@ -1266,7 +1280,9 @@ class ShadowEvalLogger(CustomLogger):
verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw))
except Exception as e: # noqa: BLE001 # malformed verdicts become error rows
verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e)
return _CallFailure(f"unparseable judge verdict: {e}", cost=_call_cost(response))
return _CallFailure(
f"unparseable judge verdict: {e}; {_judge_reply_shape(response)}", cost=_call_cost(response)
)
return _JudgeVerdict(
preference=_unmask_preference(verdict.preference, real_is_a),
confidence=max(0.0, min(1.0, verdict.confidence)),

40
litellm/litellm_core_utils/litellm_logging.py
View file

@ -78,7 +78,10 @@ from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import (
InteractionsUsageObjectTransformation,
)
from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages
from litellm.litellm_core_utils.logging_utils import (
truncate_base64_in_messages,
truncate_base64_in_messages_async,
)
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_custom_logger,
@ -538,6 +541,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
self.initialize_standard_built_in_tools_params(kwargs)
)
self.truncated_messages_for_logging: str | list | dict | None = None # mutable-ok: logged messages shape
## TIME TO FIRST TOKEN LOGGING ##
self.completion_start_time: datetime.datetime | None = None
self._llm_caching_handler: LLMCachingHandler | None = None
@ -1820,6 +1824,7 @@ class Logging(LiteLLMLoggingBaseClass):
and litellm_params.get(CallTypes.aanthropic_messages.value, False) is not True
and litellm_params.get(CallTypes.agenerate_content.value, False) is not True
and litellm_params.get(CallTypes.agenerate_content_stream.value, False) is not True
and litellm_params.get(CallTypes.arealtime.value, False) is not True
)
def _is_assembled_stream_success(self, result=None) -> bool:
@ -1913,7 +1918,9 @@ class Logging(LiteLLMLoggingBaseClass):
two paths cannot mutate it at the same time. ``prefer_async_handlers`` only
bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from
``completion()``); legacy string callbacks still run via
``executor.submit(failure_handler)`` when configured.
``executor.submit(failure_handler)`` when configured, and still get submitted
when the awaiting task is cancelled (e.g. the event loop shuts down right after
the request failed).
"""
litellm_params: Final = self.model_call_details.get("litellm_params", {}) or {}
sync_sdk: Final = self._is_sync_litellm_request(litellm_params)
@ -1922,12 +1929,11 @@ class Logging(LiteLLMLoggingBaseClass):
self.failure_handler(exception, traceback_exception)
return
await self.async_failure_handler(exception, traceback_exception)
if not self._should_run_sync_failure_callbacks_for_async_calls():
return
executor.submit(self.failure_handler, exception, traceback_exception)
try:
await self.async_failure_handler(exception, traceback_exception)
finally:
if self._should_run_sync_failure_callbacks_for_async_calls():
executor.submit(self.failure_handler, exception, traceback_exception)
def should_run_logging(
self,
@ -2932,6 +2938,11 @@ class Logging(LiteLLMLoggingBaseClass):
result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above
result.usage = batch_result.usage
self.truncated_messages_for_logging = await truncate_base64_in_messages_async(
StandardLoggingPayloadSetup.append_system_prompt_messages(
kwargs=self.model_call_details, messages=self.model_call_details.get("messages")
)
)
start_time, end_time, result = self._success_handler_helper_fn(
start_time=start_time,
end_time=end_time,
@ -3224,8 +3235,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details = {}
if (
self.model_call_details.get("log_event_type") == "failed_api_call"
and self.model_call_details.get("exception") is exception
self.model_call_details.get("exception") is exception
and self.model_call_details.get("standard_logging_object") is not None
):
return start_time, self.model_call_details["end_time"]
@ -6201,9 +6211,13 @@ def get_standard_logging_object_payload(
model_id=_model_id,
requester_ip_address=clean_metadata.get("requester_ip_address", None),
user_agent=clean_metadata.get("user_agent", None),
messages=truncate_base64_in_messages(
StandardLoggingPayloadSetup.append_system_prompt_messages(
kwargs=kwargs, messages=kwargs.get("messages")
messages=(
logging_obj.truncated_messages_for_logging
if logging_obj.truncated_messages_for_logging is not None
else truncate_base64_in_messages(
StandardLoggingPayloadSetup.append_system_prompt_messages(
kwargs=kwargs, messages=kwargs.get("messages")
)
)
),
response=final_response_obj,

40
litellm/litellm_core_utils/logging_utils.py
View file

@ -3,12 +3,15 @@ import functools
import inspect
import re
import time
from collections.abc import Mapping
from collections.abc import Iterator, Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
from litellm._logging import verbose_logger
from litellm.constants import MAX_BASE64_LENGTH_FOR_LOGGING
from litellm.constants import (
BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS,
MAX_BASE64_LENGTH_FOR_LOGGING,
)
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -141,6 +144,39 @@ def truncate_base64_in_messages(
return messages
_StringTree = str | Sequence["_StringTree"] | Mapping[str, "_StringTree"] | None
def _iter_string_leaves(value: _StringTree) -> Iterator[str]:
stack: Final[list[_StringTree]] = [value] # mutable-ok: explicit stack, recursive functions are banned in litellm/
while stack:
match stack.pop():
case str() as text:
yield text
case Mapping() as mapping:
stack.extend(mapping.values())
case Sequence() as items:
stack.extend(items)
case None:
pass
async def truncate_base64_in_messages_async(
messages: str | list | dict | None, # mutable-ok: same contract as truncate_base64_in_messages
) -> str | list | dict | None: # mutable-ok: same contract as truncate_base64_in_messages
"""
Same result as truncate_base64_in_messages, but payloads whose string content
reaches BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS are scanned in a worker
thread so the regex pass over multi-MB base64 images does not block the event loop.
"""
if messages is None or MAX_BASE64_LENGTH_FOR_LOGGING <= 0:
return messages
total_chars: Final = sum(len(leaf) for leaf in _iter_string_leaves(messages))
if total_chars < BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS:
return truncate_base64_in_messages(messages)
return await asyncio.to_thread(truncate_base64_in_messages, messages)
# Global service logger instance to avoid recreating it
_service_logger = None

8
litellm/litellm_core_utils/realtime_errors.py
View file

@ -29,3 +29,11 @@ def websocket_close_reason(message: str, fallback: str) -> str:
if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES:
return message
return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore")
def client_close_code(upstream_code: int) -> int:
from websockets.frames import EXTERNAL_CLOSE_CODES, CloseCode
if upstream_code in EXTERNAL_CLOSE_CODES or 3000 <= upstream_code < 5000:
return upstream_code
return int(CloseCode.INTERNAL_ERROR)

220
litellm/litellm_core_utils/realtime_streaming.py
View file

@ -1,12 +1,15 @@
import asyncio
import json
from collections.abc import Mapping, Sequence
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
import traceback
from collections.abc import Coroutine, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum, auto
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast
from typing_extensions import ReadOnly
import litellm
from litellm._logging import verbose_logger
from litellm._logging import redact_internal_details_from_client_message, verbose_logger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
from litellm.types.llms.openai import (
@ -19,9 +22,11 @@ from litellm.types.llms.openai import (
from litellm.types.realtime import ALL_DELTA_TYPES
from .litellm_logging import Logging as LiteLLMLogging
from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason
if TYPE_CHECKING:
from websockets.asyncio.client import ClientConnection
from websockets.exceptions import ConnectionClosed
from litellm.types.guardrails import GuardrailEventHooks
@ -30,8 +35,30 @@ else:
CLIENT_CONNECTION_CLASS = Any
class _ClientWebSocketExceptions(Protocol):
ConnectionClosed: type[Exception]
REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged"
@dataclass(frozen=True, slots=True)
class BackendClose:
code: int
reason: str
@property
def message(self) -> str:
if not self.reason:
return f"upstream websocket closed with code {self.code}"
return f"upstream websocket closed with code {self.code}: {self.reason}"
class ClientLoopExit(Enum):
CLIENT_DISCONNECTED = auto()
BACKEND_CLOSED = auto()
def backend_close_from(error: "ConnectionClosed") -> BackendClose:
if error.rcvd is None:
return BackendClose(code=1006, reason=str(error))
return BackendClose(code=error.rcvd.code, reason=error.rcvd.reason)
class _ASGIScope(TypedDict, total=False):
@ -69,10 +96,13 @@ class _ScopedWebSocket(Protocol):
class _ClientWebSocket(_ScopedWebSocket, Protocol):
exceptions: _ClientWebSocketExceptions
async def send_text(self, data: str) -> None: ...
async def receive_text(self) -> str: ...
async def close(self, code: int = 1000, reason: str | None = None) -> None: ...
class _LoggingWorker(Protocol):
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: ...
def _decode_json_object(payload: str) -> Mapping[str, object]:
@ -108,11 +138,14 @@ class RealTimeStreaming:
backend_uses_beta_protocol: bool | None = None,
force_transcription_model: str | None = None,
event_normalizer: RealtimeEventNormalizer | None = None,
logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER,
):
self.websocket: _ClientWebSocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self._logging_worker = logging_worker
self.messages: list[OpenAIRealtimeEvents] = []
self._backend_sent_frames: bool = False
self.input_message: dict = {}
self.input_messages: list[dict[str, str]] = []
self.session_tools: list[dict] = []
@ -388,9 +421,10 @@ class RealTimeStreaming:
# Route through the bounded logging worker (per-coroutine timeout +
# concurrency cap) instead of a bare create_task, so a slow callback
# can't leave suspended tasks pinning each call's response in memory.
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
self._logging_worker.ensure_initialized_and_enqueue(
self.logging_obj.dispatch_success_handlers(self.messages, prefer_async_handlers=True)
)
self.logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
async def _send_to_backend(self, message: str) -> bool:
"""Send a message to the backend WebSocket.
@ -1035,60 +1069,84 @@ class RealTimeStreaming:
return True
return False
async def backend_to_client_send_messages(self):
async def _relay_backend_messages(self) -> NoReturn:
while True:
try:
raw_response = await self.backend_ws.recv(decode=False)
except TypeError:
raw_response = await self.backend_ws.recv()
self._backend_sent_frames = True
if isinstance(raw_response, bytes):
try:
raw_response = raw_response.decode("utf-8")
except UnicodeDecodeError:
verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.")
continue
if self.provider_config:
try:
await self._handle_provider_config_message(raw_response)
except Exception as e:
verbose_logger.exception("Error processing backend message, skipping: %s", e)
continue
else:
event = self._parse_backend_event(raw_response)
if event is None:
await self.websocket.send_text(raw_response)
continue
if self._should_drop_event_from_client(event):
continue
if await self._handle_raw_backend_message(event, raw_response):
continue
event = self._normalize_event_for_ga_client(event)
self.store_message(event)
if not self._client_wants_beta:
await self.websocket.send_text(json.dumps(event))
continue
translated = self._translate_event_to_beta(event)
if translated is None:
continue
await self.websocket.send_text(json.dumps(translated))
async def backend_to_client_send_messages(self) -> BackendClose:
import websockets
try:
while True:
try:
raw_response = await self.backend_ws.recv(decode=False)
except TypeError:
raw_response = await self.backend_ws.recv()
if isinstance(raw_response, bytes):
try:
raw_response = raw_response.decode("utf-8")
except UnicodeDecodeError:
verbose_logger.warning("Received non-UTF-8 binary frame from backend, skipping.")
continue
if self.provider_config:
try:
await self._handle_provider_config_message(raw_response)
except Exception as e:
verbose_logger.exception("Error processing backend message, skipping: %s", e)
continue
else:
event = self._parse_backend_event(raw_response)
if event is None:
await self.websocket.send_text(raw_response)
continue
if self._should_drop_event_from_client(event):
continue
if await self._handle_raw_backend_message(event, raw_response):
continue
event = self._normalize_event_for_ga_client(event)
self.store_message(event)
if not self._client_wants_beta:
await self.websocket.send_text(json.dumps(event))
continue
translated = self._translate_event_to_beta(event)
if translated is None:
continue
await self.websocket.send_text(json.dumps(translated))
await self._relay_backend_messages()
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.exception("Connection closed in backend to client send messages - %s", e)
except Exception as e:
verbose_logger.exception("Error in backend to client send messages: %s", e)
finally:
close: Final = backend_close_from(e)
self._flush_unbilled_transcription_usage()
if self._backend_refused_session(close):
await self.log_backend_refusal(e)
else:
await self.log_messages()
return close
except asyncio.CancelledError:
self._flush_unbilled_transcription_usage()
await self.log_messages()
raise
except Exception as e:
verbose_logger.exception("Error in backend to client send messages: %s", e)
self._flush_unbilled_transcription_usage()
await self.log_messages()
return BackendClose(code=1011, reason="proxy failed while relaying the upstream websocket")
def _backend_refused_session(self, close: BackendClose) -> bool:
return close.code != 1000 and not self._backend_sent_frames
async def log_backend_refusal(self, error: Exception) -> None:
if not self.logging_obj:
return
self._logging_worker.ensure_initialized_and_enqueue(
self.logging_obj.dispatch_failure_handlers(error, traceback.format_exc(), prefer_async_handlers=True)
)
@staticmethod
def _detect_beta_header(websocket: _ScopedWebSocket) -> bool:
@ -1243,11 +1301,22 @@ class RealTimeStreaming:
item["content"] = new_content
return item
async def client_ack_messages(self):
async def _receive_client_message(self) -> str | None:
try:
return await self.websocket.receive_text()
except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone
verbose_logger.debug("Client disconnected: %s", e)
return None
async def client_ack_messages(self) -> ClientLoopExit:
import websockets
client_event: _ClientEventFrame
try:
while True:
message = await self.websocket.receive_text()
message = await self._receive_client_message()
if message is None:
return ClientLoopExit.CLIENT_DISCONNECTED
## GUARDRAIL: intercept conversation.item.create for text-based injection.
guardrail_turn_detection_injected = False
@ -1481,23 +1550,38 @@ class RealTimeStreaming:
if guardrail_turn_detection_injected and sent:
self._guardrail_turn_detection_update_sent = True
except websockets.exceptions.ConnectionClosed as e:
verbose_logger.debug("Backend closed while forwarding a client message: %s", e)
return ClientLoopExit.BACKEND_CLOSED
except Exception as e:
verbose_logger.debug("Error in client ack messages: %s", e)
return ClientLoopExit.CLIENT_DISCONNECTED
async def bidirectional_forward(self):
async def bidirectional_forward(self) -> None:
forward_task: Final = asyncio.create_task(self.backend_to_client_send_messages())
client_task: Final = asyncio.create_task(self.client_ack_messages())
try:
await self.client_ack_messages()
except self.websocket.exceptions.ConnectionClosed:
verbose_logger.debug("Connection closed")
forward_task.cancel()
await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED)
if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED:
return
await self._close_client(await forward_task)
finally:
if not forward_task.done():
forward_task.cancel()
try:
await forward_task
except asyncio.CancelledError:
pass
forward_task.cancel()
client_task.cancel()
await asyncio.gather(forward_task, client_task, return_exceptions=True)
async def _close_client(self, close: BackendClose) -> None:
redacted_message: Final = redact_internal_details_from_client_message(close.message)
redacted_reason: Final = redact_internal_details_from_client_message(close.reason)
try:
if close.code != 1000:
await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error"))
await self.websocket.close(
code=client_close_code(close.code),
reason=websocket_close_reason(redacted_reason, fallback=redacted_message),
)
except Exception as e: # noqa: BLE001 # the client may already be gone; the session is over either way
verbose_logger.debug("Could not relay the upstream close to the client: %s", e)
def client_sent_openai_beta_realtime_header(websocket: _ScopedWebSocket) -> bool:

14
litellm/proxy/_experimental/mcp_server/exceptions.py
View file

@ -5,6 +5,20 @@ from typing import Final
from fastapi import HTTPException
class MCPServerURLCredentialsError(HTTPException):
"""A fixed, sanitized URL-credential migration error safe for operator previews."""
def __init__(self) -> None:
super().__init__(
status_code=500,
detail=(
"misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; "
"remove them from the URL and configure Basic Auth with auth_type: basic and "
"auth_value: username:password"
),
)
class MCPUpstreamAuthError(Exception):
"""Raised when an upstream MCP server returns an authentication failure
(typically HTTP 401) and the gateway should surface it transparently to

3
litellm/proxy/_experimental/mcp_server/outbound_credentials/adapter.py
View file

@ -19,6 +19,7 @@ from pydantic import SecretStr
from typing_extensions import assert_never
from litellm.experimental_mcp_client.client import strip_auth_scheme, to_basic_credentials
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
from litellm.proxy._experimental.mcp_server.oauth_utils import resolve_upstream_resource
from litellm.proxy._experimental.mcp_server.outbound_credentials.types import (
DEFAULT_CREDENTIAL_HEADER,
@ -293,6 +294,8 @@ def raise_public(error: CredError) -> NoReturn:
)
case "misconfigured":
raise HTTPException(status_code=500, detail=error.summary)
case "url_credentials_not_allowed":
raise MCPServerURLCredentialsError()
case "upstream_unavailable":
raise HTTPException(status_code=503, detail=error.summary)
case "unsupported_mode":

11
litellm/proxy/_experimental/mcp_server/outbound_credentials/resolver.py
View file

@ -134,7 +134,7 @@ class UpstreamCredentialProvider:
async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Result[httpx.Auth, CredError]:
match server.config:
case NoneConfig():
return Ok(NoOpAuth())
return self._none(server)
case ApiKeyConfig() as config:
return self._api_key(config)
case PassthroughConfig():
@ -151,6 +151,15 @@ class UpstreamCredentialProvider:
return _not_implemented(AuthSpecKind.aws_sigv4)
assert_never(server.config)
def _none(self, server: ServerSpec) -> Result[httpx.Auth, CredError]:
try:
resource: Final = httpx.URL(server.resource)
except httpx.InvalidURL:
return Ok(NoOpAuth())
if resource.userinfo:
return Error(CredError.of_url_credentials_not_allowed())
return Ok(NoOpAuth())
async def has_user_token(self, subject: Subject, server: ServerSpec) -> bool:
"""Whether a usable per-user token exists for this server (the preemptive 401's check).

12
litellm/proxy/_experimental/mcp_server/outbound_credentials/types.py
View file

@ -95,6 +95,7 @@ class CredError:
tag: Literal[
"unauthorized",
"misconfigured",
"url_credentials_not_allowed",
"upstream_unavailable",
"unsupported_mode",
"precondition_required",
@ -103,6 +104,7 @@ class CredError:
unauthorized: Unauthorized = case() # no usable credential for this (subject, server) -> 401 challenge
misconfigured: str = case() # the declared mode is missing required config -> 5xx (operator)
url_credentials_not_allowed: None = case()
upstream_unavailable: str = case() # the IdP / token endpoint could not be reached -> 503
unsupported_mode: str = case() # a raw mode string did not parse into AuthSpecKind (boundary)
precondition_required: str = case() # a required per-user value (e.g. an env var) has not been provided -> 412
@ -129,6 +131,10 @@ class CredError:
def of_misconfigured(detail: str) -> CredError:
return CredError(misconfigured=detail)
@staticmethod
def of_url_credentials_not_allowed() -> CredError:
return CredError(url_credentials_not_allowed=None)
@staticmethod
def of_upstream_unavailable(detail: str) -> CredError:
return CredError(upstream_unavailable=detail)
@ -154,6 +160,12 @@ class CredError:
return f"unauthorized: {self.unauthorized.detail}"
case "misconfigured":
return f"misconfigured: {self.misconfigured}"
case "url_credentials_not_allowed":
return (
"misconfigured: auth_type none cannot be used with credentials embedded in the upstream URL; "
"remove them from the URL and configure Basic Auth with auth_type: basic and "
"auth_value: username:password"
)
case "upstream_unavailable":
return f"upstream unavailable: {self.upstream_unavailable}"
case "unsupported_mode":

3
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -18,6 +18,7 @@ from litellm.exceptions import (
)
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
MCPServerURLCredentialsError,
MCPUpstreamAuthError,
)
from litellm.proxy._experimental.mcp_server.faults.list_outcomes import (
@ -75,6 +76,8 @@ _MCP_GUARDRAIL_REJECTIONS: Final = (
def _connection_error_message(exc: BaseException, url: str | None, timeout_seconds: float) -> str:
if isinstance(exc, MCPServerURLCredentialsError):
return str(exc.detail)
if isinstance(exc, TimeoutError):
return (
f"Failed to connect to MCP server: no response from {url or 'the server'} "

11
litellm/proxy/auth/litellm_license.py
View file

@ -17,7 +17,7 @@ if TYPE_CHECKING:
AUTO_ROUTER_LICENSE_FEATURE: Final = "auto_router"
HEURISTIC_V2_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
AUTO_ROUTER_LICENSE_REMEDY: Final = "A LiteLLM license with the 'auto_router' feature lifts the limit."
class LicenseCheck:
@ -153,11 +153,12 @@ class LicenseCheck:
return False
return team_count > _max_teams_in_license
def heuristic_v2_router_limit(self) -> int | None:
def auto_router_capability_limit(self) -> int | None:
"""
How many heuristic_v2 auto-routers this proxy may hold: unlimited (None) only when the
signed license lists the auto_router feature, otherwise one. A license verified through
the API carries no feature list, so it does not lift the limit either.
How many auto-routers may claim each licensed capability (heuristic_v2, operator-defined
tier_definitions): unlimited (None) only when the signed license lists the auto_router
feature, otherwise one per capability. A license verified through the API carries no
feature list, so it does not lift the limit either.
"""
if self.airgapped_license_data is None:
return 1

6
litellm/proxy/db/exception_handler.py
View file

@ -199,6 +199,12 @@ class PrismaDBExceptionHandler:
return True
return False
@staticmethod
def is_prisma_error(e: Exception) -> bool:
import prisma
return isinstance(e, _exception_types(prisma.errors.PrismaError))
@staticmethod
def is_deadlock_error(e: Exception) -> bool:
"""True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""

157
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -50,7 +50,7 @@ from litellm.proxy._types import (
TeamModelDeleteRequest,
UserAPIKeyAuth,
)
from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.config_sync_pubsub import (
coordination_redis_cache,
@ -98,11 +98,13 @@ from litellm.router_strategy.complexity_router import (
normalize_classification_prompt,
)
from litellm.router_utils.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
capability_limit_violation,
carries_complexity_router_settings,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
uses_heuristic_v2_classifier,
count_capability_routers,
gated_capability_of,
is_complexity_router_model,
validate_complexity_router_config_placement,
validate_complexity_router_config_write,
validate_strategy_router_model_write,
@ -237,11 +239,13 @@ def _strategy_router_write_violation(
An auto-router deployment's ``litellm_params.model`` (``auto_router/...``) is
the discriminator the router loads it by; a write that mangles it makes the
router drop the deployment silently under ``ignore_invalid_deployments``.
Only writes that supply ``litellm_params.model`` are judged on the naming
contract, against the merged (stored + incoming) params, so partial patches
and restores of an already-corrupted row stay legal. A config is judged only
when the write carries one, for the same reason: a rename must not be held
hostage by a stored config it does not touch. Returns the violation, or None.
A patch adding auto-router settings is judged against the effective model,
decrypting the stored model when the patch omits it, so a regular deployment
cannot claim a strategy-router configuration. Unrelated partial patches and
restores that do not touch strategy-router settings stay legal. A config is
judged only when the write carries one, for the same reason: a rename must
not be held hostage by a stored config it does not touch. Returns the
violation, or None.
"""
if incoming_params is None:
return None
@ -256,14 +260,18 @@ def _strategy_router_write_violation(
for source in (incoming_params, existing_params)
if source is not None and getattr(source, field, None) is not None
)
# Scope reads the incoming model because the stored one is encrypted at rest.
if carries_complexity_router_settings(incoming_params.model, present_fields):
effective_params: Final = _effective_complexity_router_params(incoming_params, existing_params)
effective_model: Final = effective_params.get("model")
if carries_complexity_router_settings(
effective_model if isinstance(effective_model, str) else None, present_fields
):
placement_violation: Final = validate_complexity_router_config_placement(incoming_params.model_extra)
if placement_violation is not None:
return placement_violation
if incoming_params.model is None:
return None
return validate_strategy_router_model_write(model=incoming_params.model, present_fields=present_fields)
return validate_strategy_router_model_write(
model=effective_model if isinstance(effective_model, str) else "",
present_fields=present_fields,
)
def _raise_on_strategy_router_write_violation(
@ -281,14 +289,23 @@ def _raise_on_strategy_router_write_violation(
)
HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301
_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
_HEURISTIC_V2_DB_ROWS_SQL: Final = """
SELECT count(*)::int AS held FROM "LiteLLM_ProxyModelTable"
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301
_CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
_STORED_LITELLM_PARAMS_SQL: Final = (
"(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)"
)
_STORED_COMPLEXITY_CONFIG_SQL: Final = f"{_STORED_LITELLM_PARAMS_SQL} -> 'complexity_router_config'"
_CAPABILITY_DB_ROWS_SQL: Final[Mapping[str, str]] = MappingProxyType(
{
capability.key: f"""
SELECT {_STORED_LITELLM_PARAMS_SQL} ->> 'model' AS model
FROM "LiteLLM_ProxyModelTable"
WHERE model_id <> $1
AND (CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)
-> 'complexity_router_config' ->> 'classifier_type' = 'heuristic_v2'
AND ({capability.sql_config_predicate.format(config=_STORED_COMPLEXITY_CONFIG_SQL)})
"""
for capability in GATED_AUTO_ROUTER_CAPABILITIES
}
)
def _effective_complexity_router_config(
@ -301,13 +318,44 @@ def _effective_complexity_router_config(
return existing_params.complexity_router_config
@asynccontextmanager
async def _heuristic_v2_slot(
prisma_client: PrismaClient, *, effective_config: object, model_id: str | None
) -> AsyncGenerator[_ProxyModelTable, None]:
"""Hand out the model table to write through while the row's claim on a heuristic_v2 slot is settled.
def _effective_model(
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
) -> str | None:
"""The model a write leaves on the row, decrypting an existing value only when the patch omits it."""
incoming: Final = None if incoming_params is None else incoming_params.model
if incoming is not None:
return incoming
existing: Final = None if existing_params is None else existing_params.model
if existing is None:
return None
decrypted: Final = decrypt_value_helper(
value=existing,
key="model",
exception_type="debug",
return_original_value=True,
)
return decrypted if isinstance(decrypted, str) else None
A write that leaves the row on classifier_type heuristic_v2 under a limited license runs
def _effective_complexity_router_params(
incoming_params: GenericLiteLLMParams | None, existing_params: GenericLiteLLMParams | None
) -> Mapping[str, object]:
"""The model and complexity config a write leaves, for placement and capability decisions."""
return MappingProxyType(
{
"model": _effective_model(incoming_params, existing_params),
"complexity_router_config": _effective_complexity_router_config(incoming_params, existing_params),
}
)
@asynccontextmanager
async def _auto_router_capability_slot(
prisma_client: PrismaClient, *, effective_params: Mapping[str, object], model_id: str | None
) -> AsyncGenerator[_ProxyModelTable, None]:
"""Hand out the model table to write through while the row's claim on a licensed capability is settled.
A write that leaves the row claiming a licensed capability under a limited license runs
inside one transaction that takes an advisory lock in its own statement before counting
(a statement's snapshot predates anything it locks), so pods cannot both pass the count:
the DB rows (any pod, either JSON shape) plus this proxy's config.yaml routers are judged
@ -321,21 +369,37 @@ async def _heuristic_v2_slot(
"""
from litellm.proxy.proxy_server import _license_check, llm_router
limit: Final = _license_check.heuristic_v2_router_limit()
if limit is None or not uses_heuristic_v2_classifier(effective_config):
limit: Final = _license_check.auto_router_capability_limit()
capability: Final = gated_capability_of(effective_params)
if limit is None or capability is None:
yield _proxy_model_table(prisma_client)
return
async with prisma_client.db.tx() as tx_ctx:
tables: Final[_TxModelTables] = tx_ctx
await tx_ctx.query_raw(_HEURISTIC_V2_LOCK_SQL, HEURISTIC_V2_SLOT_LOCK_KEY)
rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(_HEURISTIC_V2_DB_ROWS_SQL, model_id or "")
db_held: Final = rows[0].get("held") if rows else 0
await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY)
rows: Sequence[Mapping[str, object]] = await tx_ctx.query_raw(
_CAPABILITY_DB_ROWS_SQL[capability.key], model_id or ""
)
db_held: Final = sum(
1
for row in rows
for stored_model in (row.get("model"),)
if isinstance(stored_model, str)
and is_complexity_router_model(
decrypt_value_helper(
value=stored_model,
key="model",
exception_type="debug",
return_original_value=True,
)
)
)
config_rows: Final = () if llm_router is None else tuple(llm_router.config_deployments())
held: Final = (db_held if isinstance(db_held, int) else 0) + count_heuristic_v2_routers(config_rows)
violation: Final = heuristic_v2_limit_violation(held=held + 1, limit=limit)
held: Final = db_held + count_capability_routers(config_rows, capability=capability)
violation: Final = capability_limit_violation(capability=capability, held=held + 1, limit=limit)
if violation is not None:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {HEURISTIC_V2_LICENSE_REMEDY}"
status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}"
)
yield tables.litellm_proxymodeltable
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
@ -791,6 +855,9 @@ async def patch_model(
existing_params=db_model.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
patch_data.litellm_params, db_model.litellm_params
)
requested_model_name: Final = patch_data.model_name
stored_model_name: str | None = None
@ -799,11 +866,9 @@ async def patch_model(
stored_model_name = update_data.get("model_name")
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
update_data["updated_at"] = cast(str, get_utc_datetime())
async with _heuristic_v2_slot(
async with _auto_router_capability_slot(
prisma_client,
effective_config=_effective_complexity_router_config(
patch_data.litellm_params, db_model.litellm_params
),
effective_params=effective_params,
model_id=model_id,
) as table:
return await table.update(where={"model_id": model_id}, data=update_data)
@ -1959,9 +2024,12 @@ async def add_new_model(
model_params=priced_model_params,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
slot=_heuristic_v2_slot(
slot=_auto_router_capability_slot(
prisma_client,
effective_config=priced_model_params.litellm_params.complexity_router_config,
effective_params=_effective_complexity_router_params(
priced_model_params.litellm_params,
None,
),
model_id=priced_model_params.model_info.id,
),
)
@ -2110,6 +2178,9 @@ async def update_model(
incoming_params=model_params.litellm_params,
existing_params=deployment.litellm_params,
)
effective_params: Final = _effective_complexity_router_params(
model_params.litellm_params, deployment.litellm_params
)
# update DB
if store_model_in_db is True:
@ -2147,11 +2218,9 @@ async def update_model(
"updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME,
**({} if renamed_to is None else {"model_name": renamed_to}),
}
async with _heuristic_v2_slot(
async with _auto_router_capability_slot(
prisma_client,
effective_config=_effective_complexity_router_config(
model_params.litellm_params, deployment.litellm_params
),
effective_params=effective_params,
model_id=_model_id,
) as table:
model_response: Final = await table.update(

132
litellm/proxy/proxy_server.py
View file

@ -118,10 +118,11 @@ from litellm.router_utils.add_retry_fallback_headers import (
get_hidden_params_dict,
)
from litellm.router_utils.auto_router_model_naming import (
GATED_AUTO_ROUTER_CAPABILITIES,
STRATEGY_ROUTER_PARAM_FIELDS,
capability_limit_violation,
carries_complexity_router_settings,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
count_capability_routers,
validate_complexity_router_config_placement,
)
from litellm.types.utils import (
@ -303,7 +304,7 @@ from litellm.proxy.auth.auth_utils import (
)
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.litellm_license import HEURISTIC_V2_LICENSE_REMEDY, LicenseCheck
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
from litellm.proxy.auth.model_checks import (
expand_wildcard_deployments_for_model_info,
get_all_fallbacks,
@ -4340,17 +4341,28 @@ def validate_deployment_complexity_router_placement(model: Mapping[str, object])
raise ValueError(f"model {model.get('model_name', '')!r}: {violation}")
def validate_heuristic_v2_router_limit(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None:
def validate_auto_router_capability_limits(model_list: Sequence[Mapping[str, object]], *, limit: int | None) -> None:
"""
Refuse to start when config.yaml defines more heuristic_v2 auto-routers than the license allows.
Refuse to start when config.yaml defines more auto-routers claiming a licensed capability than allowed.
Checked here rather than left to router registration for the same reason as the two
validators above: the proxy builds its router with `ignore_invalid_deployments=True`, so
the router's own refusal would turn the extra router into a silently missing model.
"""
violation: Final = heuristic_v2_limit_violation(held=count_heuristic_v2_routers(model_list), limit=limit)
if violation is not None:
raise ValueError(f"config.yaml model_list: {violation} {HEURISTIC_V2_LICENSE_REMEDY}")
violations: Final = tuple(
message
for capability in GATED_AUTO_ROUTER_CAPABILITIES
if (
message := capability_limit_violation(
capability=capability,
held=count_capability_routers(model_list, capability=capability),
limit=limit,
)
)
is not None
)
if violations:
raise ValueError(f"config.yaml model_list: {' '.join(violations)} {AUTO_ROUTER_LICENSE_REMEDY}")
def pin_complexity_router_model_id(model: dict) -> None: # mutable-ok: out-param, model_info is stamped in place
@ -5758,7 +5770,7 @@ class ProxyConfig:
model_list: Final = config.get("model_list", None)
if model_list:
router_params["model_list"] = model_list
validate_heuristic_v2_router_limit(model_list, limit=_license_check.heuristic_v2_router_limit())
validate_auto_router_capability_limits(model_list, limit=_license_check.auto_router_capability_limit())
print( # noqa: T201
"\033[32mLiteLLM: Proxy initialized with Config, Set models:\033[0m"
)
@ -5848,7 +5860,7 @@ class ProxyConfig:
),
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
fallback_access_check=router_fallback_access_check,
heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
if redis_usage_cache is not None and router.cache.redis_cache is None:
@ -6309,7 +6321,7 @@ class ProxyConfig:
search_tools=search_tools,
ignore_invalid_deployments=True,
fallback_access_check=router_fallback_access_check,
heuristic_v2_router_limit=_license_check.heuristic_v2_router_limit,
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)
else:
@ -11458,6 +11470,37 @@ def _realtime_query_params_template(model: str | None, intent: str | None) -> tu
return tuple(params)
async def _release_realtime_budget_reservation(user_api_key_dict: UserAPIKeyAuth) -> None:
from litellm.proxy.spend_tracking.budget_reservation import (
release_or_invalidate_budget_reservation,
)
await release_or_invalidate_budget_reservation(
budget_reservation=user_api_key_dict.budget_reservation,
)
async def _reject_realtime_session(
websocket: WebSocket,
user_api_key_dict: UserAPIKeyAuth,
*,
code: int,
reason: str,
error_message: str | None = None,
) -> None:
try:
if error_message is not None:
try:
await websocket.send_text(
json.dumps({"type": "error", "error": {"type": "guardrail_error", "message": error_message}})
)
except Exception: # noqa: BLE001 # best-effort notice: a dead client socket must not skip the close below
verbose_proxy_logger.debug("Could not send realtime pre-call error event to client; closing anyway")
await websocket.close(code=code, reason=reason)
finally:
await _release_realtime_budget_reservation(user_api_key_dict)
@app.websocket("/openai/v1/realtime")
@app.websocket("/v1/realtime")
@app.websocket("/realtime")
@ -11483,7 +11526,9 @@ async def realtime_websocket_endpoint(
if intent == "transcription":
route_model = "gpt-realtime-whisper"
else:
await websocket.close(code=1008, reason="model query parameter is required")
await _reject_realtime_session(
websocket, user_api_key_dict, code=1008, reason="model query parameter is required"
)
return
assert route_model is not None
try:
@ -11494,7 +11539,7 @@ async def realtime_websocket_endpoint(
llm_router=llm_router,
)
except ProxyException as e:
await websocket.close(code=1008, reason=e.message[:120])
await _reject_realtime_session(websocket, user_api_key_dict, code=1008, reason=e.message[:120])
return
await websocket.accept(**accept_kwargs)
@ -11553,21 +11598,9 @@ async def realtime_websocket_endpoint(
)
except Exception as e:
verbose_proxy_logger.exception("Realtime pre-call error")
try:
await websocket.send_text(
json.dumps(
{
"type": "error",
"error": {
"type": "guardrail_error",
"message": str(e),
},
}
)
)
except Exception:
pass
await websocket.close(code=1011, reason="Pre-call error")
await _reject_realtime_session(
websocket, user_api_key_dict, code=1011, reason="Pre-call error", error_message=str(e)
)
return
# Phase 2: route to upstream LLM.
@ -11597,6 +11630,13 @@ async def realtime_websocket_endpoint(
)
except Exception: # noqa: BLE001 # the lower layer may have closed the socket already; closing twice is not an error
verbose_proxy_logger.debug("Could not close realtime client websocket; it is already gone")
finally:
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
)
if not litellm_logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY):
await _release_realtime_budget_reservation(user_api_key_dict)
######################################################################
@ -13501,6 +13541,27 @@ def _is_auto_router_model(model: Mapping[str, object]) -> bool:
return isinstance(litellm_model, str) and litellm_model.startswith("auto_router/")
def _model_in_access_group(model: Mapping[str, object], access_group: str) -> bool:
model_info: Final = model.get("model_info")
if not isinstance(model_info, Mapping):
return False
access_groups: Final = model_info.get("access_groups")
return isinstance(access_groups, (list, tuple)) and access_group in access_groups
def _matches_model_info_filters(
model: Mapping[str, object],
exclude_auto_routers: bool | None,
access_group: str | None,
wildcard_only: bool | None,
) -> bool:
if exclude_auto_routers is True and _is_auto_router_model(model):
return False
if isinstance(access_group, str) and not _model_in_access_group(model, access_group):
return False
return wildcard_only is not True or "*" in str(model.get("model_name") or "")
def _paginate_models_response(
all_models: list[dict[str, Any]],
page: int,
@ -13811,6 +13872,14 @@ async def model_info_v2(
"existing callers are unaffected"
),
),
access_group: str | None = fastapi.Query(
None,
description="Only return deployments whose `model_info.access_groups` contains this access group",
),
wildcard_only: bool | None = fastapi.Query(
False,
description="Only return wildcard deployments, i.e. those whose `model_name` contains `*`",
),
):
"""
Paginated model metadata for proxy deployments (pricing, provider, team access).
@ -13828,6 +13897,8 @@ async def model_info_v2(
modelId: Return a single deployment by LiteLLM model id.
teamId: Filter to models with direct access or team membership for this team id.
sortBy / sortOrder: Sort by model_name, created_at, updated_at, costs, or status.
access_group: Only return deployments in this model access group.
wildcard_only: Only return deployments whose `model_name` contains `*`.
Example request:
```
@ -13981,8 +14052,9 @@ async def model_info_v2(
# `is True` because direct-call tests bypass FastAPI, so the Query default arrives as a
# truthy sentinel object rather than False.
if exclude_auto_routers is True:
all_models = [m for m in all_models if not _is_auto_router_model(m)]
all_models = [
m for m in all_models if _matches_model_info_filters(m, exclude_auto_routers, access_group, wildcard_only)
]
# Update total count to include agents
search_total_count = len(all_models)

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

@ -373,6 +373,33 @@ async def invalidate_budget_reservation_counters(
await _invalidate_spend_counter(counter_key=counter_key)
async def release_or_invalidate_budget_reservation(
budget_reservation: dict | None, # mutable-ok: stamps finalized on the caller's shared reservation dict
) -> None:
"""Reconcile a still-open reservation on a terminal path that settles no cost.
A failed or upstream-refused request never runs the success cost callback, so
its pre-call reservation stays open and keeps the spend counter pinned above
real spend until the counter's TTL expires, 429ing later requests on the same
key. Release it to zero; if the release itself fails (e.g. the counter store is
unreachable) drop the reserved counters directly and mark the reservation
finalized so nothing reprocesses it. Idempotent: the finalized guard makes a
second call a no-op once success or failure handling already reconciled.
"""
if budget_reservation is None or budget_reservation.get("finalized") is True:
return
try:
await asyncio.shield(release_budget_reservation(budget_reservation=budget_reservation))
except Exception: # noqa: BLE001 # a cleanup failure must not pin the counter; drop it directly instead
verbose_proxy_logger.exception("Failed to release budget reservation; invalidating counters")
try:
await invalidate_budget_reservation_counters(budget_reservation=budget_reservation)
except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed
verbose_proxy_logger.exception("Failed to invalidate budget reservation counters after release failed")
finally:
budget_reservation["finalized"] = True
async def _get_budget_counters(
request_body: dict,
valid_token: UserAPIKeyAuth,

17
litellm/proxy/utils.py
View file

@ -6386,10 +6386,17 @@ class ProxyUpdateSpend:
)
break
except Exception as e:
if not PrismaDBExceptionHandler.is_database_transport_error(e):
if not _is_transient_spend_log_write_error(e):
if PrismaDBExceptionHandler.is_prisma_error(e):
await enqueue_spend_logs(prisma_client, logs_to_process, at_head=True)
verbose_proxy_logger.warning(
"Spend tracking - DB error writing spend logs, requeued %d rows for the next flush. error=%s",
len(logs_to_process),
str(e),
)
raise
verbose_proxy_logger.warning(
"Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
"Spend tracking - transient DB error writing spend logs, retry %d/%d. logs_count=%d, error=%s",
i + 1,
n_retry_times,
len(logs_to_process),
@ -6732,6 +6739,10 @@ async def _monitor_spend_logs_queue(
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
def _is_transient_spend_log_write_error(e: Exception) -> bool:
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
async def _create_spend_logs_with_poison_isolation(
repo: SpendLogsRepository,
rows: Sequence[Mapping[str, object]],
@ -6767,6 +6778,8 @@ async def _create_spend_logs_with_poison_isolation(
raise
if PrismaDBExceptionHandler.is_database_service_unavailable_error(e):
raise
if PrismaDBExceptionHandler.is_deadlock_error(e):
raise
budget_left: Final = max(failure_budget - 1, 0)
if len(rows) == 1:
request_id: Final = rows[0].get("request_id")

4
litellm/realtime_api/main.py
View file

@ -27,7 +27,7 @@ from litellm.types.realtime import (
RealtimeTranscriptionSessionRequest,
)
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
from litellm.types.utils import CallTypes, LlmProviders
from litellm.utils import ProviderConfigManager
from ..litellm_core_utils.get_litellm_params import get_litellm_params
@ -360,7 +360,7 @@ async def _arealtime(
user: Final = kwargs.get("user", None)
litellm_params: Final = GenericLiteLLMParams(**kwargs)
litellm_params_dict: Final = get_litellm_params(**kwargs)
litellm_params_dict: Final = {**get_litellm_params(**kwargs), CallTypes.arealtime.value: True}
model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = get_llm_provider(
model=model,

93
litellm/responses/main.py
View file

@ -2,6 +2,7 @@ import asyncio
import contextvars
from collections.abc import Coroutine, Generator, Iterable, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast
@ -23,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig
from litellm.responses.litellm_completion_transformation.handler import (
LiteLLMCompletionTransformationHandler,
)
@ -403,8 +405,40 @@ def _bridges_to_chat_completions(
return responses_api_provider_config is None or use_chat_completions_api is True
def _deployment_passes_through_responses(model_info: object) -> bool:
"""Whether ``model_info.supported_endpoints`` opts the deployment into native ``{api_base}/responses``."""
if not isinstance(model_info, dict):
return False
supported_endpoints: Final = model_info.get("supported_endpoints")
return isinstance(supported_endpoints, (list, tuple)) and "/v1/responses" in supported_endpoints
def _deployment_model_info_after_prompt_swap(
requested_provider: str | None, resolved_provider: str | None, model_info: object
) -> object:
"""Deployment metadata only describes the upstream while the prompt manager keeps its provider."""
return model_info if resolved_provider == requested_provider else None
@dataclass(frozen=True, slots=True)
class _AsyncPromptManagementOutcome:
merged_optional_params: Mapping[str, object]
deployment_model_info: object
def _resolve_responses_api_provider_config(
model: str, custom_llm_provider: str, model_info: object
) -> BaseResponsesAPIConfig | None:
provider_config: Final = ProviderConfigManager.get_provider_responses_api_config(
model=model, provider=custom_llm_provider
)
if provider_config is not None or not _deployment_passes_through_responses(model_info):
return provider_config
return OpenAILikeResponsesConfig()
def _will_bridge_to_chat_completions(
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool
model: str, custom_llm_provider: str | None, use_chat_completions_api: bool, model_info: object
) -> bool:
"""``_bridges_to_chat_completions`` for callers running before the provider config is resolved.
@ -418,9 +452,7 @@ def _will_bridge_to_chat_completions(
if custom_llm_provider is None:
return True
return _bridges_to_chat_completions(
ProviderConfigManager.get_provider_responses_api_config(
model=normalized_model[0], provider=custom_llm_provider
),
_resolve_responses_api_provider_config(normalized_model[0], custom_llm_provider, model_info),
use_chat_completions_api or normalized_model[1],
)
@ -527,7 +559,10 @@ async def aresponses(
with _prompt_management_sees_a_provisional_message_list(
kwargs,
bridged=_will_bridge_to_chat_completions(
model, custom_llm_provider, bool(kwargs.get("use_chat_completions_api"))
model,
custom_llm_provider,
bool(kwargs.get("use_chat_completions_api")),
kwargs.get("model_info"),
),
):
(
@ -552,6 +587,7 @@ async def aresponses(
merged_input=merged_input,
),
)
requested_provider: Final = custom_llm_provider
if model != original_model:
custom_llm_provider = _resolve_prompt_swapped_provider(
original_model=original_model,
@ -561,7 +597,12 @@ async def aresponses(
prompt_id=prompt_id,
)
kwargs.pop("prompt_id", None)
kwargs["_async_prompt_merged_params"] = merged_optional_params
kwargs["_async_prompt_merged_params"] = _AsyncPromptManagementOutcome(
merged_optional_params=merged_optional_params,
deployment_model_info=_deployment_model_info_after_prompt_swap(
requested_provider, custom_llm_provider, kwargs.get("model_info")
),
)
func: Final = partial(
responses,
@ -666,12 +707,14 @@ def _apply_prompt_management_to_responses_call(
kwargs: dict[str, Any],
local_vars: dict[str, object],
use_chat_completions_api: bool,
) -> tuple[str | ResponseInputParam, str, str | None]:
async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None)
if async_merged is not None:
for key, value in async_merged.items():
) -> tuple[str | ResponseInputParam, str, str | None, object]:
"""Returns the prompt-managed input, model and provider, plus the deployment metadata that still
describes the upstream (``None`` once the prompt manager moved the request to another provider)."""
async_outcome: Final[_AsyncPromptManagementOutcome | None] = kwargs.pop("_async_prompt_merged_params", None)
if async_outcome is not None:
for key, value in async_outcome.merged_optional_params.items():
local_vars[key] = value
return input, model, custom_llm_provider
return input, model, custom_llm_provider, async_outcome.deployment_model_info
prompt_id: Final = cast(str | None, kwargs.get("prompt_id", None))
prompt_variables: Final = cast(dict | None, kwargs.get("prompt_variables", None))
@ -684,7 +727,9 @@ def _apply_prompt_management_to_responses_call(
):
with _prompt_management_sees_a_provisional_message_list(
kwargs,
bridged=_will_bridge_to_chat_completions(model, custom_llm_provider, use_chat_completions_api),
bridged=_will_bridge_to_chat_completions(
model, custom_llm_provider, use_chat_completions_api, kwargs.get("model_info")
),
):
(
model,
@ -710,19 +755,28 @@ def _apply_prompt_management_to_responses_call(
)
local_vars["input"] = input
local_vars["model"] = model
if model != original_model:
custom_llm_provider = _resolve_prompt_swapped_provider(
resolved_provider: Final = (
custom_llm_provider
if model == original_model
else _resolve_prompt_swapped_provider(
original_model=original_model,
swapped_model=model,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
prompt_id=prompt_id,
)
local_vars["custom_llm_provider"] = custom_llm_provider
)
local_vars["custom_llm_provider"] = resolved_provider
for key, value in merged_optional_params.items():
local_vars[key] = value
return (
input,
model,
resolved_provider,
_deployment_model_info_after_prompt_swap(custom_llm_provider, resolved_provider, kwargs.get("model_info")),
)
return input, model, custom_llm_provider
return input, model, custom_llm_provider, kwargs.get("model_info")
# Opt-in via model id (mirrors the `responses/` prefix pattern on chat completions).
@ -1052,7 +1106,7 @@ def responses(
)
local_vars["custom_llm_provider"] = custom_llm_provider
input, model, custom_llm_provider = _apply_prompt_management_to_responses_call(
input, model, custom_llm_provider, deployment_model_info = _apply_prompt_management_to_responses_call(
input=input,
model=model,
custom_llm_provider=custom_llm_provider,
@ -1123,9 +1177,8 @@ def responses(
if custom_llm_provider is None:
responses_api_provider_config = None
else:
responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config(
model=model,
provider=custom_llm_provider,
responses_api_provider_config = _resolve_responses_api_provider_config(
model, custom_llm_provider, deployment_model_info
)
local_vars.update(kwargs)

206
litellm/router.py
View file

@ -116,10 +116,11 @@ from litellm.router_utils.add_retry_fallback_headers import (
)
from litellm.router_utils.auto_router_model_naming import (
AUTO_ROUTER_MODEL_PREFIX,
GatedAutoRouterCapability,
capability_limit_violation,
claimed_capability,
classify_strategy_router_model,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
uses_heuristic_v2_classifier,
count_capability_routers,
)
from litellm.router_utils.batch_utils import (
_get_router_metadata_variable_name,
@ -208,6 +209,7 @@ from litellm.types.router import (
AlertingConfig,
AllowedFailsPolicy,
AssistantsTypedDict,
AutoRouterCapabilityLimit,
ConsumedRequestTagsStamp,
CredentialLiteLLMParams,
CustomRoutingStrategyBase,
@ -215,7 +217,6 @@ from litellm.types.router import (
DeploymentTypedDict,
FallbackAccessCheck,
GuardrailTypedDict,
HeuristicV2RouterLimit,
LiteLLM_Params,
MockRouterTestingParams,
ModelGroupInfo,
@ -692,7 +693,7 @@ class Router:
background_health_check_model_groups: Sequence[str] | None = None,
enable_weighted_failover: bool = False,
fallback_access_check: FallbackAccessCheck | None = None,
heuristic_v2_router_limit: HeuristicV2RouterLimit | None = None,
auto_router_capability_limit: AutoRouterCapabilityLimit | None = None,
) -> None:
"""
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
@ -769,7 +770,7 @@ class Router:
self.set_verbose = set_verbose
self.ignore_invalid_deployments = ignore_invalid_deployments
self.heuristic_v2_router_limit = heuristic_v2_router_limit
self.auto_router_capability_limit = auto_router_capability_limit
self.fallback_access_check: Final = fallback_access_check
self.debug_level = debug_level
self.enable_pre_call_checks = enable_pre_call_checks
@ -2596,14 +2597,20 @@ class Router:
model_response: CustomStreamWrapper,
messages: list[dict[str, str]],
initial_kwargs: dict,
deployment_slot: contextlib.AsyncExitStack | None = None,
) -> CustomStreamWrapper:
"""
Helper to iterate over a streaming response.
Catches errors for fallbacks using the router's fallback system
`deployment_slot` holds the deployment's max_parallel_requests semaphore; it is
released when the stream is exhausted, closed, or falls back to another deployment
"""
from litellm.exceptions import MidStreamFallbackError
held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack()
class FallbackStreamWrapper(CustomStreamWrapper):
def __init__(self, async_generator: AsyncGenerator):
# Copy attributes from the original model_response
@ -2627,12 +2634,26 @@ class Router:
async def __anext__(self):
return await self._async_generator.__anext__()
async def close_model_response() -> None:
if not hasattr(model_response, "aclose"):
return
try:
await model_response.aclose()
except BaseException as e:
verbose_router_logger.debug(
"stream_with_fallbacks: error closing model_response: %s",
e,
)
async def stream_with_fallbacks():
fallback_response = None # Track for cleanup in finally
try:
async for item in model_response:
yield item
except MidStreamFallbackError as e:
with anyio.CancelScope(shield=True):
await close_model_response()
await held_slot.aclose()
if not e.is_pre_first_chunk and (
e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)
):
@ -2706,14 +2727,8 @@ class Router:
# (e.g. on client disconnect).
# Shield from anyio cancellation so the awaits can complete.
with anyio.CancelScope(shield=True):
if hasattr(model_response, "aclose"):
try:
await model_response.aclose()
except BaseException as e:
verbose_router_logger.debug(
"stream_with_fallbacks: error closing model_response: %s",
e,
)
await close_model_response()
await held_slot.aclose()
if fallback_response is not None and hasattr(fallback_response, "aclose"):
try:
await fallback_response.aclose()
@ -3378,61 +3393,53 @@ class Router:
kwargs=kwargs,
client_type="max_parallel_requests",
)
if rpm_semaphore is not None and isinstance(rpm_semaphore, asyncio.Semaphore):
async with rpm_semaphore:
"""
- Check rpm limits before making the call
- If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe)
"""
await self.async_routing_strategy_pre_call_checks(
deployment=deployment,
logging_obj=logging_obj,
parent_otel_span=parent_otel_span,
)
response = await _response
else:
async with contextlib.AsyncExitStack() as deployment_slot:
if isinstance(rpm_semaphore, asyncio.Semaphore):
await deployment_slot.enter_async_context(rpm_semaphore)
await self.async_routing_strategy_pre_call_checks(
deployment=deployment,
logging_obj=logging_obj,
parent_otel_span=parent_otel_span,
)
response = await _response
## CHECK CONTENT FILTER ERROR ##
if isinstance(response, ModelResponse):
_should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs)
if _should_raise:
raise litellm.ContentPolicyViolationError(
message="Response output was blocked.",
model=model,
llm_provider="",
## CHECK CONTENT FILTER ERROR ##
if isinstance(response, ModelResponse):
_should_raise = self._should_raise_content_policy_error(
model=model, response=response, kwargs=kwargs
)
if _should_raise:
raise litellm.ContentPolicyViolationError(
message="Response output was blocked.",
model=model,
llm_provider="",
)
if (
isinstance(response, CustomStreamWrapper)
and response.completion_stream is None
and response.make_call is not None
):
await response.fetch_stream()
if (
isinstance(response, CustomStreamWrapper)
and response.completion_stream is None
and response.make_call is not None
):
await response.fetch_stream()
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
# debug how often this deployment picked
self._track_deployment_metrics(
deployment=deployment,
response=response,
parent_otel_span=parent_otel_span,
)
if isinstance(response, CustomStreamWrapper):
return await self._acompletion_streaming_iterator(
model_response=response,
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
self.success_calls[model_name] += 1
verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
# debug how often this deployment picked
self._track_deployment_metrics(
deployment=deployment,
response=response,
parent_otel_span=parent_otel_span,
)
return response
if isinstance(response, CustomStreamWrapper):
return await self._acompletion_streaming_iterator(
model_response=response,
messages=messages,
initial_kwargs=input_kwargs_for_streaming_fallback,
deployment_slot=deployment_slot.pop_all(),
)
return response
except litellm.Timeout as e:
deployment_request_timeout_param: Final = _timeout_debug_deployment_dict.get("litellm_params", {}).get(
"request_timeout", None
@ -8373,17 +8380,12 @@ class Router:
## LOG FAILURE EVENT
if logging_obj is not None:
asyncio.create_task(
logging_obj.async_failure_handler(
logging_obj.dispatch_failure_handlers(
exception=e,
traceback_exception=traceback.format_exc(),
end_time=time.time(),
prefer_async_handlers=True,
)
)
## LOGGING
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback.format_exc()),
).start() # log response
_set_cooldown_deployments(
litellm_router_instance=self,
exception_status=e.status_code,
@ -8396,17 +8398,12 @@ class Router:
## LOG FAILURE EVENT
if logging_obj is not None:
asyncio.create_task(
logging_obj.async_failure_handler(
logging_obj.dispatch_failure_handlers(
exception=e,
traceback_exception=traceback.format_exc(),
end_time=time.time(),
prefer_async_handlers=True,
)
)
## LOGGING
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback.format_exc()),
).start() # log response
raise e
async def async_callback_filter_deployments(
@ -8444,17 +8441,12 @@ class Router:
## LOG FAILURE EVENT
if logging_obj is not None:
asyncio.create_task(
logging_obj.async_failure_handler(
logging_obj.dispatch_failure_handlers(
exception=e,
traceback_exception=traceback.format_exc(),
end_time=time.time(),
prefer_async_handlers=True,
)
)
## LOGGING
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback.format_exc()),
).start() # log response
raise e
return returned_healthy_deployments
@ -8811,20 +8803,21 @@ class Router:
if not (isinstance(model_info, Mapping) and model_info.get("db_model")):
yield deployment
def heuristic_v2_router_limit_violation(self) -> str | None:
def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None:
"""
Why one more heuristic_v2 router cannot join this router, or None when it can.
Why one more router claiming ``capability`` cannot join this router, or None when it can.
Judged against every deployment currently on the model_list; an upsert pops the row being
edited first, so an edit of an existing heuristic_v2 router keeps its own slot. The limit is
resolved on every call through ``heuristic_v2_router_limit``; unset means unlimited, which
is the SDK default, and the proxy injects a resolver backed by its license.
edited first, so an edit of an existing gated router keeps its own slot. The limit is
resolved on every call through ``auto_router_capability_limit``; unset means unlimited,
which is the SDK default, and the proxy injects a resolver backed by its license.
"""
limit: Final = self.heuristic_v2_router_limit() if self.heuristic_v2_router_limit is not None else None
others: Final = count_heuristic_v2_routers(
deployment for deployment in self.model_list if isinstance(deployment, Mapping)
limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None
others: Final = count_capability_routers(
(deployment for deployment in self.model_list if isinstance(deployment, Mapping)),
capability=capability,
)
return heuristic_v2_limit_violation(held=others + 1, limit=limit)
return capability_limit_violation(capability=capability, held=others + 1, limit=limit)
def init_complexity_router_deployment(self, deployment: Deployment):
"""
@ -8843,8 +8836,9 @@ class Router:
)
complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config
if uses_heuristic_v2_classifier(complexity_router_config):
limit_violation: Final = self.heuristic_v2_router_limit_violation()
capability: Final = claimed_capability(complexity_router_config)
if capability is not None:
limit_violation: Final = self.auto_router_capability_violation(capability)
if limit_violation is not None:
raise ValueError(limit_violation)
@ -9674,13 +9668,13 @@ class Router:
"""Put a deployment back the way it was before a failed upsert popped it.
A rollback re-admits state that was already serving, so it does not go through the
heuristic_v2 ceiling a newcomer gets: with the ceiling tightened since the deployment first
capability ceiling a newcomer gets: with the ceiling tightened since the deployment first
registered, judging the rollback would drop a serving router over an unrelated failed edit.
"""
if previous_deployment is None or self.has_model_id(model_id):
return
limit_resolver: Final = self.heuristic_v2_router_limit
self.heuristic_v2_router_limit = None
limit_resolver: Final = self.auto_router_capability_limit
self.auto_router_capability_limit = None
try:
self.add_deployment(deployment=previous_deployment)
verbose_router_logger.info(
@ -9696,7 +9690,7 @@ class Router:
restore_error,
)
finally:
self.heuristic_v2_router_limit = limit_resolver
self.auto_router_capability_limit = limit_resolver
@staticmethod
def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:
@ -12634,13 +12628,13 @@ class Router:
logging_obj: Final = request_kwargs.get("litellm_logging_obj", None)
if logging_obj is not None:
## LOGGING
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback_exception),
).start() # log response
# Handle any exceptions that might occur during streaming
asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception))
asyncio.create_task(
logging_obj.dispatch_failure_handlers(
exception=e,
traceback_exception=traceback_exception,
prefer_async_handlers=True,
)
)
raise e
async def async_get_available_deployment_for_pass_through(
@ -12768,11 +12762,13 @@ class Router:
if request_kwargs is not None:
logging_obj: Final = request_kwargs.get("litellm_logging_obj", None)
if logging_obj is not None:
threading.Thread(
target=logging_obj.failure_handler,
args=(e, traceback_exception),
).start()
asyncio.create_task(logging_obj.async_failure_handler(e, traceback_exception))
asyncio.create_task(
logging_obj.dispatch_failure_handlers(
exception=e,
traceback_exception=traceback_exception,
prefer_async_handlers=True,
)
)
raise e
async def _run_routing_plugins(

134
litellm/router_utils/auto_router_model_naming.py
View file

@ -10,7 +10,7 @@ the router silently dropping the deployment at load time under
``ignore_invalid_deployments``.
"""
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, TypeAlias
@ -81,6 +81,11 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None:
return "semantic"
def is_complexity_router_model(model: str | None) -> bool:
"""Whether ``model`` selects the complexity-router implementation."""
return classify_strategy_router_model(model or "") == "complexity"
def _named(value: object, role: StrategyRouterDependencyRole) -> tuple[StrategyRouterDependency, ...]:
"""One dependency from a scalar field, or none when it is absent or not a name."""
return (StrategyRouterDependency(value, role),) if isinstance(value, str) and value else ()
@ -168,20 +173,121 @@ def uses_heuristic_v2_classifier(complexity_router_config: object) -> bool:
return _mapping(complexity_router_config).get("classifier_type") == "heuristic_v2"
def is_heuristic_v2_router(litellm_params: Mapping[str, object]) -> bool:
"""Whether this deployment is a complexity router that classifies with heuristic_v2."""
return classify_strategy_router_model(str(litellm_params.get("model") or "")) == "complexity" and (
uses_heuristic_v2_classifier(litellm_params.get("complexity_router_config"))
def defines_custom_tiers(complexity_router_config: object) -> bool:
"""Whether this complexity config replaces the built-in tier ladder with operator-defined tier_definitions.
Mirrors the SQL spelling on the capability record: only an actual array claims the capability,
so an explicit JSON null or a malformed value does not.
"""
return isinstance(_mapping(complexity_router_config).get("tier_definitions"), (list, tuple))
OPERATOR_CLASSIFIER_PROMPT_FIELDS: Final = ("classification_prompt", "classification_examples")
def defines_custom_classifier_prompt(complexity_router_config: object) -> bool:
"""Whether an operator wrote any part of this router's classifier prompt themselves.
Three spellings, all metered: a whole replacement prompt (``classifier_llm_config.system_prompt``),
replacement opening instructions (``classification_prompt``), and replacement calibration examples
(``classification_examples``). Choosing a shipped ``classification_rubric`` preset is not authoring.
Scoped to the classifier types that actually call an LLM, which is also where the config validator
accepts these fields: the heuristic scorers never read them.
"""
config: Final = _mapping(complexity_router_config)
if config.get("classifier_type") not in LLM_CLASSIFIER_TYPES:
return False
return _mapping(config.get("classifier_llm_config")).get("system_prompt") is not None or any(
config.get(field) is not None for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
)
def count_heuristic_v2_routers(deployments: Iterable[Mapping[str, object]]) -> int:
"""How many of ``deployments`` (router model_list entries or config.yaml rows) are heuristic_v2 routers."""
return sum(1 for deployment in deployments if is_heuristic_v2_router(_mapping(deployment.get("litellm_params"))))
def uses_custom_tier_or_classifier_prompt(complexity_router_config: object) -> bool:
"""Whether this router replaces shipped tiers or its shipped classifier prompt."""
return defines_custom_tiers(complexity_router_config) or defines_custom_classifier_prompt(complexity_router_config)
def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None:
"""Why holding ``held`` heuristic_v2 routers exceeds ``limit``, or None when it fits.
_LLM_CLASSIFIER_TYPES_SQL: Final = ", ".join(f"'{name}'" for name in sorted(LLM_CLASSIFIER_TYPES))
@dataclass(frozen=True, slots=True)
class GatedAutoRouterCapability:
"""A complexity-router capability the license meters, in every spelling an enforcement point needs.
``uses`` and ``sql_config_predicate`` answer the same question, in process and in a DB count over
stored ``litellm_params`` (``{config}`` is the caller's expression for the normalized
``complexity_router_config`` jsonb, substituted as many times as the predicate needs); they live
on one record so they cannot drift apart. ``subject`` and ``remedy`` build the shared refusal
message. A validated config claims at most one capability, and the validator is what makes that
true: tier_definitions rejects every heuristic classifier_type, and it also rejects the
classifier system_prompt, which in turn only applies to the classifier types heuristic_v2 is not.
"""
key: str
subject: str
remedy: str
uses: Callable[[object], bool]
sql_config_predicate: str
HEURISTIC_V2_CAPABILITY: Final = GatedAutoRouterCapability(
key="heuristic_v2",
subject="with classifier_type 'heuristic_v2'",
remedy="Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router.",
uses=uses_heuristic_v2_classifier,
sql_config_predicate="{config} ->> 'classifier_type' = 'heuristic_v2'",
)
_OPERATOR_PROMPT_FIELDS_SQL: Final = " OR ".join(
f"{{config}} ->> '{field}' IS NOT NULL" for field in OPERATOR_CLASSIFIER_PROMPT_FIELDS
)
CUSTOMIZATION_CAPABILITY: Final = GatedAutoRouterCapability(
key="tier_or_classifier_prompt",
subject="with operator-defined tier_definitions or an operator-written classifier prompt",
remedy=(
"Use the shipped tiers and classifier prompt for this router or remove an existing router "
"with tier_definitions or its own classifier prompt."
),
uses=uses_custom_tier_or_classifier_prompt,
sql_config_predicate=(
"jsonb_typeof({config} -> 'tier_definitions') = 'array' OR "
f"({{config}} ->> 'classifier_type' IN ({_LLM_CLASSIFIER_TYPES_SQL}) AND ("
"{config} -> 'classifier_llm_config' ->> 'system_prompt' IS NOT NULL OR "
f"{_OPERATOR_PROMPT_FIELDS_SQL}))"
),
)
GATED_AUTO_ROUTER_CAPABILITIES: Final = (HEURISTIC_V2_CAPABILITY, CUSTOMIZATION_CAPABILITY)
def claimed_capability(complexity_router_config: object) -> GatedAutoRouterCapability | None:
"""The licensed capability this complexity config claims, or None."""
return next(
(capability for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(complexity_router_config)),
None,
)
def gated_capability_of(litellm_params: Mapping[str, object]) -> GatedAutoRouterCapability | None:
"""The licensed capability this deployment claims, or None unless it is a complexity router."""
model: Final = litellm_params.get("model")
if not is_complexity_router_model(model if isinstance(model, str) else None):
return None
return claimed_capability(litellm_params.get("complexity_router_config"))
def count_capability_routers(
deployments: Iterable[Mapping[str, object]], *, capability: GatedAutoRouterCapability
) -> int:
"""How many of ``deployments`` (router model_list entries or config.yaml rows) claim ``capability``."""
return sum(
1 for deployment in deployments if gated_capability_of(_mapping(deployment.get("litellm_params"))) is capability
)
def capability_limit_violation(*, capability: GatedAutoRouterCapability, held: int, limit: int | None) -> str | None:
"""Why holding ``held`` routers claiming ``capability`` exceeds ``limit``, or None when it fits.
``limit`` None means unlimited. The message is shared by every enforcement point (config
load, model writes, router registration) and stays SDK-neutral: it names the cap and what
@ -190,8 +296,8 @@ def heuristic_v2_limit_violation(*, held: int, limit: int | None) -> str | None:
if limit is None or held <= limit:
return None
return (
f"At most {limit} auto-router(s) with classifier_type 'heuristic_v2' can be registered but this would make "
f"{held}. Use classifier_type 'heuristic' for this router or remove an existing heuristic_v2 router."
f"At most {limit} auto-router(s) {capability.subject} can be registered but this would make "
f"{held}. {capability.remedy}"
)
@ -237,9 +343,7 @@ def carries_complexity_router_settings(model: str | None, present_fields: frozen
``validate_strategy_router_model_write`` is judged on, so a router named only by its
default model is in scope, and a field added to the table above is covered here for free.
"""
return classify_strategy_router_model(model or "") == "complexity" or bool(
present_fields & _COMPLEXITY_ROUTER_FIELDS
)
return is_complexity_router_model(model) or bool(present_fields & _COMPLEXITY_ROUTER_FIELDS)
def validate_complexity_router_config_placement(litellm_params: Mapping[str, object] | None) -> str | None:

17
litellm/router_utils/pattern_match_deployments.py
View file

@ -56,8 +56,9 @@ class PatternMatchRouter:
This class will store a mapping for regex pattern: List[Deployments]
"""
def __init__(self):
def __init__(self, pattern_utils: type[PatternUtils] = PatternUtils):
self.patterns: dict[str, list] = {}
self._pattern_utils: Final = pattern_utils
def add_pattern(self, pattern: str, llm_deployment: dict):
"""
@ -69,9 +70,10 @@ class PatternMatchRouter:
"""
# Convert the pattern to a regex
regex: Final = self._pattern_to_regex(pattern)
if regex not in self.patterns:
self.patterns[regex] = []
self.patterns[regex].append(llm_deployment)
if regex in self.patterns:
self.patterns[regex].append(llm_deployment)
return
self.patterns = dict(self._pattern_utils.sorted_patterns({**self.patterns, regex: [llm_deployment]}))
def remove_deployment(self, model_id: str) -> None:
"""
@ -138,11 +140,12 @@ class PatternMatchRouter:
if request is None:
return None
sorted_patterns: Final = PatternUtils.sorted_patterns(self.patterns)
regex_filtered_model_names: Final = (
[self._pattern_to_regex(m) for m in filtered_model_names] if filtered_model_names is not None else []
tuple(self._pattern_to_regex(m) for m in filtered_model_names)
if filtered_model_names is not None
else ()
)
for pattern, llm_deployments in sorted_patterns:
for pattern, llm_deployments in self.patterns.items():
if filtered_model_names is not None and pattern not in regex_filtered_model_names:
continue
pattern_match = re.match(pattern, request)

4
litellm/types/router.py
View file

@ -887,9 +887,9 @@ class FallbackAccessCheck(Protocol):
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: "Router") -> bool: ...
class HeuristicV2RouterLimit(Protocol):
class AutoRouterCapabilityLimit(Protocol):
"""
Resolves how many heuristic_v2 complexity routers the Router may hold right now; None means unlimited.
Resolves how many complexity routers may claim each licensed capability right now; None means unlimited.
The Router calls it on every registration and limit query instead of caching the answer, so the
proxy can keep the limit on its license object (re-verified on config load) rather than hand

4
pyproject.toml
View file

@ -67,8 +67,8 @@ proxy = [
"azure-identity>=1.25.2,<2.0",
"azure-storage-blob>=12.28.0,<13.0",
"mcp>=1.28.1,<2.0",
"litellm-proxy-extras==0.4.93",
"litellm-enterprise==0.1.64",
"litellm-proxy-extras==0.4.94",
"litellm-enterprise==0.1.65",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",

98
tests/e2e/llm_translation/test_chat_completions_regression_e2e.py
View file

@ -11,7 +11,7 @@ fails that provider's row here.
The per-provider classes below cover the OpenAI-compatible /chat/completions
translation for providers customers reach by registering their own deployment
via /model/new (Cohere, Gemini, hosted_vllm), each deleted on teardown.
via /model/new (Cohere, Gemini, hosted_vllm, Anthropic), each deleted on teardown.
"""
from __future__ import annotations
@ -46,6 +46,7 @@ pytestmark = pytest.mark.e2e
COHERE_BACKEND = "cohere/command-r-08-2024"
GEMINI_BACKEND = "gemini/gemini-2.5-flash"
OPENAI_BACKEND = "openai/gpt-5.6"
ANTHROPIC_BACKEND = "anthropic/claude-haiku-4-5-20251001"
BEDROCK_CONVERSE_BACKEND = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
@ -746,3 +747,98 @@ class TestBedrockConverseChatCompletions:
response = unwrap(client.proxy.chat(key, ChatBody(model=model, messages=_vision_messages(), max_tokens=32)))
_assert_describes_cat(response)
class TestAnthropicChatCompletions:
"""Anthropic via the OpenAI-compatible /chat/completions path, the translation
customers on the OpenAI SDK rely on when they route to Claude. The streamed call
must deliver real content deltas, and a tool-forced call must come back as a
well-formed tool_call on both the non-streamed and streamed paths.
"""
def _register(self, client: PassthroughClient, resources: ResourceManager, prefix: str) -> str:
model = f"{prefix}-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=ANTHROPIC_BACKEND, api_key="os.environ/ANTHROPIC_API_KEY")
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model
@pytest.mark.covers(
"llm.chat_completions.anthropic.basic.stream.works",
exercised_on=["chat_completions"],
)
def test_anthropic_chat_streams_real_content(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
model = self._register(client, resources, "e2e-anthropic-stream")
key = resources.key()
result = client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[
ChatMessage(role="user", content=f"Count from 1 to 5, one number per line. {unique_marker()}")
],
max_tokens=64,
stream=True,
),
)
_assert_streamed_completion(result)
@pytest.mark.covers(
"llm.chat_completions.anthropic.tool_use.nonstream.works",
exercised_on=["chat_completions"],
)
def test_anthropic_chat_returns_tool_call(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
model = self._register(client, resources, "e2e-anthropic-tool")
key = resources.key()
response = unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[
ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.")
],
tools=[_WEATHER_TOOL],
tool_choice="required",
max_tokens=128,
),
)
)
_assert_weather_tool_call(response)
@pytest.mark.covers(
"llm.chat_completions.anthropic.tool_use.stream.works",
exercised_on=["chat_completions"],
)
def test_anthropic_chat_streams_tool_call(
self, client: PassthroughClient, resources: ResourceManager
) -> None:
model = self._register(client, resources, "e2e-anthropic-tool-stream")
key = resources.key()
result = client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[
ChatMessage(role="user", content="What is the weather in San Francisco? Use the get_weather tool.")
],
tools=[_WEATHER_TOOL],
tool_choice="required",
max_tokens=128,
stream=True,
),
)
assert result.ok and result.is_streaming, f"tool stream was not established: {result}"
assert result.stream_error is None, f"tool stream carried an error event: {result.stream_error}"
name, arguments = _streamed_tool_call(result.stream_events)
assert name == "get_weather", f"streamed tool call named {name!r}: {result.stream_events[:5]}"
args = _WeatherArgs.model_validate_json(arguments)
assert args.location.strip(), f"streamed tool call arguments missing location: {arguments!r}"

9
tests/e2e/models.py
View file

@ -851,6 +851,15 @@ class ModelListEntry(BaseModel):
id: str
class ModelsListParams(BaseModel):
"""Query for GET /v1/models. A wildcard route such as ``openai/gpt-5.4*`` is
listed only under ``return_wildcard_routes``; without it the route is dropped
and only its expansions remain, so a readiness poll for the pattern itself
never resolves."""
return_wildcard_routes: bool = True
class ModelsListResponse(BaseModel):
"""GET /v1/models on the data plane: the deployments the gateway can actually
serve right now. Used to confirm a freshly created model has propagated from

3
tests/e2e/proxy_client.py
View file

@ -55,6 +55,7 @@ from models import (
ModelMode,
ModelNewBody,
ModelNewResponse,
ModelsListParams,
ModelsListResponse,
ModelUpdateBody,
OcrBody,
@ -336,7 +337,7 @@ class ProxyClient:
lambda poll_timeout: self.transport.get(
"/v1/models",
headers=headers,
params=NoBody(),
params=ModelsListParams(),
response_type=ModelsListResponse,
timeout=poll_timeout,
),

15
tests/e2e/router/test_auto_router_regressions_e2e.py
View file

@ -597,9 +597,20 @@ class TestSemanticAutoRouterResponses:
)
)
assert answer.id, "/v1/responses through the semantic auto-router returned no response id"
rows: Final = proxy.poll_logs_for_key(key, min_rows=1)
rows: Final = proxy.poll_logs_for_key(
key,
min_rows=2,
predicate=lambda logged: any(row.model == EMBEDDING_MODEL for row in logged),
)
embedding_rows: Final = tuple(row for row in rows if row.model == EMBEDDING_MODEL)
assert embedding_rows, (
"the routing embedding was not billed to the caller's key; "
f"spend logs show {tuple(row.model for row in rows)}"
)
_assert_served_only_by(
rows, CHEAP_SERVED | {semantic_auto_router.target}, "semantic auto-router /v1/responses string input"
[row for row in rows if row.model != EMBEDDING_MODEL],
CHEAP_SERVED | {semantic_auto_router.target},
"semantic auto-router /v1/responses string input",
)

101
tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py
View file

@ -10,6 +10,8 @@ Covers the three defects from the ticket:
handling live only on the native path).
"""
import time
import pytest
from litellm_enterprise.enterprise_callbacks.secret_detection import (
@ -19,12 +21,16 @@ from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
AWS_KEY = "AKIAIOSFODNN7EXAMPLE"
OPENAI_KEY = "sk-test-abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGH"
SHORT_OPENAI_KEY = "sk-12345"
UNICODE_DIGIT_SUFFIX = "sk-notification٣"
STRIPE_LIVE_KEY = f"sk_live_{'1234567890' * 3}"
URL_ENCODED_KEY = "Bearer%20sk-Ab3dEf6Gh7Ij8Kl9Mn0Pq2Rs3Tu4Vw5X"
AWS_KEYS = [f"AKIAIOSFODNN7EXAMPL{suffix}" for suffix in "FEDCBA"]
def _guardrail() -> _ENTERPRISE_SecretDetection:
return _ENTERPRISE_SecretDetection(
guardrail_name="hide-secrets", event_hook="pre_call", default_on=True
)
return _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets", event_hook="pre_call", default_on=True)
def _recorded(request_data: dict) -> dict:
@ -33,6 +39,91 @@ def _recorded(request_data: dict) -> dict:
return entries[0]
def test_scan_message_preserves_benign_identifiers_and_xml_tags():
guardrail = _guardrail()
content = "<task-notification> model: claude-sonnet-4-5-20250929 </task-notification>"
assert guardrail.scan_message_for_secrets(content) == []
assert guardrail.redact_text(content) == content
assert guardrail.redact_text("result = compute(x) </task-notification>") == (
"result = compute(x) </task-notification>"
)
def test_scan_message_preserves_quoted_benign_identifiers():
guardrail = _guardrail()
content = '{"content-type": "application/json", "model": "claude-sonnet-4-5-20250929"}'
assert guardrail.scan_message_for_secrets(content) == []
assert guardrail.redact_text(content) == content
def test_scan_message_redacts_every_openai_key_occurrence():
guardrail = _guardrail()
content = f"first {OPENAI_KEY}, second {OPENAI_KEY}"
assert guardrail.redact_text(content) == "first [REDACTED], second [REDACTED]"
def test_scan_message_redacts_short_numeric_openai_like_values():
guardrail = _guardrail()
assert guardrail.redact_text(f"value {SHORT_OPENAI_KEY}") == "value [REDACTED]"
def test_scan_message_requires_ascii_digits_for_openai_like_values():
guardrail = _guardrail()
assert guardrail.scan_message_for_secrets(UNICODE_DIGIT_SUFFIX) == []
assert guardrail.redact_text(UNICODE_DIGIT_SUFFIX) == UNICODE_DIGIT_SUFFIX
def test_scan_message_redacts_openai_key_after_separator():
guardrail = _guardrail()
assert guardrail.redact_text(f"openai_{OPENAI_KEY} key-{OPENAI_KEY}") == (
"openai_[REDACTED] key-[REDACTED]"
)
assert guardrail.redact_text(URL_ENCODED_KEY) == "Bearer%20[REDACTED]"
def test_scan_message_does_not_stop_openai_key_at_token_characters():
guardrail = _guardrail()
assert guardrail.redact_text("key sk-proj-abcde12345/extra") == "key [REDACTED]/extra"
def test_scan_message_stays_linear_on_repeated_sk_separators():
guardrail = _guardrail()
content = "-sk-" * 25_000
started = time.perf_counter()
assert guardrail.scan_message_for_secrets(content) == []
assert time.perf_counter() - started < 2.0
def test_scan_message_redacts_whole_stripe_live_key():
guardrail = _guardrail()
assert guardrail.redact_text(f"stripe {STRIPE_LIVE_KEY} end") == "stripe [REDACTED] end"
def test_scan_message_returns_matches_in_stable_order():
guardrail = _guardrail()
detected = guardrail.scan_message_for_secrets(" ".join(AWS_KEYS))
assert [secret["value"] for secret in detected] == sorted(AWS_KEYS)
def test_scan_message_replaces_longest_overlapping_match_first():
guardrail = _guardrail()
content = f'token = "{OPENAI_KEY}/extra"'
detected = guardrail.scan_message_for_secrets(content)
assert [secret["value"] for secret in detected] == [f"{OPENAI_KEY}/extra", OPENAI_KEY]
assert guardrail.redact_text(content) == 'token = "[REDACTED]"'
@pytest.mark.asyncio
async def test_apply_guardrail_redacts_secrets():
"""Playground path: the returned texts must carry [REDACTED], not the secret."""
@ -199,9 +290,7 @@ async def test_apply_guardrail_without_texts_records_nothing():
"messages": [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "https://x/y.png"}}
],
"content": [{"type": "image_url", "image_url": {"url": "https://x/y.png"}}],
}
],
"metadata": {},

24
tests/test_litellm/integrations/cloudzero/test_cz_stream_api.py
View file

@ -69,6 +69,30 @@ class TestCloudZeroStreamer:
assert "2025-01-19" in result
assert len(result["2025-01-19"]) == 1
def test_group_by_date_infers_schema_from_every_row(self):
"""Test daily batches retain optional string columns that are null for thousands of leading rows."""
streamer = CloudZeroStreamer("test-key", "test-connection")
leading_nulls = 10_000
rows = [
{"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": None}
for _ in range(leading_nulls)
]
rows.append(
{"time/usage_start": "2025-01-19T10:30:00Z", "resource/tag:team_alias": "team-alias"}
)
data = pl.DataFrame(
rows,
schema={"time/usage_start": pl.String, "resource/tag:team_alias": pl.String},
)
result = streamer._group_by_date(data)
batch = result["2025-01-19"]
assert len(batch) == leading_nulls + 1
assert batch.schema["resource/tag:team_alias"] == pl.String
assert batch["resource/tag:team_alias"].null_count() == leading_nulls
assert batch.tail(1).item(0, "resource/tag:team_alias") == "team-alias"
def test_parse_and_convert_timestamp_utc(self):
"""Test _parse_and_convert_timestamp method with UTC timestamp."""
streamer = CloudZeroStreamer("test-key", "test-connection")

27
tests/test_litellm/integrations/cloudzero/test_transform.py
View file

@ -86,6 +86,33 @@ class TestCBFTransformer:
assert result.is_empty()
def test_transform_keeps_tags_first_seen_after_row_100(self):
transformer = CBFTransformer()
teamless_rows = 101
team_rows = 2
total_rows = teamless_rows + team_rows
data = pl.DataFrame(
{
"date": ["2025-01-19"] * total_rows,
"successful_requests": [1] * total_rows,
"spend": [0.5] * total_rows,
"prompt_tokens": [10] * total_rows,
"completion_tokens": [5] * total_rows,
"model": ["gpt-4"] * total_rows,
"custom_llm_provider": ["openai"] * total_rows,
"api_key": ["sk-late-team"] * total_rows,
"team_id": pl.Series([None] * teamless_rows + ["team-late"] * team_rows, dtype=pl.String),
"team_alias": pl.Series([None] * teamless_rows + ["Late Team"] * team_rows, dtype=pl.String),
}
)
result = transformer.transform(data)
assert len(result) == total_rows
assert "resource/tag:team_alias" in result.columns
assert result["resource/tag:team_alias"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows
assert result["resource/tag:entity_id"].to_list() == [None] * teamless_rows + ["Late Team"] * team_rows
def test_create_cbf_record(self):
"""Test _create_cbf_record method with valid row data."""
transformer = CBFTransformer()

168
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -1,4 +1,5 @@
import asyncio
from typing import TYPE_CHECKING, Literal, Optional
from unittest.mock import AsyncMock
import pytest
@ -11,6 +12,9 @@ from litellm.integrations.custom_guardrail import (
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
class TestCustomGuardrailDeploymentHook:
@ -2239,6 +2243,170 @@ class TestRecordsOwnGuardrailInformation:
assert _guardrail_entries(request_data) == []
class _UndecoratedGuardrail(CustomGuardrail):
"""apply_guardrail written like the docs example: no @log_guardrail_information."""
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
from litellm.exceptions import GuardrailRaisedException
if any("forbidden" in text for text in inputs.get("texts") or []):
raise GuardrailRaisedException(guardrail_name=self.guardrail_name, message="Content blocked")
return inputs
class _UndecoratedSelfRecordingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"custom": True},
request_data=request_data,
guardrail_status="success",
start_time=0.0,
end_time=0.0,
duration=0.0,
)
return inputs
class _InheritedApplyGuardrail(_UndecoratedGuardrail):
pass
class TestUndecoratedApplyGuardrailIsLogged:
"""LIT-5983 regression: a custom guardrail that overrides apply_guardrail without the
@log_guardrail_information decorator must still record guardrail information, and the
auto-wrap must not double-record decorated or self-recording implementations."""
@pytest.mark.asyncio
async def test_undecorated_success_is_recorded(self):
from litellm.types.guardrails import GuardrailEventHooks
guardrail = _UndecoratedGuardrail(guardrail_name="docs-style", event_hook=GuardrailEventHooks.pre_call)
request_data: dict = {"model": "gpt-4o"}
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
request_data=request_data,
input_type="request",
)
entries = _guardrail_entries(request_data)
assert len(entries) == 1
assert entries[0]["guardrail_name"] == "docs-style"
assert entries[0]["guardrail_mode"] == "pre_call"
assert entries[0]["guardrail_status"] == "success"
@pytest.mark.asyncio
async def test_undecorated_block_is_recorded_and_reraised(self):
from litellm.exceptions import GuardrailRaisedException
guardrail = _UndecoratedGuardrail(guardrail_name="docs-style")
request_data: dict = {"model": "gpt-4o"}
with pytest.raises(GuardrailRaisedException):
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["forbidden"]),
request_data=request_data,
input_type="request",
)
entries = _guardrail_entries(request_data)
assert len(entries) == 1
assert entries[0]["guardrail_name"] == "docs-style"
assert entries[0]["guardrail_status"] == "guardrail_intervened"
@pytest.mark.asyncio
async def test_undecorated_bare_exception_is_recorded_as_failed_to_respond(self):
class _BareExceptionGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
raise Exception("Content blocked: policy violation")
guardrail = _BareExceptionGuardrail(guardrail_name="docs-style")
request_data: dict = {"model": "gpt-4o"}
with pytest.raises(Exception, match="Content blocked"):
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["x"]),
request_data=request_data,
input_type="request",
)
entries = _guardrail_entries(request_data)
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond"
@pytest.mark.asyncio
async def test_inherited_apply_guardrail_is_recorded_once(self):
guardrail = _InheritedApplyGuardrail(guardrail_name="child")
request_data: dict = {"model": "gpt-4o"}
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
request_data=request_data,
input_type="request",
)
assert len(_guardrail_entries(request_data)) == 1
@pytest.mark.asyncio
async def test_undecorated_self_recording_apply_guardrail_is_recorded_once(self):
guardrail = _UndecoratedSelfRecordingGuardrail(guardrail_name="self-recording")
request_data: dict = {"model": "gpt-4o"}
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
request_data=request_data,
input_type="request",
)
entries = _guardrail_entries(request_data)
assert len(entries) == 1
assert entries[0]["guardrail_response"] == {"custom": True}
@pytest.mark.asyncio
async def test_base_apply_guardrail_is_not_recorded(self):
guardrail = CustomGuardrail(guardrail_name="base")
request_data: dict = {"model": "gpt-4o"}
await guardrail.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=["hello"]),
request_data=request_data,
input_type="request",
)
assert _guardrail_entries(request_data) == []
def test_subclass_keywords_reach_cooperative_init_subclass(self):
class _LabelMixin:
seen_label: str = ""
def __init_subclass__(cls, label: str = "", **kwargs: object) -> None:
super().__init_subclass__(**kwargs)
cls.seen_label = label
class _Labelled(CustomGuardrail, _LabelMixin, label="docs-style"):
pass
assert _Labelled.seen_label == "docs-style"
class _ApplyOnlyObserver(CustomGuardrail):
"""Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook."""

144
tests/test_litellm/integrations/test_shadow_eval_logger.py
View file

@ -3,6 +3,7 @@ the detached pipeline's single attempt-row write, and the cache-first job lookup
import asyncio
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -145,6 +146,48 @@ def _shadow_reply_router(message, finish_reason="stop", routed_model="cheap-mode
return router
def _reasoning_judge_router(
reasoning_tokens: int, verdict: str = '{"preference": "A", "confidence": 0.9}'
) -> MagicMock:
"""A router whose judge arm reasons before it answers, the way a deployment carrying an
elevated reasoning_effort does: reasoning bills against the caller's own max_tokens and
the reply is cut off at that cap. One character stands in for one token."""
router = MagicMock()
router.model_group_alias = {}
router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}])
async def acompletion(**kwargs):
if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN:
kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"}
return {"choices": [{"message": {"content": "shadow answer"}}]}
budget_for_the_answer: Final = kwargs["max_tokens"] - reasoning_tokens
return {"choices": [{"message": {"content": verdict[: max(0, budget_for_the_answer)]}}]}
router.acompletion = MagicMock(side_effect=acompletion)
return router
def _judge_reply_router(content: str | None, finish_reason: str = "stop", served_model: str = "judge-pick") -> MagicMock:
"""A router whose judge arm returns a caller-shaped reply, so the shapes that all land
on the same parser error can be posed apart: no content at all, versus JSON cut off
mid-object."""
router = MagicMock()
router.model_group_alias = {}
router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}])
async def acompletion(**kwargs):
if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN:
kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"}
return {"choices": [{"message": {"content": "shadow answer"}}]}
return ModelResponse(
model=served_model,
choices=[{"index": 0, "finish_reason": finish_reason, "message": {"role": "assistant", "content": content}}],
)
router.acompletion = MagicMock(side_effect=acompletion)
return router
TOOL_CALL_MESSAGE = {
"content": None,
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
@ -1258,6 +1301,79 @@ class TestShadowPipeline:
assert row["judge_cost"] == expected_cost
assert row["shadow_cost"] == expected_shadow_cost
async def _judge_error(self, router: MagicMock, monkeypatch: pytest.MonkeyPatch) -> str:
import litellm as litellm_module
monkeypatch.setattr(litellm_module, "completion_cost", lambda completion_response: 0.007)
prisma = _prisma()
await _logger(router=router, prisma=prisma)._run_shadow_eval(
job=_job(),
request_id="req-1",
messages=({"role": "user", "content": "hi"},),
real_text="real answer",
real_model="claude-opus",
real_cost=0.0,
real_classifier_cost=0.0,
real_cache_hit=False,
control_tier=None,
shadow_params={},
parent_metadata={},
)
return prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]["error"]
async def test_a_judge_that_answered_nothing_is_told_apart_from_one_cut_off(
self, monkeypatch: pytest.MonkeyPatch
):
"""Both land on the same parser message, and they want opposite fixes: a judge
returning no content points at the reply never being text, while one cut off
mid-object points at the output cap. The row has to say which."""
truncated = '{"preference": "A", "confidence": 0.9, "reasoning": "'
answered_nothing = await self._judge_error(_judge_reply_router(None), monkeypatch)
cut_off = await self._judge_error(
_judge_reply_router(truncated, finish_reason="length"), monkeypatch
)
assert "content=no content" in answered_nothing
assert "finish_reason=stop" in answered_nothing
assert f"content={len(truncated)} chars" in cut_off
assert "finish_reason=length" in cut_off
async def test_an_unparseable_verdict_names_the_model_that_served_it(self, monkeypatch: pytest.MonkeyPatch):
"""A judge_model that fans out over deployments hides which one truncates: without
the served model the operator cannot tell a bad deployment from a bad cap."""
error = await self._judge_error(_judge_reply_router(None, served_model="claude-sonnet-5"), monkeypatch)
assert "model=claude-sonnet-5" in error
async def test_a_diagnosed_verdict_error_stays_groupable(self, monkeypatch: pytest.MonkeyPatch):
"""The customer groups attempt rows by error text. Every varying part has to sit
after the first semicolon or each row becomes its own group."""
first = await self._judge_error(_judge_reply_router(None, served_model="model-a"), monkeypatch)
second = await self._judge_error(_judge_reply_router(None, served_model="model-b"), monkeypatch)
assert first != second
assert first.split(";")[0] == second.split(";")[0]
async def test_a_judge_reply_that_cannot_be_read_still_records_an_error(self, monkeypatch: pytest.MonkeyPatch):
"""The shape reader runs inside the failure path: it must never raise a second time
and cost the row entirely."""
router = MagicMock()
router.model_group_alias = {}
router.get_model_list = MagicMock(return_value=[{"litellm_params": {"model": "openai/gpt-4o-mini"}}])
async def acompletion(**kwargs):
if kwargs["metadata"].get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == SHADOW_EVAL_ROUTER_CALL_ORIGIN:
kwargs["metadata"]["routing_decision"] = {"tier_label": "SIMPLE", "routed_model": "cheap-model"}
return {"choices": [{"message": {"content": "shadow answer"}}]}
return {"choices": []}
router.acompletion = MagicMock(side_effect=acompletion)
error = await self._judge_error(router, monkeypatch)
assert "unparseable judge verdict" in error
assert "unreadable judge reply" in error
async def test_an_empty_shadow_reply_still_bills_its_cost(self, monkeypatch: pytest.MonkeyPatch):
"""A shadow call that returns no extractable text has still billed; pricing it at
zero would keep the dollar gate open while shadow calls keep charging the key."""
@ -1287,6 +1403,34 @@ class TestShadowPipeline:
assert row["shadow_cost"] == 0.007
assert logger._test_counter["spend:shadow_eval:job-1"] == 0.007
async def test_the_judge_output_cap_leaves_room_for_a_reasoning_judge(self):
"""The output cap covers reasoning tokens as well as the answer, and a judge_model
deployment carrying an elevated reasoning_effort spends that budget before it writes
anything. A cap sized for the verdict JSON alone goes entirely to reasoning and the
reply arrives empty, which the attempt records as an unparseable verdict rather than
a result. The judge here burns a reasoning budget a live claude-sonnet-5 call was
measured at, so the cap has to clear it for the verdict to survive."""
reasoning_tokens = 2000
logger = _logger(router=_reasoning_judge_router(reasoning_tokens), prisma=(prisma := _prisma()))
await logger._run_shadow_eval(
job=_job(),
request_id="req-1",
messages=({"role": "user", "content": "hi"},),
real_text="real answer",
real_model="claude-opus",
real_cost=0.0,
real_classifier_cost=0.0,
real_cache_hit=False,
control_tier=None,
shadow_params={},
parent_metadata={},
)
row = prisma.db.litellm_shadowevalattempt.create.call_args.kwargs["data"]
assert row["outcome"] in ("real", "shadow", "tie"), row["error"]
assert row["error"] is None
async def _no_text_error(self, router) -> str:
prisma = _prisma()
await _logger(router=router, prisma=prisma)._run_shadow_eval(

164
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -995,6 +995,35 @@ async def test_anthropic_messages_marks_litellm_params_async():
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_arealtime_marks_litellm_params_async(monkeypatch):
"""LIT-6973: ``_arealtime`` must plant ``_arealtime`` in ``litellm_params`` so
``_is_sync_litellm_request`` classifies the session async and a failed session
reaches a CustomLogger's failure hook once, through the async path only, even
though the sync ``failure_handler`` still runs ahead of the async one."""
captured = {}
async_logged = asyncio.Event()
class CaptureLogger(CustomLogger):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
captured["litellm_params"] = kwargs.get("litellm_params", {})
async_logged.set()
logger = CaptureLogger()
logger.log_failure_event = MagicMock()
monkeypatch.setattr(litellm, "callbacks", [logger])
monkeypatch.setattr(litellm, "failure_callback", [])
monkeypatch.setattr(litellm, "_async_failure_callback", [])
monkeypatch.setattr(litellm, "success_callback", [])
monkeypatch.setattr(litellm, "_async_success_callback", [])
with pytest.raises(ValueError, match="Unsupported model"):
await litellm._arealtime(model="anthropic/claude-x", websocket=MagicMock())
await asyncio.wait_for(async_logged.wait(), timeout=10)
logger.log_failure_event.assert_not_called()
assert captured["litellm_params"].get("_arealtime") is True
assert LitellmLogging._is_sync_litellm_request(captured["litellm_params"]) is False
@pytest.mark.asyncio
async def test_agenerate_content_marks_litellm_params_async():
"""LIT-4475: the async ``agenerate_content`` entrypoint must plant
@ -1085,6 +1114,56 @@ async def test_logging_non_streaming_request():
litellm.callbacks = original_callbacks
@pytest.mark.asyncio
async def test_async_success_handler_truncates_large_base64_off_the_event_loop(monkeypatch):
"""The standard logging payload's base64 scan of a large multimodal request must not run on the loop thread."""
import threading
from litellm.litellm_core_utils import logging_utils
loop_thread = threading.get_ident()
scan_threads: list[int] = []
original_scan = logging_utils._truncate_base64_in_string
def recording_scan(value: str) -> str:
scan_threads.append(threading.get_ident())
return original_scan(value)
monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan)
monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000)
logged = asyncio.Event()
captured: dict = {}
class CaptureLogger(CustomLogger):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
captured["standard_logging_object"] = kwargs["standard_logging_object"]
logged.set()
monkeypatch.setattr(litellm, "callbacks", [CaptureLogger()])
payload = "L" * 20_000
await litellm.acompletion(
model="openai/gpt-5.6",
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": "describe"},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}},
],
}
],
mock_response="ok",
)
await asyncio.wait_for(logged.wait(), timeout=10)
logged_url = captured["standard_logging_object"]["messages"][0]["content"][1]["image_url"]["url"]
assert "base64_data truncated" in logged_url
assert payload not in logged_url
assert scan_threads
assert loop_thread not in scan_threads
@pytest.mark.parametrize(
"async_flag",
[
@ -1180,6 +1259,7 @@ def test_is_sync_litellm_request():
assert LitellmLogging._is_sync_litellm_request({}) is True
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
assert LitellmLogging._is_sync_litellm_request({"allm_passthrough_route": True}) is False
assert LitellmLogging._is_sync_litellm_request({"_arealtime": True}) is False
assert LitellmLogging._is_sync_litellm_request({"aanthropic_messages": True}) is False
assert LitellmLogging._is_sync_litellm_request({"agenerate_content": True}) is False
assert LitellmLogging._is_sync_litellm_request({"agenerate_content_stream": True}) is False
@ -1466,6 +1546,62 @@ async def test_dispatch_failure_handlers_async_completes_before_sync_submit(
assert events == ["async_start", "async_end", "sync_submit"]
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_submits_sync_handler_when_task_is_cancelled(
logging_obj,
):
"""Cancelling the dispatch task mid-await still submits the sync failure_handler.
Router failure paths fire the dispatcher with ``asyncio.create_task`` and raise
right away. When the event loop is torn down before the task finishes (a short
``asyncio.run`` in the SDK), the cancelled task must still hand the sync callbacks
to the executor, as the old raw-thread path did, and only once the async handler
has stopped.
"""
exception = ValueError("boom")
traceback_exception = "traceback"
events: list[str] = []
async_started = asyncio.Event()
async def _async_failure(exc, tb, **kwargs):
events.append("async_start")
async_started.set()
await asyncio.sleep(10)
events.append("async_end")
def _submit(*args, **kwargs):
events.append("sync_submit")
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure),
patch.object(logging_obj, "failure_handler", new_callable=MagicMock),
patch.object(
logging_obj,
"_should_run_sync_failure_callbacks_for_async_calls",
return_value=True,
),
patch( # test-quality-ok: the executor submit is the observable
"litellm.litellm_core_utils.litellm_logging.executor.submit",
side_effect=_submit,
),
):
task = asyncio.create_task(
logging_obj.dispatch_failure_handlers(
exception,
traceback_exception,
prefer_async_handlers=True,
)
)
await async_started.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert events == ["async_start", "sync_submit"]
@pytest.mark.asyncio
async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks(
logging_obj,
@ -5997,6 +6133,34 @@ def test_failure_handler_helper_fn_builds_payload_once_per_exception():
assert obj.model_call_details["standard_logging_object"] is not first_payload
@pytest.mark.asyncio
async def test_sync_failure_handler_reuses_payload_after_callable_async_callback():
"""Regression for LIT-6886: the proxy runs async_failure_handler, then the threaded
failure_handler, for every rejected request. A plain-function async callback (the
Router registers one) is dispatched through CustomLogger.async_log_event, which
restamps log_event_type on the shared model_call_details; the sync handler then
rebuilt the standardized payload, doubling the redaction and payload cost of a 403."""
router_style_callback = AsyncMock()
obj = LitellmLogging(
model="gpt-4o",
messages=[{"role": "user", "content": "Hey"}],
stream=False,
call_type="acompletion",
start_time=time.time(),
litellm_call_id="lit-6886-1",
function_id="f",
dynamic_async_failure_callbacks=[router_style_callback],
)
exc = _raise_and_catch(_ClientError(status_code=403, message="key not allowed to access model"))
await obj.async_failure_handler(exception=exc, traceback_exception="")
first_payload = obj.model_call_details["standard_logging_object"]
assert first_payload is not None
assert router_style_callback.await_count == 1
obj.failure_handler(exc, "")
assert obj.model_call_details["standard_logging_object"] is first_payload
@pytest.mark.asyncio
async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_obj):
"""The savings gate reads litellm_gateway_injected_cache from the request's

71
tests/test_litellm/litellm_core_utils/test_logging_utils.py
View file

@ -2,12 +2,16 @@
Tests for litellm.litellm_core_utils.logging_utils — base64 truncation helpers.
"""
import threading
import pytest
from litellm.litellm_core_utils import logging_utils
from litellm.litellm_core_utils.logging_utils import (
_format_base64_size,
_truncate_base64_in_string,
truncate_base64_in_messages,
truncate_base64_in_messages_async,
)
# ---------------------------------------------------------------------------
@ -157,3 +161,70 @@ class TestTruncateBase64InMessages:
result[0]["content"][0]["image_url"]["url"]
== f"data:image/png;base64,{short}"
)
# ---------------------------------------------------------------------------
# truncate_base64_in_messages_async
# ---------------------------------------------------------------------------
def _image_messages(payload: str) -> list:
return [
{
"role": "user",
"content": [
{"type": "text", "text": "describe"},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{payload}"}},
],
}
]
@pytest.fixture
def scan_threads(monkeypatch):
"""Record the thread that runs every base64 regex scan."""
threads: list[int] = []
original = logging_utils._truncate_base64_in_string
def recording_scan(value: str) -> str:
threads.append(threading.get_ident())
return original(value)
monkeypatch.setattr(logging_utils, "_truncate_base64_in_string", recording_scan)
return threads
class TestTruncateBase64InMessagesAsync:
@pytest.mark.asyncio
async def test_large_payload_is_scanned_off_the_event_loop(self, monkeypatch, scan_threads):
monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000)
payload = "I" * 20_000
messages = _image_messages(payload)
result = await truncate_base64_in_messages_async(messages)
offload_threads = tuple(scan_threads)
assert result == truncate_base64_in_messages(messages)
assert payload not in result[0]["content"][1]["image_url"]["url"]
assert payload in messages[0]["content"][1]["image_url"]["url"]
assert offload_threads
assert threading.get_ident() not in offload_threads
@pytest.mark.asyncio
async def test_small_payload_stays_on_the_calling_thread(self, monkeypatch, scan_threads):
monkeypatch.setattr(logging_utils, "BASE64_TRUNCATION_OFFLOAD_THRESHOLD_CHARS", 1_000)
messages = _image_messages("J" * 200)
result = await truncate_base64_in_messages_async(messages)
assert result == truncate_base64_in_messages(messages)
assert scan_threads
assert set(scan_threads) == {threading.get_ident()}
@pytest.mark.asyncio
async def test_none_and_disabled_truncation_short_circuit(self, monkeypatch, scan_threads):
assert await truncate_base64_in_messages_async(None) is None
monkeypatch.setattr(logging_utils, "MAX_BASE64_LENGTH_FOR_LOGGING", 0)
messages = _image_messages("K" * 20_000)
assert await truncate_base64_in_messages_async(messages) is messages
assert scan_threads == []

10
tests/test_litellm/litellm_core_utils/test_realtime_errors.py
View file

@ -1,8 +1,10 @@
import json
import pytest
from litellm.litellm_core_utils.realtime_errors import (
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
client_close_code,
realtime_error_event,
websocket_close_reason,
)
@ -42,3 +44,11 @@ def test_websocket_close_reason_truncates_multibyte_message_by_bytes():
assert len(reason.encode("utf-8")) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES
assert reason == "あ" * (WEBSOCKET_CLOSE_REASON_MAX_BYTES // 3)
assert "<EFBFBD>" not in reason
@pytest.mark.parametrize(
("upstream_code", "expected"),
[(1000, 1000), (1008, 1008), (1011, 1011), (4001, 4001), (1005, 1011), (1006, 1011), (1015, 1011), (2999, 1011)],
)
def test_client_close_code_only_forwards_codes_a_server_may_send(upstream_code, expected):
assert client_close_code(upstream_code) == expected

317
tests/test_litellm/litellm_core_utils/test_realtime_streaming.py
View file

@ -1,14 +1,20 @@
import asyncio
import json
from collections.abc import Coroutine
from dataclasses import dataclass
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from websockets.exceptions import ConnectionClosed
from websockets.frames import Close
import litellm
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.realtime_streaming import (
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
RealTimeStreaming,
client_sent_openai_beta_realtime_header,
)
@ -2941,13 +2947,11 @@ async def test_log_messages_routes_async_logging_through_bounded_worker():
realtime turn leaves a suspended task pinning its response in memory -> an
unbounded leak. Regression for that fix."""
logging_obj = MagicMock()
streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj)
mock_worker = MagicMock()
streaming = RealTimeStreaming(MagicMock(), MagicMock(), logging_obj, logging_worker=mock_worker)
streaming.messages = [{"type": "session.created"}]
with (
patch("litellm.litellm_core_utils.realtime_streaming.GLOBAL_LOGGING_WORKER") as mock_worker,
patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task,
):
with patch("litellm.litellm_core_utils.realtime_streaming.asyncio.create_task") as mock_create_task:
await streaming.log_messages()
mock_worker.ensure_initialized_and_enqueue.assert_called_once()
@ -3028,12 +3032,12 @@ async def test_session_close_flushes_unbilled_transcription_usage():
messages before log_messages runs, and never forwarded to the client."""
from typing import Final
from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage
from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeResponseTypedDict
client_ws: Final = MagicMock()
client_ws.send_text = AsyncMock()
backend_ws: Final = MagicMock()
backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None))
backend_ws.recv = AsyncMock(side_effect=[b'{"serverContent": {}}', ConnectionClosed(None, None)])
logging_obj: Final = MagicMock()
logging_obj.async_success_handler = AsyncMock()
logging_obj.success_handler = MagicMock()
@ -3045,7 +3049,24 @@ async def test_session_close_flushes_unbilled_transcription_usage():
"total_tokens": 171,
"input_token_details": {"text_tokens": 0, "audio_tokens": 153},
}
transcript_frame: Final[RealtimeResponseTypedDict] = {
"response": {
"type": "conversation.item.input_audio_transcription.completed",
"event_id": "event_1",
"transcript": "ahoy",
"item_id": "item_1",
"content_index": 0,
},
"current_output_item_id": None,
"current_response_id": None,
"current_delta_chunks": None,
"current_conversation_id": None,
"current_item_chunks": None,
"current_delta_type": None,
"session_configuration_request": None,
}
provider_config: Final = MagicMock()
provider_config.transform_realtime_response = MagicMock(return_value=transcript_frame)
provider_config.unbilled_usage_on_session_close = MagicMock(return_value=usage)
streaming: Final = RealTimeStreaming(
@ -3077,7 +3098,9 @@ async def test_session_close_flushes_unbilled_transcription_usage():
)
assert len(flushed) == 1
assert flushed[0] in logged_snapshots[0]
assert not client_ws.send_text.called
forwarded: Final = tuple(json.loads(call.args[0]) for call in client_ws.send_text.await_args_list)
assert [event.get("transcript") for event in forwarded] == ["ahoy"]
assert all("usage" not in event for event in forwarded)
@pytest.mark.asyncio
@ -3111,3 +3134,281 @@ async def test_session_close_flush_noop_without_unbilled_usage():
isinstance(message, dict) and message.get("type") == "conversation.item.input_audio_transcription.completed"
for message in streaming.messages
)
_UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found"
class _InlineLoggingWorker:
def __init__(self) -> None:
self.enqueued: tuple[Coroutine[object, object, None], ...] = ()
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None:
self.enqueued = (*self.enqueued, async_coroutine)
async def drain(self) -> None:
for coroutine in self.enqueued:
await coroutine
class _RecordingLogging:
def __init__(self) -> None:
self.model_call_details: dict[str, object] = {}
self.logged_sessions: tuple[tuple[dict, ...], ...] = ()
self.logged_failures: tuple[Exception, ...] = ()
def pre_call(self, input: str | dict, api_key: str) -> None:
return None
async def dispatch_success_handlers(self, result: list[dict], prefer_async_handlers: bool = False) -> None:
self.logged_sessions = (*self.logged_sessions, tuple(result))
async def dispatch_failure_handlers(
self, exception: Exception, traceback_exception: str, prefer_async_handlers: bool = False
) -> None:
self.logged_failures = (*self.logged_failures, exception)
@dataclass(frozen=True, slots=True)
class _RelaySession:
streaming: RealTimeStreaming
logging: _RecordingLogging
worker: _InlineLoggingWorker
async def run(self) -> None:
await asyncio.wait_for(self.streaming.bidirectional_forward(), timeout=2)
await self.worker.drain()
async def _wait_forever() -> str:
await asyncio.Event().wait()
raise AssertionError("unreachable")
def _client_ws_that_never_sends() -> MagicMock:
client_ws: Final = MagicMock()
client_ws.headers = {}
client_ws.receive_text = AsyncMock(side_effect=_wait_forever)
client_ws.send_text = AsyncMock()
client_ws.close = AsyncMock()
return client_ws
def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock:
backend_ws: Final = MagicMock()
backend_ws.recv = AsyncMock(side_effect=list(frames))
return backend_ws
def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession:
logging: Final = _RecordingLogging()
worker: Final = _InlineLoggingWorker()
streaming: Final = RealTimeStreaming(
client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker
)
return _RelaySession(streaming=streaming, logging=logging, worker=worker)
def _error_events_sent_to(client_ws: MagicMock) -> list[dict]:
events: Final = (json.loads(call.args[0]) for call in client_ws.send_text.await_args_list)
return [event for event in events if event.get("type") == "error"]
@pytest.mark.asyncio
async def test_bidirectional_forward_relays_upstream_policy_close_to_client():
client_ws: Final = _client_ws_that_never_sends()
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close))
await session.run()
(error_event,) = _error_events_sent_to(client_ws)
assert error_event["error"]["type"] == "server_error"
assert "1008" in error_event["error"]["message"]
assert _UPSTREAM_REFUSAL in error_event["error"]["message"]
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
@pytest.mark.parametrize(
"leaked_detail",
(
pytest.param("sk-live-abcdef0123456789abcdef0123", id="credential"),
pytest.param("vertex-int.svc.cluster.local", id="internal-hostname"),
pytest.param("/etc/litellm/service-account.json", id="filesystem-path"),
),
)
@pytest.mark.asyncio
async def test_upstream_close_details_are_scrubbed_before_reaching_the_client(leaked_detail: str):
"""LIT-6973: the relayed close goes through the proxy's client-facing redaction, so an upstream
error echoing a credential, an internal host, or a server path never reaches the client verbatim."""
client_ws: Final = _client_ws_that_never_sends()
upstream_close: Final = ConnectionClosed(Close(1008, f"upstream rejected: {leaked_detail}"), None)
session: Final = _relay_session(client_ws, _backend_ws_closing_with(upstream_close))
await session.run()
(error_event,) = _error_events_sent_to(client_ws)
assert leaked_detail not in error_event["error"]["message"]
relayed_reason: Final = client_ws.close.await_args.kwargs["reason"]
assert leaked_detail not in relayed_reason
assert "REDACTED" in relayed_reason
@pytest.mark.asyncio
async def test_bidirectional_forward_maps_abnormal_upstream_close_to_internal_error():
client_ws: Final = _client_ws_that_never_sends()
session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(None, None)))
await session.run()
(error_event,) = _error_events_sent_to(client_ws)
assert "1006" in error_event["error"]["message"]
client_ws.close.assert_awaited_once()
assert client_ws.close.await_args.kwargs["code"] == 1011
@pytest.mark.asyncio
async def test_bidirectional_forward_relays_normal_upstream_close_without_error_event():
client_ws: Final = _client_ws_that_never_sends()
session: Final = _relay_session(client_ws, _backend_ws_closing_with(ConnectionClosed(Close(1000, ""), None)))
await session.run()
assert _error_events_sent_to(client_ws) == []
client_ws.close.assert_awaited_once()
assert client_ws.close.await_args.kwargs["code"] == 1000
@pytest.mark.asyncio
async def test_upstream_refusal_before_any_frame_logs_a_failure_not_a_success():
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
await session.run()
assert session.logging.logged_failures == (upstream_close,)
assert session.logging.logged_sessions == ()
@pytest.mark.asyncio
async def test_upstream_refusal_after_a_synthetic_session_created_still_logs_a_failure():
"""LIT-6973: deferred Gemini Live setup stores a synthetic ``session.created`` before
the relay starts. It is not an upstream frame, so a refusal after it is still a refusal."""
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
session.streaming.store_message(json.dumps({"type": "session.created", "session": {"id": "sess_synthetic"}}))
await session.run()
assert session.logging.logged_failures == (upstream_close,)
assert session.logging.logged_sessions == ()
@pytest.mark.asyncio
async def test_upstream_close_after_relayed_events_still_logs_the_session_as_success():
client_ws: Final = _client_ws_that_never_sends()
session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close))
await session.run()
(logged_session,) = session.logging.logged_sessions
assert [event["type"] for event in logged_session] == ["session.created"]
assert session.logging.logged_failures == ()
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
@pytest.mark.asyncio
async def test_upstream_closing_while_a_client_message_is_forwarded_still_reaches_the_client():
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
backend_closed: Final = asyncio.Event()
client_messages: Final = iter((json.dumps({"type": "response.create"}),))
async def receive_text() -> str:
message = next(client_messages, None)
return message if message is not None else await _wait_forever()
async def send_to_backend(_message: str) -> None:
backend_closed.set()
raise upstream_close
async def recv_from_backend() -> bytes:
await backend_closed.wait()
raise upstream_close
client_ws: Final = _client_ws_that_never_sends()
client_ws.receive_text = receive_text
backend_ws: Final = MagicMock()
backend_ws.send = send_to_backend
backend_ws.recv = recv_from_backend
session: Final = _relay_session(client_ws, backend_ws)
await session.run()
(error_event,) = _error_events_sent_to(client_ws)
assert _UPSTREAM_REFUSAL in error_event["error"]["message"]
client_ws.close.assert_awaited_once_with(code=1008, reason=_UPSTREAM_REFUSAL)
assert session.logging.logged_failures == (upstream_close,)
@pytest.mark.asyncio
async def test_client_hanging_up_first_ends_the_session_without_a_relayed_close():
client_ws: Final = _client_ws_that_never_sends()
client_ws.receive_text = AsyncMock(side_effect=RuntimeError("client went away"))
backend_ws: Final = MagicMock()
backend_ws.recv = AsyncMock(side_effect=_wait_forever)
session: Final = _relay_session(client_ws, backend_ws)
await session.run()
assert session.logging.logged_sessions == ((),)
assert session.logging.logged_failures == ()
client_ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_client_hanging_up_with_a_websockets_close_is_not_mistaken_for_the_backend_closing():
client_ws: Final = _client_ws_that_never_sends()
client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(None, None))
backend_ws: Final = MagicMock()
backend_ws.recv = AsyncMock(side_effect=_wait_forever)
session: Final = _relay_session(client_ws, backend_ws)
await session.run()
assert session.logging.logged_sessions == ((),)
assert session.logging.logged_failures == ()
client_ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_success_logging_stamps_the_reservation_ownership_marker():
"""LIT-6973: only the success path enqueues the cost callback that settles the
session's budget reservation, so it stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on
the shared logging object. The proxy endpoint reads that stamp to decide whether to
release the reservation itself, so a logged-as-success session must carry it."""
client_ws: Final = _client_ws_that_never_sends()
session_created: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
upstream_close: Final = ConnectionClosed(Close(1000, ""), None)
session: Final = _relay_session(client_ws, _backend_ws_closing_with(session_created, upstream_close))
await session.run()
assert session.logging.logged_sessions != ()
assert session.logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
@pytest.mark.asyncio
async def test_refused_session_does_not_stamp_the_reservation_ownership_marker():
"""A refused session logs a failure, not a success, so it must not stamp
REALTIME_SESSION_SUCCESS_LOGGED_KEY. If it did, the proxy endpoint would skip its
own reservation release and the refused session's reservation would stay pinned."""
upstream_close: Final = ConnectionClosed(Close(1008, _UPSTREAM_REFUSAL), None)
session: Final = _relay_session(_client_ws_that_never_sends(), _backend_ws_closing_with(upstream_close))
await session.run()
assert session.logging.logged_failures == (upstream_close,)
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details

31
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -1075,6 +1075,37 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
assert result == responses_so_far
class TestUndecoratedGuardrailIsRecorded:
"""LIT-5983 regression: the handler calls apply_guardrail bare, so a custom guardrail
without @log_guardrail_information must still end up in the request's guardrail
information on both the request and response paths."""
@pytest.mark.asyncio
async def test_request_path_records_undecorated_guardrail(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="docs-style")
data = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
await handler.process_input_messages(data, guardrail)
entries = data["metadata"]["standard_logging_guardrail_information"]
assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")]
@pytest.mark.asyncio
async def test_response_path_records_undecorated_guardrail(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="docs-style")
response = ModelResponse(
choices=[Choices(finish_reason="stop", index=0, message=Message(content="hi", role="assistant"))]
)
request_data: dict = {"metadata": {}}
await handler.process_output_response(response, guardrail, request_data=request_data)
entries = request_data["metadata"]["standard_logging_guardrail_information"]
assert [(e["guardrail_name"], e["guardrail_status"]) for e in entries] == [("docs-style", "success")]
class TestGetStructuredMessages:
"""Test the get_structured_messages method."""

12
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_adapter.py
View file

@ -13,6 +13,7 @@ from fastapi import HTTPException
from pydantic import ValidationError
from litellm.experimental_mcp_client.client import MCPClient
from litellm.proxy._experimental.mcp_server.exceptions import MCPServerURLCredentialsError
from litellm.proxy._experimental.mcp_server.outbound_credentials.adapter import (
oauth_protected_resource_path,
raise_public,
@ -442,6 +443,17 @@ def test_raise_public_maps_each_error_to_its_status(error, status):
assert exc_info.value.status_code == status
def test_raise_public_marks_only_url_credentials_error_as_safe_for_preview():
with pytest.raises(HTTPException) as generic_exc_info:
raise_public(CredError.of_misconfigured("private operator detail"))
assert not isinstance(generic_exc_info.value, MCPServerURLCredentialsError)
error = CredError.of_url_credentials_not_allowed()
with pytest.raises(MCPServerURLCredentialsError) as url_exc_info:
raise_public(error)
assert url_exc_info.value.detail == error.summary
def test_raise_public_emits_unauthorized_challenge():
body = {"error": "byok_auth_required", "server_id": "s1"}
error = CredError.of_unauthorized("needs key", www_authenticate='Bearer resource_metadata="/x"', body=body)

29
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_resolver.py
View file

@ -116,6 +116,35 @@ async def test_none_mode_yields_a_no_op_auth():
assert isinstance(result.ok, NoOpAuth)
@pytest.mark.asyncio
async def test_none_mode_rejects_url_userinfo():
spec = ServerSpec(
server_id="s",
resource="https://lit-user:s3cr3t@upstream.example.com/mcp",
config=NoneConfig(),
)
result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec)
assert isinstance(result, Error)
assert result.error.tag == "url_credentials_not_allowed"
assert "Basic Auth" in result.error.summary
assert "auth_type: basic" in result.error.summary
assert "auth_value: username:password" in result.error.summary
assert "lit-user" not in result.error.summary
assert "s3cr3t" not in result.error.summary
@pytest.mark.asyncio
async def test_none_mode_does_not_validate_non_credential_resource():
spec = ServerSpec(server_id="s", resource="https://[::1", config=NoneConfig())
result = await UpstreamCredentialProvider().resolve_credentials(_SUBJECT, spec)
assert isinstance(result, Ok)
assert isinstance(result.ok, NoOpAuth)
@pytest.mark.asyncio
async def test_api_key_shared_emits_the_configured_header():
config = ApiKeyConfig(

9
tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_types.py
View file

@ -75,6 +75,15 @@ def test_crederror_factory_sets_the_matching_tag(factory, expected_tag):
assert "detail text" in err.summary
def test_url_credentials_error_has_a_fixed_actionable_summary():
err = CredError.of_url_credentials_not_allowed()
assert err.tag == "url_credentials_not_allowed"
assert "Basic Auth" in err.summary
assert "auth_type: basic" in err.summary
assert "auth_value: username:password" in err.summary
def test_apikeyconfig_requires_a_key_source():
with pytest.raises(ValidationError):
ApiKeyConfig() # type: ignore[call-arg]

42
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -8966,6 +8966,24 @@ class TestCreateMcpClientV2Graft:
assert isinstance(client._resolved_auth, NoOpAuth)
assert client._mcp_auth_value is None
@pytest.mark.parametrize("auth_type", [None, MCPAuth.none])
async def test_none_mode_rejects_url_userinfo(self, auth_type):
with pytest.raises(HTTPException) as exc_info:
await MCPServerManager()._create_mcp_client(
self._http_server(
auth_type=auth_type,
url="https://lit-user:s3cr3t@upstream.example.com/mcp",
)
)
detail = str(exc_info.value.detail)
assert exc_info.value.status_code == 500
assert "Basic Auth" in detail
assert "auth_type: basic" in detail
assert "auth_value: username:password" in detail
assert "lit-user" not in detail
assert "s3cr3t" not in detail
@pytest.mark.parametrize(
"auth_type, token, expected_name, expected_value",
[
@ -11369,6 +11387,30 @@ class TestResolveOpenapiToolAuth:
assert "Authorization" not in (forwarded or {})
@pytest.mark.asyncio
async def test_none_mode_without_url_keeps_spec_path_server_unauthenticated(self):
server = MCPServer(
server_id="openapi-only",
name="report_api",
server_name="report_api",
url=None,
transport=MCPTransport.http,
auth_type=MCPAuth.none,
spec_path="https://api.example.com/openapi.json",
)
resolved, forwarded = await MCPServerManager().resolve_openapi_upstream_auth(
mcp_server=server,
oauth2_headers=None,
raw_headers=None,
mcp_auth_header=None,
user_api_key_auth=None,
forwarded_headers={"X-Trace": "trace-id"},
)
assert resolved is None
assert forwarded == {"X-Trace": "trace-id"}
class TestOpenApiHandlerRelaysUpstreamAuth:
"""`_call_openapi_tool_handler` must not flatten a re-auth signal into a generic message.

30
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -177,6 +177,36 @@ class TestExecuteWithMcpClient:
assert "https://api.example.com/mcp/" in message
assert "30s" in message
def test_connection_error_message_hides_arbitrary_http_exception_detail(self):
message = rest_endpoints._connection_error_message(
HTTPException(status_code=500, detail="secret upstream detail"),
"https://api.example.com/mcp/",
30.0,
)
assert "secret upstream detail" not in message
@pytest.mark.asyncio
async def test_none_mode_url_credentials_returns_actionable_redacted_error(self):
async def unreached_operation(client):
raise AssertionError("operation must not run for an invalid server configuration")
payload = NewMCPServerRequest(
server_name="example",
url="https://lit-user:s3cr3t@upstream.example.com/mcp",
auth_type=MCPAuth.none,
)
result = await rest_endpoints._execute_with_mcp_client(payload, unreached_operation)
message = str(result["message"])
assert result["error"] is True
assert "Basic Auth" in message
assert "auth_type: basic" in message
assert "auth_value: username:password" in message
assert "lit-user" not in message
assert "s3cr3t" not in message
@pytest.mark.asyncio
async def test_forwards_static_headers(self, monkeypatch):
"""Ensure static_headers are forwarded to the MCP client during test calls.

18
tests/test_litellm/proxy/auth/test_litellm_license.py
View file

@ -34,27 +34,27 @@ def test_is_over_limit():
assert license_check.is_over_limit(99) is False
def test_heuristic_v2_router_limit() -> None:
def test_auto_router_capability_limit() -> None:
"""Only the signed license's auto_router feature lifts the one-router limit; an API-verified
license (no airgapped data) and an airgapped license without the feature keep it."""
license_check = LicenseCheck()
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["auto_router"]}
assert license_check.heuristic_v2_router_limit() is None
assert license_check.auto_router_capability_limit() is None
license_check.airgapped_license_data = {
"expiration_date": "2999-01-01",
"allowed_features": ["sso", "auto_router", "audit_logs"],
}
assert license_check.heuristic_v2_router_limit() is None
assert license_check.auto_router_capability_limit() is None
license_check.airgapped_license_data = {"expiration_date": "2999-01-01", "allowed_features": ["sso"]}
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
license_check.airgapped_license_data = {"expiration_date": "2999-01-01"}
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
license_check.airgapped_license_data = None
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
def _signed_license(expiration_date: str) -> tuple[RSAPublicKey, str]:
@ -81,12 +81,12 @@ def test_expired_or_unreadable_license_grants_no_features() -> None:
license_check = LicenseCheck()
public_key, valid_key = _signed_license("2999-01-01")
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
assert license_check.heuristic_v2_router_limit() is None
assert license_check.auto_router_capability_limit() is None
_, expired_key = _signed_license("2000-01-01")
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=expired_key) is not True
assert license_check.airgapped_license_data is None
assert license_check.heuristic_v2_router_limit() == 1
assert license_check.auto_router_capability_limit() == 1
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=valid_key) is True
assert license_check.verify_license_without_api_request(public_key=public_key, license_key="not-a-license") is not True
@ -98,4 +98,4 @@ def test_valid_signed_license_with_auto_router_lifts_the_limit() -> None:
public_key, license_key = _signed_license("2999-01-01")
assert license_check.verify_license_without_api_request(public_key=public_key, license_key=license_key) is True
assert license_check.heuristic_v2_router_limit() is None
assert license_check.auto_router_capability_limit() is None

6
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -1137,7 +1137,11 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source():
mock_api.assert_called_once()
kwargs = mock_api.call_args.kwargs
assert kwargs["source"] == "OUTPUT"
assert kwargs["request_data"] == {"model": "gpt-4o"}
assert kwargs["request_data"]["model"] == "gpt-4o"
recorded = kwargs["request_data"]["metadata"]["standard_logging_guardrail_information"]
assert [(e["guardrail_name"], e["guardrail_status"]) for e in recorded] == [
(guardrail.guardrail_name, "success")
]
synthetic = kwargs["response"]
assert isinstance(synthetic, ModelResponse)
assert len(synthetic.choices) == 2

292
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -2,6 +2,7 @@ import inspect
import asyncio
import contextlib
import json
from collections.abc import Mapping
from typing import Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch
@ -4048,6 +4049,72 @@ class TestStrategyRouterWriteValidation:
)
assert _strategy_router_write_violation(incoming_params=None, existing_params=None) is None
@pytest.mark.parametrize(
"config",
[
{"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
{
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
],
)
def test_model_less_patch_cannot_attach_router_config_to_a_regular_model(self, config: dict[str, object]) -> None:
"""The license gate applies only to complexity routers, so a partial PATCH cannot poison a regular
model with a capability-shaped config and make it occupy a slot."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
)
from litellm.types.router import updateLiteLLMParams
violation = _strategy_router_write_violation(
incoming_params=updateLiteLLMParams(complexity_router_config=config),
existing_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
)
assert violation is not None
assert "does not start with 'auto_router/'" in violation
assert "complexity_router_config" in violation
def test_effective_params_decrypts_a_stored_complexity_router_model(self, monkeypatch) -> None:
"""A database row encrypts model, so the model-aware gate must not accidentally rely on plaintext mocks."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.management_endpoints.model_management_endpoints import (
_effective_complexity_router_params,
)
from litellm.types.router import updateLiteLLMParams
monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt")
encrypted_model = encrypt_value_helper("auto_router/complexity_router")
effective_params = _effective_complexity_router_params(
updateLiteLLMParams(complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini"}}),
LiteLLM_Params(model=encrypted_model),
)
assert effective_params["model"] == "auto_router/complexity_router"
def test_model_less_patch_keeps_a_complexity_router_in_scope(self) -> None:
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
)
from litellm.types.router import updateLiteLLMParams
assert (
_strategy_router_write_violation(
incoming_params=updateLiteLLMParams(
complexity_router_config={"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
),
existing_params=self._stored_complexity_params(),
)
is None
)
def test_restore_of_corrupted_row_is_allowed(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
_strategy_router_write_violation,
@ -4354,33 +4421,33 @@ class TestStrategyRouterWriteValidation:
)
@staticmethod
def _live_router_holding_one_heuristic_v2(limit: int | None) -> Router:
def _live_router_holding_one_capability(limit: int | None, config: Mapping[str, object]) -> Router:
return Router(
model_list=[
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k"}},
{
"model_name": "held-v2",
"model_name": "held",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}},
"complexity_router_config": config,
},
"model_info": {"id": "held-id"},
},
],
heuristic_v2_router_limit=lambda: limit,
auto_router_capability_limit=lambda: limit,
)
class _FakeTx:
"""Stands in for a prisma transaction: records the raw statements and exposes the model table."""
"""Stands in for a prisma transaction: records raw statements and returns encrypted-model candidates."""
def __init__(self, db_held: int) -> None:
self.db_held = db_held
def __init__(self, db_models: list[str]) -> None:
self.db_models = db_models
self.raw_calls: list[tuple[str, tuple[object, ...]]] = []
self.litellm_proxymodeltable = MagicMock(create=AsyncMock(), update=AsyncMock())
async def query_raw(self, sql: str, *args: object) -> list[dict[str, object]]:
self.raw_calls.append((sql, args))
return [{"held": self.db_held}] if "count(*)" in sql else []
return [{"model": model} for model in self.db_models] if "AS model" in sql else []
async def __aenter__(self) -> "TestStrategyRouterWriteValidation._FakeTx":
return self
@ -4391,9 +4458,9 @@ class TestStrategyRouterWriteValidation:
class _FakeDb:
"""Stands in for prisma_client: the plain client and the transaction it opens are told apart by identity."""
def __init__(self, db_held: int, existing_row: object = None) -> None:
def __init__(self, db_models: list[str], existing_row: object = None) -> None:
self.db = self
self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_held)
self.tx_obj = TestStrategyRouterWriteValidation._FakeTx(db_models)
self.litellm_proxymodeltable = MagicMock(
create=AsyncMock(), update=AsyncMock(), find_unique=AsyncMock(return_value=existing_row)
)
@ -4403,6 +4470,43 @@ class TestStrategyRouterWriteValidation:
_V2 = {"classifier_type": "heuristic_v2", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_V1 = {"classifier_type": "heuristic", "tiers": {"SIMPLE": "gpt-4o-mini"}}
_CUSTOM_TIERS = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
}
_TIER_LABELS_ONLY = {
"classifier_type": "heuristic",
"tiers": {"SIMPLE": "gpt-4o-mini"},
"tier_labels": {"SIMPLE": "Cheap"},
}
_CUSTOM_PROMPT = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
}
_OPERATOR_EXAMPLES = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_examples": '- "reset my password" -> SIMPLE',
}
_OPERATOR_OPENING_PROMPT = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_prompt": "Grade by data sensitivity",
}
_SHIPPED_RUBRIC = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "classification_rubric": "agentic"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
}
@pytest.mark.parametrize(
"incoming,existing,expected",
@ -4431,41 +4535,55 @@ class TestStrategyRouterWriteValidation:
@pytest.mark.asyncio
@pytest.mark.parametrize(
"limit,effective_config,db_held,config_holds_one,model_id,expected",
"limit,effective_params,db_models,config_config,model_id,expected",
[
(1, _V2, 1, False, None, "refused"),
(1, _V2, 0, True, None, "refused"),
(1, _V2, 0, False, None, "reserved"),
(1, _V2, 0, False, "held-id", "reserved"),
(2, _V2, 1, False, None, "reserved"),
(1, _V1, 5, True, None, "plain"),
(1, None, 5, True, None, "plain"),
(None, _V2, 5, True, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], _V2, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, [], None, "held-id", "reserved"),
(2, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, ["openai/gpt-4o"], None, None, "reserved"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIERS}, [], _CUSTOM_PROMPT, None, "refused"),
(1, {"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIERS}, ["auto_router/complexity_router"], None, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _V1}, ["auto_router/complexity_router"], _V2, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": None}, ["auto_router/complexity_router"], _V2, None, "plain"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _V2}, ["auto_router/complexity_router"], _V2, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _TIER_LABELS_ONLY}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_PROMPT}, ["auto_router/complexity_router"], None, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, [], _CUSTOM_TIERS, None, "refused"),
(1, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_OPENING_PROMPT}, [], _CUSTOM_PROMPT, None, "refused"),
(None, {"model": "auto_router/complexity_router", "complexity_router_config": _OPERATOR_EXAMPLES}, ["auto_router/complexity_router"], _CUSTOM_TIERS, None, "plain"),
],
)
async def test_heuristic_v2_slot_matrix(
async def test_auto_router_capability_slot_matrix(
self,
limit: int | None,
effective_config: object,
db_held: int,
config_holds_one: bool,
effective_params: Mapping[str, object],
db_models: list[str],
config_config: Mapping[str, object] | None,
model_id: str | None,
expected: str,
) -> None:
"""The slot is claimed inside a locked transaction only for a heuristic_v2 write under a limit; the DB rows
(other pods included) plus config.yaml routers decide, the row being edited is excluded through the SQL
parameter, and every other write runs on the plain client with no lock."""
"""The slot is claimed inside a locked transaction only for a write that claims a licensed capability
under a limit; the DB rows (other pods included) plus config.yaml routers decide, the row being edited
is excluded through the SQL parameter, and every other write runs on the plain client with no lock.
heuristic_v2 has its own slot, while custom tier definitions and custom prompts count into one shared
customization slot. Renaming built-in tiers through tier_labels claims nothing at all."""
from fastapi import HTTPException
from litellm.proxy.management_endpoints.model_management_endpoints import (
HEURISTIC_V2_SLOT_LOCK_KEY,
_heuristic_v2_slot,
AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,
_auto_router_capability_slot,
)
from litellm.router_utils.auto_router_model_naming import gated_capability_of
fake = self._FakeDb(db_held)
live_router = self._live_router_holding_one_heuristic_v2(limit) if config_holds_one else None
capability = gated_capability_of(effective_params)
fake = self._FakeDb(db_models)
live_router = self._live_router_holding_one_capability(limit, config_config) if config_config is not None else None
with (
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam
patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: limit), # test-quality-ok: the guard reads the proxy license singleton with no injection seam
patch("litellm.proxy.proxy_server.llm_router", live_router), # test-quality-ok: the guard reads the proxy router global with no injection seam
patch( # test-quality-ok: the cross-pod publish is the side effect under test; redis is not configured here
"litellm.proxy.management_endpoints.model_management_endpoints.publish_config_change",
@ -4474,13 +4592,15 @@ class TestStrategyRouterWriteValidation:
):
if expected == "refused":
with pytest.raises(HTTPException) as exc_info:
async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id):
async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id):
pass
assert exc_info.value.status_code == 403
assert capability is not None
assert "At most 1 auto-router" in str(exc_info.value.detail)
assert capability.subject in str(exc_info.value.detail)
assert "'auto_router' feature lifts the limit" in str(exc_info.value.detail)
return
async with _heuristic_v2_slot(fake, effective_config=effective_config, model_id=model_id) as tables:
async with _auto_router_capability_slot(fake, effective_params=effective_params, model_id=model_id) as tables:
handle = tables
if expected == "plain":
await handle.create(data={})
@ -4489,10 +4609,13 @@ class TestStrategyRouterWriteValidation:
return
assert handle is fake.tx_obj.litellm_proxymodeltable
published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable")
(lock_sql, lock_params), (_count_sql, count_params) = fake.tx_obj.raw_calls
(lock_sql, lock_params), (count_sql, count_params) = fake.tx_obj.raw_calls
assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql
assert lock_params == (HEURISTIC_V2_SLOT_LOCK_KEY,)
assert lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,)
assert count_params == (model_id or "",)
assert "AS model" in count_sql
assert capability is not None
assert capability.sql_config_predicate.split("{config}")[-1].strip() in count_sql
@pytest.mark.asyncio
async def test_team_model_bookkeeping_runs_after_the_slot_is_released(self) -> None:
@ -4549,14 +4672,14 @@ class TestStrategyRouterWriteValidation:
)
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
fake = self._FakeDb(db_held=1)
fake = self._FakeDb(["auto_router/complexity_router"])
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
@ -4579,6 +4702,93 @@ class TestStrategyRouterWriteValidation:
fake.tx_obj.litellm_proxymodeltable.create.assert_not_awaited()
fake.litellm_proxymodeltable.create.assert_not_awaited()
@pytest.mark.asyncio
async def test_model_less_patch_rejects_router_config_on_a_regular_model(self) -> None:
"""PATCH rejects the poison before its row write or the capability slot."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import patch_model
from litellm.types.router import updateLiteLLMParams
model_id = "regular-model"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
regular = Deployment(
model_name="regular-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
model_info={"id": model_id},
)
fake = self._FakeDb([])
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag
patch( # test-quality-ok: inject stored regular row without a database
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
new=AsyncMock(return_value=regular),
),
patch( # test-quality-ok: endpoint must reject before database authorization needs a live store
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
):
with pytest.raises(ProxyException) as exc_info:
await patch_model(
model_id=model_id,
patch_data=updateDeployment(
litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS)
),
user_api_key_dict=admin,
)
assert exc_info.value.code == "400"
assert "does not start with 'auto_router/'" in str(exc_info.value.message)
assert fake.tx_obj.raw_calls == []
assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0
assert fake.litellm_proxymodeltable.update.await_count == 0
@pytest.mark.asyncio
async def test_model_less_legacy_update_rejects_router_config_on_a_regular_model(self) -> None:
"""The legacy update endpoint enforces the same boundary before its row write or slot."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import update_model
from litellm.types.router import ModelInfo, updateLiteLLMParams
model_id = "regular-model"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
regular = Deployment(
model_name="regular-model",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"),
model_info={"id": model_id},
)
existing_row = MagicMock()
existing_row.model_dump.return_value = regular.model_dump()
existing_row.litellm_params = regular.litellm_params.model_dump()
fake = self._FakeDb([], existing_row=existing_row)
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: config-based lookup must be absent to drive the stored-row branch
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reaches its DB-write branch only with this process setting
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: authorization branch reads the proxy-wide premium flag
patch( # test-quality-ok: endpoint must reject before database authorization needs a live store
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),
),
):
with pytest.raises(ProxyException) as exc_info:
await update_model(
model_params=updateDeployment(
litellm_params=updateLiteLLMParams(complexity_router_config=self._CUSTOM_TIERS),
model_info=ModelInfo(id=model_id),
),
user_api_key_dict=admin,
)
assert exc_info.value.code == "400"
assert "does not start with 'auto_router/'" in str(exc_info.value.message)
assert fake.tx_obj.raw_calls == []
assert fake.tx_obj.litellm_proxymodeltable.update.await_count == 0
assert fake.litellm_proxymodeltable.update.await_count == 0
@pytest.mark.asyncio
async def test_patch_model_refuses_switching_another_router_to_heuristic_v2(self) -> None:
"""patch_model relays HTTPException as-is, so the license refusal reaches the client as a plain 403."""
@ -4591,14 +4801,14 @@ class TestStrategyRouterWriteValidation:
model_id = "other-id"
admin = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
fake = self._FakeDb(db_held=1)
fake = self._FakeDb(["auto_router/complexity_router"])
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # test-quality-ok: the write must be refused before this DB step runs
"litellm.proxy.management_endpoints.model_management_endpoints.get_db_model",
new=AsyncMock(return_value=self._db_complexity_router(model_id)),
@ -4643,14 +4853,14 @@ class TestStrategyRouterWriteValidation:
"model_info": {"id": model_id},
}
existing_row.litellm_params = existing_row.model_dump.return_value["litellm_params"]
fake = self._FakeDb(db_held=1, existing_row=existing_row)
fake = self._FakeDb(["auto_router/complexity_router"], existing_row=existing_row)
with (
patch("litellm.proxy.proxy_server.prisma_client", fake), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1), # test-quality-ok: endpoint reads proxy server globals with no injection seam
patch( # test-quality-ok: prior auth check needs a live DB; only the license limit is under test
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
new=AsyncMock(return_value=None),

14
tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py
View file

@ -1181,3 +1181,17 @@ async def test_find_member_if_email_missing_row_raises_documented_400():
"non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
)
}
def test_v2_update_organization_is_in_openapi_schema():
"""PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec."""
from fastapi import FastAPI
from litellm.proxy.management_endpoints.organization_endpoints import router
app = FastAPI()
app.include_router(router)
v2_path = app.openapi()["paths"]["/v2/organization/{organization_id}"]
assert v2_path["patch"]["tags"] == ["organization management"]
assert "OrganizationUpdateRequestV2" in json.dumps(v2_path["patch"]["requestBody"])

91
tests/test_litellm/proxy/proxy_server/test_proxy_config.py
View file

@ -28,7 +28,7 @@ from litellm.proxy.proxy_server import (
resolve_routing_plugins,
validate_deployment_complexity_router_placement,
validate_deployment_max_agentic_loops,
validate_heuristic_v2_router_limit,
validate_auto_router_capability_limits,
)
from .conftest import normalize
@ -204,13 +204,71 @@ def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") ->
}
def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() -> None:
def _custom_tier_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
},
}
def _operator_examples_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_examples": '- "reset my password" -> SIMPLE',
},
},
}
def _custom_prompt_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
},
},
}
@pytest.mark.parametrize(
"over_limit_rows,subject",
[
([_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], "heuristic_v2"),
([_custom_tier_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "tier_definitions"),
([_custom_prompt_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
([_custom_tier_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
([_operator_examples_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "operator-written classifier prompt"),
],
)
def test_validate_auto_router_capability_limits_refuses_to_start_over_the_limit(
over_limit_rows: list[dict[str, object]], subject: str
) -> None:
"""Same reason as the two validators above: the proxy router swallows registration errors, so
an over-limit config.yaml must fail here instead of booting with a silently missing router."""
with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info:
validate_heuristic_v2_router_limit(
[_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], limit=1
)
validate_auto_router_capability_limits(over_limit_rows, limit=1)
assert subject in str(exc_info.value)
assert "'auto_router' feature lifts the limit" in str(exc_info.value)
@ -220,12 +278,15 @@ def test_validate_heuristic_v2_router_limit_refuses_to_start_over_the_limit() ->
([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None),
([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1),
([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1),
([_custom_tier_row("a"), _custom_tier_row("b")], None),
([_custom_tier_row("a"), _heuristic_v2_row("b")], 1),
],
)
def test_validate_heuristic_v2_router_limit_leaves_configs_within_the_limit_alone(
def test_validate_auto_router_capability_limits_leaves_configs_within_the_limit_alone(
model_list: list[dict[str, object]], limit: int | None
) -> None:
assert validate_heuristic_v2_router_limit(model_list, limit=limit) is None
"""The last case is the separate-ceiling invariant: one router of each capability fits under a limit of one."""
assert validate_auto_router_capability_limits(model_list, limit=limit) is None
_TWO_HEURISTIC_V2_ROUTERS_YAML = (
@ -247,7 +308,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = (
" classifier_type: heuristic_v2\n"
" tiers: {SIMPLE: gpt-4o-mini}\n"
"router_settings:\n"
" heuristic_v2_router_limit: 99\n"
" auto_router_capability_limit: 99\n"
)
@ -256,7 +317,7 @@ _TWO_HEURISTIC_V2_ROUTERS_YAML = (
async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_license_only(
tmp_path, monkeypatch, license_limit: int | None
) -> None:
"""`router_settings.heuristic_v2_router_limit` is managed outside config.yaml: an operator
"""`router_settings.auto_router_capability_limit` is managed outside config.yaml: an operator
cannot grant the entitlement by editing the config, and a licensed proxy boots both routers."""
f = tmp_path / "c.yaml"
f.write_text(_TWO_HEURISTIC_V2_ROUTERS_YAML)
@ -264,15 +325,15 @@ async def test_ProxyConfig_load_config_takes_the_heuristic_v2_limit_from_the_lic
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(
"litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: license_limit
"litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: license_limit
)
if license_limit is None:
router, _model_list, _general_settings = await ProxyConfig().load_config(
router=None, config_file_path=str(f)
)
assert router.heuristic_v2_router_limit is not None
assert router.heuristic_v2_router_limit() is None
assert router.auto_router_capability_limit is not None
assert router.auto_router_capability_limit() is None
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
return
@ -296,12 +357,12 @@ async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_b
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.heuristic_v2_router_limit", lambda: 1)
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1)
router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert router.heuristic_v2_router_limit is not None
assert router.heuristic_v2_router_limit() == 1
assert router.auto_router_capability_limit is not None
assert router.auto_router_capability_limit() == 1
assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"})
assert router.upsert_deployment(db_row) is None

89
tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
View file

@ -452,3 +452,92 @@ async def test_model_info_v2_query_sentinel_does_not_filter(monkeypatch, mixed_a
)
assert "tri-tier-router" in [m["model_name"] for m in resp["data"]]
# ---------------------------------------------------------------------------
# GET /v2/model/info?access_group / ?wildcard_only
# ---------------------------------------------------------------------------
@pytest.fixture
def access_group_router(monkeypatch):
"""Router with one sales-team deployment, one wildcard sales-team deployment and one ungrouped one."""
model_list = [
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
"model_info": {"id": "sales-1", "db_model": False, "access_groups": ["sales-team"]},
},
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*"},
"model_info": {"id": "sales-wildcard", "db_model": False, "access_groups": ["sales-team", "eng"]},
},
{
"model_name": "claude-opus",
"litellm_params": {"model": "anthropic/claude-opus-4-6"},
"model_info": {"id": "plain-1", "db_model": False},
},
]
from unittest.mock import AsyncMock
router = MagicMock()
router.model_list = model_list
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", model_list)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "user_model", None)
monkeypatch.setattr(proxy_server.proxy_config, "get_config", AsyncMock(return_value={}))
monkeypatch.setattr(
proxy_server,
"_apply_search_filter_to_models",
AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))),
)
monkeypatch.setattr(proxy_server, "_enrich_model_info_with_litellm_data", lambda model, **kw: model)
import litellm.proxy.agent_endpoints.model_list_helpers as mlh
monkeypatch.setattr(mlh, "append_agents_to_model_info", AsyncMock(side_effect=lambda models, **kw: models))
yield router
def test_v2_model_info_without_new_filters_returns_everything(client, auth_as, access_group_router):
with auth_as():
response = client.get("/v2/model/info")
payload = response.json()
assert payload["total_count"] == 3
assert len(payload["data"]) == 3
def test_v2_model_info_access_group_filters_rows_and_total(client, auth_as, access_group_router):
"""The table pages off total_count, so the filter must shrink the total, not only the page."""
with auth_as():
response = client.get("/v2/model/info", params={"access_group": "sales-team"})
payload = response.json()
assert _model_names(payload) == ["gpt-4o-mini", "openai/*"]
assert payload["total_count"] == 2
def test_v2_model_info_unknown_access_group_is_empty(client, auth_as, access_group_router):
with auth_as():
response = client.get("/v2/model/info", params={"access_group": "nobody"})
payload = response.json()
assert payload["data"] == []
assert payload["total_count"] == 0
def test_v2_model_info_wildcard_only_filters_rows_and_total(client, auth_as, access_group_router):
with auth_as():
response = client.get("/v2/model/info", params={"wildcard_only": "true"})
payload = response.json()
assert _model_names(payload) == ["openai/*"]
assert payload["total_count"] == 1
def test_v2_model_info_access_group_paginates_over_the_filtered_set(client, auth_as, access_group_router):
with auth_as():
response = client.get("/v2/model/info", params={"access_group": "sales-team", "page": 2, "size": 1})
payload = response.json()
assert _model_names(payload) == ["openai/*"]
assert payload["total_count"] == 2
assert payload["total_pages"] == 2

216
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -9521,6 +9521,222 @@ def test_realtime_websocket_route_aliases_registered():
)
def _lit6973_fake_realtime_ws() -> MagicMock:
ws = MagicMock()
ws.headers = {}
ws.scope = {"headers": [], "type": "websocket"}
ws.url = "ws://testserver/v1/realtime"
ws.accept = AsyncMock()
ws.send_text = AsyncMock()
ws.close = AsyncMock()
return ws
async def _lit6973_drive_realtime_session(
reservation: dict,
*,
backend_logged_success: bool,
phase_one_exit: str | None = None,
websocket: MagicMock | None = None,
) -> MagicMock:
"""Drive realtime_websocket_endpoint through one of its reservation-settling exits.
phase_one_exit picks a rejection before the relay: "model_access" makes the
key/model check raise ProxyException, "pre_call" makes pre-call processing
(rate limits, guardrails) raise. Neither reaches route_request, so no success
log can own the reservation and the endpoint has to release it on that exit.
route_request resolves normally in both cases: the relay owns the session
once route_request returns. A successful session enqueues its success cost
callback and stamps REALTIME_SESSION_SUCCESS_LOGGED_KEY on the shared logging
object; a refused one does neither. The endpoint keys its reservation cleanup
off that stamp, so backend_logged_success reproduces both branches. The fake
logging object carries a real model_call_details dict so the stamp is
observable, and the reservation has empty entries so the real release touches
no counter store."""
from litellm.litellm_core_utils.realtime_streaming import REALTIME_SESSION_SUCCESS_LOGGED_KEY
from litellm.proxy import proxy_server as ps
user_api_key_dict: Final = UserAPIKeyAuth(api_key="sk-test", token="hashed-token")
user_api_key_dict.budget_reservation = reservation
logging_obj: Final = MagicMock()
logging_obj.model_call_details = {}
async def fake_llm_call() -> None:
if backend_logged_success:
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
from litellm.proxy._types import ProxyException
model_access_error: Final = (
ProxyException(message="key cannot access model", type="auth_error", param="model", code=401)
if phase_one_exit == "model_access"
else None
)
pre_call_error: Final = Exception("Rate limit exceeded") if phase_one_exit == "pre_call" else None
pre_call: Final = AsyncMock(
side_effect=pre_call_error, return_value=({"model": "vertex_ai/gemini-live-2.5-flash"}, logging_obj)
)
ws: Final = websocket if websocket is not None else _lit6973_fake_realtime_ws()
can_call = patch.object(ps, "can_key_call_resolved_model", new=AsyncMock(side_effect=model_access_error)) # test-quality-ok: no HTTP boundary; fakes in-process auth to reach the exit under test
pre = patch.object(ps.ProxyBaseLLMRequestProcessing, "common_processing_pre_call_logic", new=pre_call) # test-quality-ok: fakes phase-1 wiring; assertion checks observable reservation state
route = patch.object(ps, "route_request", new=AsyncMock(return_value=fake_llm_call())) # test-quality-ok: fakes the relay whose success/refusal outcome the endpoint reads off the logging object
with can_call, pre, route:
await ps.realtime_websocket_endpoint(
websocket=ws,
model="vertex_ai/gemini-live-2.5-flash",
intent=None,
guardrails=None,
user_api_key_dict=user_api_key_dict,
)
return ws
@pytest.mark.asyncio
async def test_refused_realtime_session_releases_the_budget_reservation():
"""LIT-6973: a refused realtime session enqueues no success cost callback, so
the pre-call reservation would stay open and pin the key/team/user spend
counters, locking the key after a couple of refusals. The endpoint sees no
success stamp and reconciles it: the reservation ends up finalized."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
await _lit6973_drive_realtime_session(reservation, backend_logged_success=False)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_realtime_session_rejected_in_pre_call_releases_the_budget_reservation():
"""A rate-limit or guardrail rejection happens before route_request, so the
relay never runs and no success log can own the reservation. The endpoint
must release it on that exit too, or the key stays pinned at the reserved
amount and its next requests 429 with budget_exceeded while /key/info shows
spend 0 (reproduced live with rpm_limit=1). The client still gets the
pre-call error event and the 1011 close it got before."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
ws: Final = await _lit6973_drive_realtime_session(
reservation, backend_logged_success=False, phase_one_exit="pre_call"
)
assert reservation["finalized"] is True
assert json.loads(ws.send_text.await_args.args[0])["error"]["message"] == "Rate limit exceeded"
ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error")
@pytest.mark.asyncio
async def test_realtime_session_denied_model_access_releases_the_budget_reservation():
"""The key/model access check rejects before the socket is even accepted;
that exit skipped the release as well, pinning the reservation."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
ws: Final = await _lit6973_drive_realtime_session(
reservation, backend_logged_success=False, phase_one_exit="model_access"
)
assert reservation["finalized"] is True
ws.close.assert_awaited_once_with(code=1008, reason="key cannot access model")
@pytest.mark.asyncio
async def test_rejected_realtime_session_closes_the_client_before_releasing_the_reservation():
"""The counter release can block on a slow or unreachable store, and a
rejected client must not sit behind it: the relay's own failure path closes
the client first and releases in its finally, so the pre-relay rejection
has to close first as well. The fake close checks the reservation is still
open when the client is closed."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
ws: Final = _lit6973_fake_realtime_ws()
async def close_while_reservation_is_still_open(**_: object) -> None:
assert reservation["finalized"] is False, "client was closed only after the reservation release"
ws.close = AsyncMock(side_effect=close_while_reservation_is_still_open)
await _lit6973_drive_realtime_session(
reservation, backend_logged_success=False, phase_one_exit="pre_call", websocket=ws
)
ws.close.assert_awaited_once_with(code=1011, reason="Pre-call error")
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_rejected_realtime_session_releases_the_reservation_when_the_client_is_already_gone():
"""A client that hung up before the rejection makes the close raise; the
reservation must still be released, or the key stays pinned."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
ws: Final = _lit6973_fake_realtime_ws()
ws.close = AsyncMock(side_effect=RuntimeError("client already disconnected"))
with pytest.raises(RuntimeError, match="client already disconnected"):
await _lit6973_drive_realtime_session(
reservation, backend_logged_success=False, phase_one_exit="model_access", websocket=ws
)
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_successful_realtime_session_leaves_the_reservation_for_the_cost_callback():
"""A billable realtime session settles its reservation through the enqueued
success cost callback, not the endpoint. The endpoint must not finalize it in
its finally, or it would reconcile the reservation to zero before the cost
callback applies real spend, so billable sessions stop counting against budget.
With the success stamp present, the endpoint leaves the reservation untouched."""
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
await _lit6973_drive_realtime_session(reservation, backend_logged_success=True)
assert reservation["finalized"] is False
@pytest.mark.asyncio
async def test_release_or_invalidate_falls_back_to_invalidating_the_counters():
"""If releasing the reservation itself fails (e.g. the counter store is down),
the reserved counters must be invalidated directly so the estimate does not
stay pinned, and the reservation is finalized so nothing reprocesses it."""
from litellm.proxy import proxy_server as ps
from litellm.proxy.spend_tracking import budget_reservation as br
reservation: Final = {
"reserved_cost": 0.55,
"input_cost": 0.0,
"finalized": False,
"entries": [{"counter_key": "spend:key:hashed-token"}],
}
invalidated: Final[list[str]] = []
async def _record(counter_key: str) -> None:
invalidated.append(counter_key)
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the failure branch; assertion observes which counter key got invalidated
sink = patch.object(ps, "_invalidate_spend_counter", new=_record) # test-quality-ok: fakes the counter-store sink so the invalidated key is observable
with failing_release, sink:
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
assert invalidated == ["spend:key:hashed-token"]
assert reservation["finalized"] is True
@pytest.mark.asyncio
async def test_release_or_invalidate_finalizes_even_when_the_invalidate_fallback_fails():
"""Both counter-store calls failing must not raise out of the realtime
endpoint's finally (it would mask the session's own outcome) and must still
stamp finalized so nothing retries the same reservation."""
from litellm.proxy.spend_tracking import budget_reservation as br
reservation: Final = {"reserved_cost": 0.55, "input_cost": 0.0, "finalized": False, "entries": []}
failing_release = patch.object(br, "release_budget_reservation", new=AsyncMock(side_effect=RuntimeError("counter store down"))) # test-quality-ok: forces the fallback branch
failing_invalidate = patch.object(br, "invalidate_budget_reservation_counters", new=AsyncMock(side_effect=RuntimeError("still down"))) # test-quality-ok: forces the fallback itself to fail
with failing_release, failing_invalidate:
await br.release_or_invalidate_budget_reservation(budget_reservation=reservation)
assert reservation["finalized"] is True
class TestTransformRequestBannedParams:
"""
/utils/transform_request applies the same banned-param check as LLM endpoints.

109
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -473,6 +473,110 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage(
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
def _deadlock_error() -> Exception:
return _data_error(
'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, '
'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })'
)
@pytest.mark.asyncio
async def test_update_spend_logs_retries_deadlock_and_keeps_every_row(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A 40P01 deadlock aborts the whole insert, so the same rows succeed on replay.
Before the fix the deadlock surfaced as a plain ``DataError`` and went through
poison-row isolation, which bisected the batch and dropped every row the
deadlock happened to hit as if Postgres had rejected it.
"""
async def _fake_sleep(_: float) -> None:
return None
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
create_many = AsyncMock(side_effect=[_deadlock_error(), _deadlock_error(), None])
mock_prisma_client.db.litellm_spendlogs.create_many = create_many
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
mock_prisma_client.spend_log_transactions = []
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=2,
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging,
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
)
attempts = tuple(
tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list
)
assert attempts == (("a", "b"), ("a", "b"), ("a", "b"))
assert mock_prisma_client.spend_log_transactions == []
@pytest.mark.asyncio
async def test_update_spend_logs_requeues_batch_once_deadlock_retries_exhaust(
mock_prisma_client: Any, make_spend_log_row: Any, monkeypatch: pytest.MonkeyPatch
) -> None:
"""If every retry deadlocks, the batch goes back to the head of the queue for
the next flush instead of being dropped.
"""
async def _fake_sleep(_: float) -> None:
return None
monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_deadlock_error())
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
with pytest.raises(type(_deadlock_error())):
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=1,
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging,
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
)
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 2
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
@pytest.mark.asyncio
async def test_update_spend_logs_requeues_batch_on_non_transport_db_error(
mock_prisma_client: Any, make_spend_log_row: Any
) -> None:
"""A DB error that is neither transport nor deadlock (here P2021, the table is
gone mid-migration) is not retried in place, but the dequeued batch must not
be lost either: it goes back to the head of the queue so it lands once the
DB is healthy again.
"""
from prisma.errors import TableNotFoundError
err = TableNotFoundError(
{"user_facing_error": {"error_code": "P2021", "message": "The table does not exist", "meta": {}}}
)
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=err)
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="c")]
with pytest.raises(TableNotFoundError):
await ProxyUpdateSpend.update_spend_logs(
n_retry_times=2,
prisma_client=mock_prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging,
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
)
assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 1
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["a", "b", "c"]
@pytest.mark.asyncio
async def test_requeue_after_outage_drops_oldest_logs_past_the_byte_budget(
mock_prisma_client: Any, make_spend_log_row: Any
@ -549,8 +653,9 @@ async def test_flush_returns_the_bytes_it_took_off_the_queue(mock_prisma_client:
async def test_update_spend_logs_does_not_requeue_non_transport_failures(
mock_prisma_client: Any, make_spend_log_row: Any
) -> None:
"""Only transport failures are worth replaying. A rejection the DB will keep
rejecting must not be requeued, or it would wedge the queue forever.
"""Only DB failures are worth replaying. A row the proxy itself cannot
serialize would fail the same way on every flush, so requeueing it would
wedge the head of the queue forever.
"""
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=ValueError("bad payload"))
proxy_logging = MagicMock()

254
tests/test_litellm/responses/test_responses_supported_endpoints_passthrough.py Normal file
View file

@ -0,0 +1,254 @@
"""
A deployment with `model_info.supported_endpoints` containing `/v1/responses` forwards
`/v1/responses` natively to `{api_base}/responses`. Without it, generic OpenAI-compatible
providers such as `custom_openai` keep bridging through `/v1/chat/completions`.
"""
import json
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
import respx
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig
from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig
from litellm.responses.main import _resolve_responses_api_provider_config
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import ModelResponse
API_BASE = "https://backend.example/v1"
RESPONSES_URL = f"{API_BASE}/responses"
CHAT_URL = f"{API_BASE}/chat/completions"
OPT_IN = {"supported_endpoints": ["/v1/chat/completions", "/v1/responses"]}
RESPONSES_BODY = {
"id": "resp_native",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "my-model",
"output": [
{
"type": "message",
"id": "msg_1",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "native", "annotations": []}],
}
],
"parallel_tool_calls": True,
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
}
CHAT_BODY = {
"id": "chatcmpl_bridged",
"object": "chat.completion",
"created": 1741476542,
"model": "my-model",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "bridged"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
SSE_BODY = (
"event: response.created\n"
f"data: {json.dumps({'type': 'response.created', 'response': RESPONSES_BODY})}\n\n"
"event: response.completed\n"
f"data: {json.dumps({'type': 'response.completed', 'response': RESPONSES_BODY})}\n\n"
)
def _mock_backend(router: respx.MockRouter) -> tuple[respx.Route, respx.Route]:
responses_route = router.post(RESPONSES_URL).mock(return_value=httpx.Response(200, json=RESPONSES_BODY))
chat_route = router.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY))
return responses_route, chat_route
SWAPPED_MODEL = "deepseek/deepseek-chat"
SWAPPED_API_BASE = "https://api.deepseek.com/beta"
def _prompt_manager_swapping_to(model: str) -> MagicMock:
"""A logging object whose prompt hook rewrites the request's model, as a prompt manager does."""
prompt_return = (model, [{"role": "user", "content": "hi"}], {})
logging_obj = MagicMock()
logging_obj.__class__ = LiteLLMLoggingObj
logging_obj.should_run_prompt_management_hooks.return_value = True
logging_obj.get_chat_completion_prompt.return_value = prompt_return
logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return)
logging_obj.model_call_details = {}
return logging_obj
def _mock_swap_targets(router: respx.MockRouter, monkeypatch) -> tuple[respx.Route, respx.Route]:
"""The swapped provider's chat endpoint, plus the `/responses` it does not serve but a stale
opt-in would send to."""
monkeypatch.setenv("DEEPSEEK_API_KEY", "sk-deepseek")
swapped_chat_route = router.post(f"{SWAPPED_API_BASE}/chat/completions").mock(
return_value=httpx.Response(200, json=CHAT_BODY)
)
stale_responses_route = router.post(f"{SWAPPED_API_BASE}/responses").mock(
return_value=httpx.Response(200, json=RESPONSES_BODY)
)
return swapped_chat_route, stale_responses_route
@pytest.fixture(autouse=True)
def _respx_interceptable_httpx_client(monkeypatch):
monkeypatch.setattr(litellm, "num_retries", 0)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
yield
litellm.in_memory_llm_clients_cache.flush_cache()
@pytest.mark.parametrize(
"model_info, expected_type",
[
(OPT_IN, OpenAILikeResponsesConfig),
({"supported_endpoints": ["/v1/chat/completions"]}, type(None)),
({}, type(None)),
(None, type(None)),
("/v1/responses", type(None)),
],
)
def test_resolver_opt_in_gates_openai_like_config(model_info, expected_type):
config = _resolve_responses_api_provider_config("my-model", "custom_openai", model_info)
assert type(config) is expected_type
def test_resolver_keeps_native_provider_config():
"""`openai/` already routes /v1/responses natively; the opt-in must not swap its config."""
config = _resolve_responses_api_provider_config("gpt-4.1", "openai", OPT_IN)
assert type(config) is OpenAIResponsesAPIConfig
@respx.mock
async def test_opt_in_forwards_responses_natively():
responses_route, chat_route = _mock_backend(respx.mock)
result = await litellm.aresponses(
model="custom_openai/my-model",
input="hi",
api_base=API_BASE,
api_key="sk-backend",
model_info=OPT_IN,
)
assert responses_route.call_count == 1
assert chat_route.call_count == 0
request = responses_route.calls.last.request
assert request.headers["authorization"] == "Bearer sk-backend"
assert json.loads(request.content)["input"] == "hi"
assert isinstance(result, ResponsesAPIResponse)
assert result.output[0].content[0].text == "native"
@respx.mock
async def test_opt_in_forwards_streaming_responses_natively(monkeypatch):
"""The router registers each deployment in `litellm.model_cost`; an unregistered model is
treated as non-streaming and would be faked, so mirror that registration here."""
monkeypatch.setitem(litellm.model_cost, "custom_openai/my-model", {"litellm_provider": "custom_openai"})
responses_route = respx.post(RESPONSES_URL).mock(
return_value=httpx.Response(200, text=SSE_BODY, headers={"content-type": "text/event-stream"})
)
chat_route = respx.post(CHAT_URL).mock(return_value=httpx.Response(200, json=CHAT_BODY))
stream = await litellm.aresponses(
model="custom_openai/my-model",
input="hi",
stream=True,
api_base=API_BASE,
api_key="sk-backend",
model_info=OPT_IN,
)
events = [event async for event in stream]
assert responses_route.call_count == 1
assert chat_route.call_count == 0
assert json.loads(responses_route.calls.last.request.content)["stream"] is True
assert [event.type for event in events] == ["response.created", "response.completed"]
@respx.mock
async def test_without_opt_in_still_bridges_through_chat_completions():
responses_route, chat_route = _mock_backend(respx.mock)
result = await litellm.aresponses(
model="custom_openai/my-model",
input="hi",
api_base=API_BASE,
api_key="sk-backend",
model_info={"supported_endpoints": ["/v1/chat/completions"]},
)
assert chat_route.call_count == 1
assert responses_route.call_count == 0
assert isinstance(result, ResponsesAPIResponse)
assert result.output[0].content[0].text == "bridged"
@respx.mock
async def test_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch):
"""When a prompt manager moves the request to another provider, the original deployment's
`supported_endpoints` no longer describes the upstream, so the swapped provider bridges."""
swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch)
result = await litellm.aresponses(
model="custom_openai/my-model",
input="hi",
prompt_id="p1",
litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL),
model_info=OPT_IN,
)
assert swapped_chat_route.call_count == 1
assert stale_responses_route.call_count == 0
assert isinstance(result, ResponsesAPIResponse)
assert result.output[0].content[0].text == "bridged"
@respx.mock
def test_sync_prompt_swap_to_other_provider_drops_deployment_opt_in(monkeypatch):
swapped_chat_route, stale_responses_route = _mock_swap_targets(respx.mock, monkeypatch)
result = litellm.responses(
model="custom_openai/my-model",
input="hi",
prompt_id="p1",
litellm_logging_obj=_prompt_manager_swapping_to(SWAPPED_MODEL),
model_info=OPT_IN,
)
assert swapped_chat_route.call_count == 1
assert stale_responses_route.call_count == 0
assert isinstance(result, ResponsesAPIResponse)
assert result.output[0].content[0].text == "bridged"
@respx.mock
async def test_mode_responses_chat_completion_reaches_native_responses(monkeypatch):
"""A `mode: responses` deployment bridges chat completions into the Responses API; with
the opt-in that inner call must reach `{api_base}/responses` instead of bouncing back
to `/chat/completions`."""
responses_route, chat_route = _mock_backend(respx.mock)
monkeypatch.setitem(
litellm.model_cost,
"custom_openai/my-model",
{"mode": "responses", "litellm_provider": "custom_openai"},
)
result = await litellm.acompletion(
model="custom_openai/my-model",
messages=[{"role": "user", "content": "hi"}],
api_base=API_BASE,
api_key="sk-backend",
model_info={"mode": "responses", **OPT_IN},
)
assert responses_route.call_count == 1
assert chat_route.call_count == 0
assert isinstance(result, ModelResponse)
assert result.choices[0].message.content == "native"

232
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -15,6 +15,12 @@ from pydantic import ValidationError
import litellm
from litellm import Router
from litellm.router_utils.auto_router_model_naming import (
CUSTOMIZATION_CAPABILITY,
GATED_AUTO_ROUTER_CAPABILITIES,
HEURISTIC_V2_CAPABILITY,
count_capability_routers,
)
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY
@ -46,7 +52,6 @@ from litellm.router_strategy.complexity_router.tier_predictor import (
TierGlobalStatistic,
TrainedTierArtifact,
)
from litellm.router_utils.auto_router_model_naming import count_heuristic_v2_routers
from litellm.types.router import (
Deployment,
LiteLLM_Params,
@ -1124,7 +1129,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-b", "id-b", "heuristic_v2"),
self._router_row("v1-c", "id-c", "heuristic"),
],
heuristic_v2_router_limit=lambda: 1,
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
@ -1139,7 +1144,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._router_row("v2-b", "id-b", "heuristic_v2"),
],
heuristic_v2_router_limit=lambda: 1,
auto_router_capability_limit=lambda: 1,
)
def test_heuristic_v2_limit_is_resolved_on_every_registration(self) -> None:
@ -1152,14 +1157,14 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._router_row("v2-b", "id-b", "heuristic_v2"),
],
heuristic_v2_router_limit=lambda: limits["value"],
auto_router_capability_limit=lambda: limits["value"],
ignore_invalid_deployments=True,
)
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
assert router.heuristic_v2_router_limit_violation() is None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None
limits["value"] = 1
assert router.heuristic_v2_router_limit_violation() is not None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
assert router.upsert_deployment(Deployment(**self._router_row("v2-c", "id-c", "heuristic_v2"))) is None
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
@ -1174,7 +1179,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._router_row("v2-b", "id-b", "heuristic_v2"),
],
heuristic_v2_router_limit=lambda: limits["value"],
auto_router_capability_limit=lambda: limits["value"],
ignore_invalid_deployments=True,
)
limits["value"] = 1
@ -1195,7 +1200,7 @@ class TestRouterComplexityDeploymentMethods:
assert router.upsert_deployment(Deployment(**db_row)) is not None
assert sorted(str(row["model_name"]) for row in router.config_deployments()) == ["gpt-4o-mini", "v2-a"]
assert count_heuristic_v2_routers(router.config_deployments()) == 1
assert count_capability_routers(router.config_deployments(), capability=HEURISTIC_V2_CAPABILITY) == 1
def test_failed_edit_of_a_live_v2_router_rolls_back_without_the_ceiling(self) -> None:
"""A rollback after a failed upsert re-admits state that was already serving, so it must not be
@ -1208,7 +1213,7 @@ class TestRouterComplexityDeploymentMethods:
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._router_row("v2-b", "id-b", "heuristic_v2"),
],
heuristic_v2_router_limit=lambda: limits["value"],
auto_router_capability_limit=lambda: limits["value"],
ignore_invalid_deployments=True,
)
limits["value"] = 1
@ -1220,7 +1225,7 @@ class TestRouterComplexityDeploymentMethods:
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
live = router.get_deployment(model_id="id-a")
assert live is not None and live.litellm_params.complexity_router_config["classifier_type"] == "heuristic_v2"
assert router.heuristic_v2_router_limit_violation() is not None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
def test_heuristic_v2_routers_are_unlimited_by_default(self) -> None:
router = Router(
@ -1232,18 +1237,18 @@ class TestRouterComplexityDeploymentMethods:
)
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
assert router.heuristic_v2_router_limit_violation() is None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is None
def test_heuristic_v2_router_limit_violation_frees_the_slot_of_the_router_being_edited(self) -> None:
def test_auto_router_capability_violation_frees_the_slot_of_the_router_being_edited(self) -> None:
"""A DB reload upserts the existing heuristic_v2 router again; that edit must keep its own slot
while a different deployment switching to heuristic_v2 is refused."""
router = Router(
model_list=[self._POOL, self._router_row("v2-a", "id-a", "heuristic_v2")],
heuristic_v2_router_limit=lambda: 1,
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
assert router.heuristic_v2_router_limit_violation() is not None
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
edited = self._router_row("v2-a-renamed", "id-a", "heuristic_v2")
assert router.upsert_deployment(Deployment(**edited)) is not None
@ -1254,6 +1259,205 @@ class TestRouterComplexityDeploymentMethods:
assert router.upsert_deployment(Deployment(**self._router_row("v1-c", "id-c", "heuristic"))) is not None
assert sorted(router.complexity_routers) == ["v1-c", "v2-a-renamed"]
@staticmethod
def _custom_tier_row(model_name: str, model_id: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tier_definitions": [
{"name": "routine", "description": "routine drafting and lookups"},
{"name": "hard", "description": "multi-step reasoning under tradeoffs"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
},
"model_info": {"id": model_id},
}
@staticmethod
def _custom_prompt_row(model_name: str, model_id: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
},
},
} | {"model_info": {"id": model_id}}
def test_a_second_custom_prompt_router_is_refused_under_the_ceiling(self) -> None:
"""An operator-written classifier system_prompt is metered like the other licensed capabilities."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_prompt_row("prompt-a", "id-a"),
self._custom_prompt_row("prompt-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None:
"""Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no
prompt at all, leaves a router unmetered, so several of them register under a ceiling of one."""
def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]:
llm_config: dict[str, object] = {"model": "gpt-4o-mini"}
if preset is not None:
llm_config["classification_rubric"] = preset
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": llm_config,
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
},
},
"model_info": {"id": model_id},
}
router = Router(
model_list=[
self._POOL,
rubric("default-a", "id-a", None),
rubric("preset-b", "id-b", "agentic"),
rubric("preset-c", "id-c", "chat"),
],
auto_router_capability_limit=lambda: 1,
)
assert sorted(router.complexity_routers) == ["default-a", "preset-b", "preset-c"]
def test_a_second_custom_tier_router_is_refused_under_the_ceiling(self) -> None:
"""Operator-defined tier sets are metered like heuristic_v2: one per proxy without the license."""
with pytest.raises(ValueError, match="tier_definitions"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_tier_row("tiers-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_custom_tier_routers_are_unlimited_with_the_license_feature(self) -> None:
router = Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_tier_row("tiers-b", "id-b"),
],
auto_router_capability_limit=lambda: None,
)
assert sorted(router.complexity_routers) == ["tiers-a", "tiers-b"]
assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is None
def test_each_capability_holds_its_own_slot(self) -> None:
"""heuristic_v2 has its own slot, while custom tiers and custom prompts share one customization
slot: one v2 plus EITHER customization fits, but a second customization of any form is refused."""
router = Router(
model_list=[
self._POOL,
self._router_row("v2-a", "id-a", "heuristic_v2"),
self._custom_tier_row("tiers-a", "id-t"),
],
auto_router_capability_limit=lambda: 1,
ignore_invalid_deployments=True,
)
assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"]
assert router.auto_router_capability_violation(HEURISTIC_V2_CAPABILITY) is not None
assert router.auto_router_capability_violation(CUSTOMIZATION_CAPABILITY) is not None
assert router.upsert_deployment(Deployment(**self._custom_tier_row("tiers-b", "id-t2"))) is None
assert router.upsert_deployment(Deployment(**self._custom_prompt_row("prompt-b", "id-p2"))) is None
assert router.upsert_deployment(Deployment(**self._router_row("v2-b", "id-b", "heuristic_v2"))) is None
assert sorted(router.complexity_routers) == ["tiers-a", "v2-a"]
@staticmethod
def _operator_prompt_row(model_name: str, model_id: str, field: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "gpt-4o-mini",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
field: '- "reset my password" -> SIMPLE',
},
},
"model_info": {"id": model_id},
}
@pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"])
def test_operator_written_prompt_sections_claim_the_customization_slot(self, field: str) -> None:
"""The dashboard prompt editor writes opening instructions and calibration examples as their own
fields on a BUILT-IN tier router, so each must claim the slot on its own."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._operator_prompt_row("prompt-a", "id-a", field),
self._operator_prompt_row("prompt-b", "id-b", field),
],
auto_router_capability_limit=lambda: 1,
)
@pytest.mark.parametrize("field", ["classification_prompt", "classification_examples"])
def test_an_operator_prompt_section_claims_the_slot_held_by_custom_tiers(self, field: str) -> None:
"""Switching the FORM of customization cannot buy a second unlicensed router."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._operator_prompt_row("prompt-b", "id-b", field),
],
auto_router_capability_limit=lambda: 1,
)
def test_a_custom_prompt_claims_the_slot_held_by_custom_tiers(self) -> None:
"""The customization ceiling is shared: changing its form cannot get a second unlicensed router."""
with pytest.raises(ValueError, match="operator-written classifier prompt"):
Router(
model_list=[
self._POOL,
self._custom_tier_row("tiers-a", "id-a"),
self._custom_prompt_row("prompt-b", "id-b"),
],
auto_router_capability_limit=lambda: 1,
)
def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None:
"""tier_labels renames the built-in ladder without defining one, so it stays ungated: two such
routers register under a ceiling of one."""
def labeled(model_name: str, model_id: str) -> dict[str, object]:
row = self._router_row(model_name, model_id, "heuristic")
row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"}
return row
router = Router(
model_list=[self._POOL, labeled("labels-a", "id-a"), labeled("labels-b", "id-b")],
auto_router_capability_limit=lambda: 1,
)
assert sorted(router.complexity_routers) == ["labels-a", "labels-b"]
def test_hybrid_initialization_waits_for_later_pool_deployments(self):
router = Router(
model_list=[

172
tests/test_litellm/router_utils/test_auto_router_model_naming.py
View file

@ -5,9 +5,11 @@ import pytest
from litellm.router_utils.auto_router_model_naming import (
carries_complexity_router_settings,
classify_strategy_router_model,
count_heuristic_v2_routers,
heuristic_v2_limit_violation,
is_heuristic_v2_router,
GATED_AUTO_ROUTER_CAPABILITIES,
capability_limit_violation,
claimed_capability,
count_capability_routers,
gated_capability_of,
strategy_router_dependencies,
validate_complexity_router_config_placement,
validate_complexity_router_config_write,
@ -376,38 +378,122 @@ def test_placement_is_scoped_to_complexity_router_deployments(model, present_fie
assert carries_complexity_router_settings(model, present_fields) is scoped
_HV2_CONFIG: Mapping[str, object] = {"classifier_type": "heuristic_v2"}
_CUSTOM_TIER_CONFIG: Mapping[str, object] = {
"classifier_type": "llm",
"tier_definitions": [{"name": "routine", "description": "easy"}, {"name": "hard", "description": "hard"}],
}
_CUSTOM_PROMPT_CONFIG: Mapping[str, object] = {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
}
@pytest.mark.parametrize(
"litellm_params,expected",
"config,expected_key",
[
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, True),
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, False),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, False),
({"model": "auto_router/complexity_router"}, False),
({"model": "auto_router/quality_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
({"model": "openai/gpt-4o", "complexity_router_config": {"classifier_type": "heuristic_v2"}}, False),
({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, False),
({}, False),
(_CUSTOM_PROMPT_CONFIG, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": "grade it"}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_examples": '- "x" -> SIMPLE'}, "tier_or_classifier_prompt"),
({"classifier_type": "hybrid", "classification_examples": "- y -> MEDIUM"}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}, "classification_prompt": None, "classification_examples": None}, None),
({"classifier_type": "heuristic", "classification_examples": "- x -> SIMPLE"}, None),
({"classifier_type": "hybrid", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"),
({"classifier_type": "heuristic_first", "classifier_llm_config": {"system_prompt": "p"}}, "tier_or_classifier_prompt"),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "classification_rubric": "chat"}}, None),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m"}}, None),
({"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": None}}, None),
({"classifier_type": "heuristic", "classifier_llm_config": {"system_prompt": "p"}}, None),
({"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}}, "heuristic_v2"),
({"classifier_type": "llm", "classifier_llm_config": "not a mapping"}, None),
],
)
def test_is_heuristic_v2_router(litellm_params: Mapping[str, object], expected: bool) -> None:
"""Only a complexity router whose config selects heuristic_v2 counts toward the license limit."""
assert is_heuristic_v2_router(litellm_params) is expected
def test_custom_classifier_prompt_capability(config: Mapping[str, object], expected_key: str | None) -> None:
"""Every operator-written part of the classifier prompt claims the customization slot: a whole
replacement system_prompt, replacement opening instructions (classification_prompt), or replacement
calibration examples (classification_examples).
A shipped rubric preset stays free, and the heuristic scorers never read system_prompt, so a
value sitting on one is inert and claims nothing (heuristic_v2 still claims its own capability).
"""
claimed = claimed_capability(config)
assert (None if claimed is None else claimed.key) == expected_key
def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_ones() -> None:
v2 = {"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic_v2"}}
@pytest.mark.parametrize(
"model,expected",
[
("auto_router/complexity_router", True),
("auto_router/complexity_router-eu", True),
("auto_router/semantic_router", False),
("auto_router/adaptive_router", False),
("auto_router/quality_router", False),
("openai/gpt-4o", False),
(None, False),
],
)
def test_is_complexity_router_model(model: str | None, expected: bool) -> None:
from litellm.router_utils.auto_router_model_naming import is_complexity_router_model
assert is_complexity_router_model(model) is expected
@pytest.mark.parametrize(
"litellm_params,expected_key",
[
({"model": "auto_router/complexity_router", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _HV2_CONFIG}, "heuristic_v2"),
({"model": "auto_router/complexity_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"),
({"model": "auto_router/complexity_router-eu", "complexity_router_config": _CUSTOM_TIER_CONFIG}, "tier_or_classifier_prompt"),
({"model": "auto_router/complexity_router", "complexity_router_config": {"classifier_type": "heuristic"}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tiers": {"SIMPLE": "a"}}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_definitions": None}}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": {"tier_labels": {"SIMPLE": "Cheap"}}}, None),
({"model": "auto_router/complexity_router"}, None),
({"model": "auto_router/quality_router", "complexity_router_config": _HV2_CONFIG}, None),
({"model": "auto_router/quality_router", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None),
({"model": "openai/gpt-4o", "complexity_router_config": _HV2_CONFIG}, None),
({"model": "openai/gpt-4o", "complexity_router_config": _CUSTOM_TIER_CONFIG}, None),
({"model": "auto_router/complexity_router", "complexity_router_config": "heuristic_v2"}, None),
({}, None),
],
)
def test_gated_capability_of(litellm_params: Mapping[str, object], expected_key: str | None) -> None:
"""Only a complexity router claiming a licensed capability counts toward that capability's limit.
Renaming the built-in tiers through tier_labels is not a custom tier set, so it stays ungated.
"""
capability = gated_capability_of(litellm_params)
assert (None if capability is None else capability.key) == expected_key
@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key)
def test_count_capability_routers_counts_only_its_own_capability(capability) -> None:
"""Each capability has its own ceiling, so a router claiming the sibling capability never counts,
while a custom tier set and a custom classifier prompt count into the SAME customization slot."""
def row(name: str, config: Mapping[str, object] | None) -> Mapping[str, object]:
params = {"model": "auto_router/complexity_router"} | ({} if config is None else {"complexity_router_config": config})
return {"model_name": name, "litellm_params": params}
by_key = {
"heuristic_v2": (_HV2_CONFIG, _HV2_CONFIG),
"tier_or_classifier_prompt": (_CUSTOM_TIER_CONFIG, _CUSTOM_PROMPT_CONFIG),
}
mine_first, mine_second = by_key[capability.key]
theirs = next(configs[0] for key, configs in by_key.items() if key != capability.key)
rows: list[Mapping[str, object]] = [
{"model_name": "a", "litellm_params": v2},
{"model_name": "b", "litellm_params": {"model": "openai/gpt-4o"}},
{"model_name": "c", "litellm_params": v2},
{"model_name": "d"},
{"model_name": "e", "litellm_params": "not a mapping"},
row("a", mine_first),
row("b", theirs),
{"model_name": "c", "litellm_params": {"model": "openai/gpt-4o"}},
row("d", mine_second),
{"model_name": "e"},
{"model_name": "f", "litellm_params": "not a mapping"},
]
assert count_heuristic_v2_routers(rows) == 2
assert count_heuristic_v2_routers(()) == 0
assert count_capability_routers(rows, capability=capability) == 2
assert count_capability_routers((), capability=capability) == 0
@pytest.mark.parametrize("capability", GATED_AUTO_ROUTER_CAPABILITIES, ids=lambda c: c.key)
@pytest.mark.parametrize(
"held,limit,violates",
[
@ -419,10 +505,42 @@ def test_count_heuristic_v2_routers_reads_model_list_rows_and_ignores_malformed_
(4, 3, True),
],
)
def test_heuristic_v2_limit_violation(held: int, limit: int | None, violates: bool) -> None:
violation = heuristic_v2_limit_violation(held=held, limit=limit)
def test_capability_limit_violation(held: int, limit: int | None, violates: bool, capability) -> None:
violation = capability_limit_violation(capability=capability, held=held, limit=limit)
assert (violation is not None) is violates
if violation is not None:
assert f"At most {limit} auto-router" in violation
assert f"would make {held}" in violation
assert capability.subject in violation
assert capability.remedy in violation
assert "license" not in violation
def test_every_gated_capability_has_a_distinct_predicate_and_sql_spelling() -> None:
"""The in-process and SQL halves of a capability must stay paired, and no two capabilities may collide."""
keys = tuple(capability.key for capability in GATED_AUTO_ROUTER_CAPABILITIES)
assert len(set(keys)) == len(keys)
for capability in GATED_AUTO_ROUTER_CAPABILITIES:
assert "{config}" in capability.sql_config_predicate
assert capability.uses is not None
@pytest.mark.parametrize(
"config",
[
_HV2_CONFIG,
_CUSTOM_TIER_CONFIG,
_CUSTOM_PROMPT_CONFIG,
{"classifier_type": "heuristic"},
{"classifier_type": "heuristic_v2", "classifier_llm_config": {"system_prompt": "p"}},
{"classifier_type": "llm", "classifier_llm_config": {"model": "m", "system_prompt": "p"}, "tier_labels": {"SIMPLE": "Cheap"}},
],
)
def test_capabilities_are_mutually_exclusive_on_one_config(config: Mapping[str, object]) -> None:
"""No config claims two capabilities, which is what lets one lock and one count serve them all.
The config validator is what makes this true and is pinned separately in test_complexity_router:
tier_definitions rejects every heuristic classifier_type and rejects the classifier system_prompt,
and system_prompt only counts for the classifier types heuristic_v2 is not one of.
"""
assert sum(1 for capability in GATED_AUTO_ROUTER_CAPABILITIES if capability.uses(config)) <= 1

32
tests/test_litellm/router_utils/test_pattern_match_deployments.py
View file

@ -2,8 +2,10 @@
from __future__ import annotations
from unittest.mock import Mock
from litellm.router_utils import pattern_match_deployments
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter, PatternUtils
def _wildcard_deployment(model_name: str) -> dict:
@ -76,3 +78,31 @@ def test_get_pattern_still_resolves_unqualified_names(monkeypatch):
router = PatternMatchRouter()
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"]
class _CountingPatternUtils(PatternUtils):
sorted_patterns = staticmethod(Mock(wraps=PatternUtils.sorted_patterns))
def test_route_never_sorts_and_the_most_specific_pattern_still_wins_after_registry_changes():
"""Regression for LIT-6886: the auth layer walks the wildcard registry for every request, so an
unmatched model name (an invalid-model 403) re-sorted every pattern by specificity per request and
a burst of rejections saturated the worker CPU. Lookups must not sort; adding a pattern or removing
a deployment must still leave the most specific pattern winning."""
router = PatternMatchRouter(pattern_utils=_CountingPatternUtils)
router.add_pattern("openai/*", _wildcard_deployment("openai/*"))
router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*"))
router.add_pattern("openai/gpt-*", {"model_name": "openai/gpt-*", "litellm_params": {"model": "azure/gpt-*"}})
sorts_after_setup = _CountingPatternUtils.sorted_patterns.call_count
for _ in range(3):
assert router.route("does-not-exist") is None
assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"]
assert _matched_models(router.route("openai/o3")) == ["openai/o3"]
assert _CountingPatternUtils.sorted_patterns.call_count == sorts_after_setup
router.add_pattern("openai/*", {**_wildcard_deployment("openai/*"), "model_info": {"id": "id-1"}})
assert len(_matched_models(router.route("openai/o3"))) == 2
router.remove_deployment("id-1")
assert _matched_models(router.route("openai/gpt-4o")) == ["azure/gpt-4o"]
assert _matched_models(router.route("openai/o3")) == ["openai/o3"]

273
tests/test_litellm/test_router.py
View file

@ -5,6 +5,7 @@ import json
import logging
import os
import threading
from datetime import datetime
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -20,6 +21,7 @@ import litellm
from litellm import Router
from litellm.exceptions import MidStreamFallbackError
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES,
@ -12938,3 +12940,274 @@ async def test_router_retry_policy_controls_upstream_attempt_count(
await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}])
assert upstream.call_count == expected_upstream_calls
def _make_failure_logging_obj():
return LiteLLMLogging(
model="gpt-5.6",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type="acompletion",
start_time=datetime.now(),
litellm_call_id="lit-6960",
function_id="f",
)
async def _assert_router_failure_logging_is_coordinated(logging_obj, trigger, expected_exception):
"""The sync failure_handler must not start until async_failure_handler has finished on the shared logging_obj."""
events: list[str] = []
sync_done = threading.Event()
async def _async_failure(*args, **kwargs):
events.append("async_start")
await asyncio.sleep(0.05)
events.append("async_end")
def _sync_failure(*args, **kwargs):
events.append("sync_start")
sync_done.set()
with (
patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure),
patch.object(logging_obj, "failure_handler", side_effect=_sync_failure),
patch.object(logging_obj, "_should_run_sync_failure_callbacks_for_async_calls", return_value=True),
):
with pytest.raises(expected_exception):
await trigger()
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
await asyncio.gather(*pending)
assert await asyncio.to_thread(sync_done.wait, 5), "failure_handler never ran"
assert events == ["async_start", "async_end", "sync_start"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"hook_error",
[
litellm.RateLimitError(message="rpm exceeded", llm_provider="openai", model="gpt-5.6"),
RuntimeError("pre call check blew up"),
],
)
async def test_async_routing_strategy_pre_call_checks_failure_logging_is_coordinated(hook_error):
class _RaisingPreCallCheck(CustomLogger):
async def async_pre_call_check(self, deployment, parent_otel_span):
raise hook_error
router = litellm.Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}]
)
deployment = router.model_list[0]
logging_obj = _make_failure_logging_obj()
with patch.object(litellm, "callbacks", [_RaisingPreCallCheck()]): # test-quality-ok: router reads this global
await _assert_router_failure_logging_is_coordinated(
logging_obj,
lambda: router.async_routing_strategy_pre_call_checks(
deployment=deployment, parent_otel_span=None, logging_obj=logging_obj
),
type(hook_error),
)
@pytest.mark.asyncio
async def test_async_callback_filter_deployments_failure_logging_is_coordinated():
class _RaisingFilter(CustomLogger):
async def async_filter_deployments(self, *args, **kwargs):
raise RuntimeError("filter blew up")
router = litellm.Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}]
)
logging_obj = _make_failure_logging_obj()
with patch.object(litellm, "callbacks", [_RaisingFilter()]): # test-quality-ok: router reads this global
await _assert_router_failure_logging_is_coordinated(
logging_obj,
lambda: router.async_callback_filter_deployments(
model="gpt-5.6",
healthy_deployments=router.model_list,
messages=None,
parent_otel_span=None,
request_kwargs={},
logging_obj=logging_obj,
),
RuntimeError,
)
@pytest.mark.asyncio
async def test_async_get_available_deployment_failure_logging_is_coordinated():
router = litellm.Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}]
)
logging_obj = _make_failure_logging_obj()
await _assert_router_failure_logging_is_coordinated(
logging_obj,
lambda: router.async_get_available_deployment(
model="model-that-is-not-configured",
request_kwargs={"litellm_logging_obj": logging_obj},
messages=[{"role": "user", "content": "hi"}],
),
litellm.BadRequestError,
)
@pytest.mark.asyncio
async def test_async_get_available_deployment_for_pass_through_failure_logging_is_coordinated():
router = litellm.Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}}]
)
logging_obj = _make_failure_logging_obj()
await _assert_router_failure_logging_is_coordinated(
logging_obj,
lambda: router.async_get_available_deployment_for_pass_through(
model="gpt-5.6",
request_kwargs={"litellm_logging_obj": logging_obj},
),
litellm.BadRequestError,
)
class _InFlightTracker:
def __init__(self) -> None:
self.current = 0
self.peak = 0
def enter(self) -> None:
self.current += 1
self.peak = max(self.peak, self.current)
def exit(self) -> None:
self.current -= 1
_SSE_CHUNKS: Final[tuple[bytes, ...]] = tuple(
b'data: {"id":"c","object":"chat.completion.chunk","created":1,"model":"gpt-5.6",'
b'"choices":[{"index":0,"delta":{"content":"x"},"finish_reason":null}]}\n\n'
for _ in range(5)
)
class _CountingSSEStream(httpx.AsyncByteStream):
def __init__(self, tracker: _InFlightTracker) -> None:
self._tracker = tracker
self._in_flight = False
def _finish(self) -> None:
if self._in_flight:
self._in_flight = False
self._tracker.exit()
async def __aiter__(self):
self._in_flight = True
self._tracker.enter()
try:
for chunk in _SSE_CHUNKS:
await asyncio.sleep(0.02)
yield chunk
finally:
await self.aclose()
yield b"data: [DONE]\n\n"
async def aclose(self) -> None:
await asyncio.sleep(0.02)
self._finish()
def _max_parallel_router(max_parallel_requests: int) -> Router:
return Router(
model_list=[
{
"model_name": "gpt-5.6",
"litellm_params": {
"model": "openai/gpt-5.6",
"api_key": "sk-fake",
"api_base": "https://max-parallel.local/v1",
"max_parallel_requests": max_parallel_requests,
},
}
],
num_retries=0,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
async def test_router_max_parallel_requests_bounds_in_flight_upstream_calls(
monkeypatch: pytest.MonkeyPatch, stream: bool
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
tracker: Final = _InFlightTracker()
router: Final = _max_parallel_router(max_parallel_requests=2)
async def upstream(request: httpx.Request) -> httpx.Response:
if stream:
return httpx.Response(
200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker)
)
tracker.enter()
await asyncio.sleep(0.05)
tracker.exit()
return httpx.Response(
200,
json={
"id": "c",
"object": "chat.completion",
"created": 1,
"model": "gpt-5.6",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "x"}, "finish_reason": "stop"}],
},
)
async def one_call() -> None:
response = await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=stream
)
if stream:
async for _ in response:
pass
with respx.mock(assert_all_called=True) as respx_mock:
respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(side_effect=upstream)
await asyncio.wait_for(asyncio.gather(*(one_call() for _ in range(10))), timeout=10)
assert tracker.peak <= 2
assert tracker.current == 0
@pytest.mark.asyncio
async def test_router_max_parallel_requests_slot_released_when_stream_closed_early(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
tracker: Final = _InFlightTracker()
router: Final = _max_parallel_router(max_parallel_requests=1)
with respx.mock() as respx_mock:
respx_mock.post("https://max-parallel.local/v1/chat/completions").mock(
side_effect=lambda request: httpx.Response(
200, headers={"content-type": "text/event-stream"}, stream=_CountingSSEStream(tracker)
)
)
first: Final = await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True
)
await first.__anext__()
async def second_call() -> None:
second = await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], stream=True
)
async for _ in second:
pass
second_task: Final = asyncio.create_task(second_call())
await asyncio.sleep(0.05)
assert tracker.current == 1
await first.aclose()
await asyncio.wait_for(second_task, timeout=2)
assert tracker.peak == 1
assert tracker.current == 0

2
ui/litellm-dashboard/eslint-budgets.json
View file

@ -3,7 +3,7 @@
"no-console": { "max": 12, "target": 0 },
"complexity": { "max": 140, "target": 80 },
"max-depth": { "max": 70, "target": 30 },
"local/no-large-inline-object-arg": { "max": 555, "target": 300 },
"local/no-large-inline-object-arg": { "max": 554, "target": 300 },
"local/no-long-condition-chain": { "max": 265, "target": 120 },
"testing-library/no-container": { "max": 133, "target": 50 },
"testing-library/no-node-access": { "max": 716, "target": 500 },

1
ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx
View file

@ -72,6 +72,7 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
return (
<DataTable
data={filteredAgents}
paginationMode="client"
columns={columns}
getRowId={(agent, index) => agent.agent_id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx
View file

@ -46,6 +46,7 @@ const GuardrailTable: React.FC<GuardrailTableProps> = ({
return (
<DataTable
data={guardrailsList}
paginationMode="client"
columns={columns}
getRowId={(guardrail, index) => guardrail.guardrail_id || String(index)}
sortingMode="client"

4
ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.test.ts
View file

@ -119,6 +119,8 @@ describe("useModelsInfo", () => {
// every other consumer of this hook keeps seeing auto-routers.
false,
undefined,
undefined,
false,
);
expect(modelInfoCall).toHaveBeenCalledTimes(1);
});
@ -147,6 +149,8 @@ describe("useModelsInfo", () => {
// every other consumer of this hook keeps seeing auto-routers.
false,
undefined,
undefined,
false,
);
});

6
ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModels.ts
View file

@ -39,6 +39,8 @@ export const useModelsInfo = (
sortOrder?: string,
excludeAutoRouters: boolean = false,
modelName?: string,
accessGroup?: string,
wildcardOnly: boolean = false,
) => {
const { accessToken, userId, userRole } = useAuthorized();
return useQuery<PaginatedModelInfoResponse>({
@ -57,6 +59,8 @@ export const useModelsInfo = (
// Part of the key: callers that exclude auto-routers must not share a cache entry
// with callers that keep them.
...(excludeAutoRouters && { excludeAutoRouters: "true" }),
...(accessGroup && { accessGroup }),
...(wildcardOnly && { wildcardOnly: "true" }),
},
}),
queryFn: async () =>
@ -73,6 +77,8 @@ export const useModelsInfo = (
sortOrder,
excludeAutoRouters,
modelName,
accessGroup,
wildcardOnly,
),
enabled: Boolean(accessToken && userId && userRole),
});

24
ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts
View file

@ -671,7 +671,7 @@ describe("useDeletedTeams", () => {
it("should return deleted teams data when query is successful", async () => {
(global.fetch as any).mockResolvedValue({
ok: true,
json: async () => ({ teams: mockDeletedTeams }),
json: async () => ({ teams: mockDeletedTeams, total: 2, page: 1, page_size: 10, total_pages: 1 }),
});
const { result } = renderHook(() => useDeletedTeams(1, 10, {}), { wrapper });
@ -684,10 +684,26 @@ describe("useDeletedTeams", () => {
expect(result.current.isSuccess).toBe(true);
});
expect(result.current.data).toEqual(mockDeletedTeams);
expect(result.current.data).toEqual({ teams: mockDeletedTeams, total: 2 });
expect(result.current.error).toBeNull();
});
it("should keep the server total so the table can paginate beyond the current page", async () => {
(global.fetch as any).mockResolvedValue({
ok: true,
json: async () => ({ teams: mockDeletedTeams, total: 137, page: 1, page_size: 2, total_pages: 69 }),
});
const { result } = renderHook(() => useDeletedTeams(1, 2, {}), { wrapper });
await waitFor(() => {
expect(result.current.isSuccess).toBe(true);
});
expect(result.current.data?.total).toBe(137);
expect((global.fetch as any).mock.calls[0][0]).toContain("page_size=2");
});
it("should handle error when API call fails", async () => {
(global.fetch as any).mockResolvedValue({
ok: false,
@ -744,7 +760,7 @@ describe("useDeletedTeams", () => {
rerender({ page: 2 });
expect(result.current.data).toEqual(mockDeletedTeams);
expect(result.current.data?.teams).toEqual(mockDeletedTeams);
});
it("should pass options to API call", async () => {
@ -785,7 +801,7 @@ describe("useDeletedTeams", () => {
expect(result.current.isSuccess).toBe(true);
});
expect(result.current.data).toEqual(mockDeletedTeams);
expect(result.current.data).toEqual({ teams: mockDeletedTeams, total: 2 });
expect(result.current.error).toBeNull();
});
});

21
ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.ts
View file

@ -20,6 +20,11 @@ export interface DeletedTeam extends Team {
deleted_by: string;
}
export interface DeletedTeamsResponse {
teams: DeletedTeam[];
total: number;
}
export interface TeamListCallOptions {
organizationID?: string | null;
teamID?: string | null;
@ -209,7 +214,7 @@ const deletedTeamListCall = async (
page: number,
pageSize: number,
options: TeamListCallOptions = {},
) => {
): Promise<DeletedTeamsResponse> => {
/**
* Get deleted teams from proxy
*/
@ -251,14 +256,12 @@ const deletedTeamListCall = async (
throw new Error(errorMessage);
}
const data = await response.json();
const data: DeletedTeam[] | (Partial<DeletedTeamsResponse> & { teams: DeletedTeam[] }) = await response.json();
// Extract teams array from response if it's wrapped in a response object
// Otherwise return the data directly if it's already an array
if (data && typeof data === "object" && "teams" in data) {
return data.teams as DeletedTeam[];
if (Array.isArray(data)) {
return { teams: data, total: data.length };
}
return data as DeletedTeam[];
return { teams: data.teams, total: data.total ?? data.teams.length };
} catch (error) {
console.error("Failed to list deleted teams:", error);
throw error;
@ -270,10 +273,10 @@ export const useDeletedTeams = (
page: number,
pageSize: number,
options: TeamListCallOptions = {},
): UseQueryResult<DeletedTeam[]> => {
): UseQueryResult<DeletedTeamsResponse> => {
const { accessToken } = useAuthorized();
return useQuery<DeletedTeam[]>({
return useQuery<DeletedTeamsResponse>({
queryKey: deletedTeamKeys.list({ page, limit: pageSize, ...options }),
queryFn: async () => await deletedTeamListCall(accessToken!, page, pageSize, options),
enabled: Boolean(accessToken),

1
ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.tsx
View file

@ -451,6 +451,7 @@ export function MCPToolsetsTab({ accessToken, userRole }: MCPToolsetsTabProps) {
<DataTable
data={toolsets}
paginationMode="client"
columns={columns}
getRowId={(toolset, index) => toolset.toolset_id || String(index)}
sortingMode="client"

49
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx
View file

@ -34,6 +34,8 @@ interface ModelsInfoArgs {
sortBy?: string;
sortOrder?: string;
modelName?: string;
accessGroup?: string;
wildcardOnly?: boolean;
}
const modelsInfoCalls: ModelsInfoArgs[] = [];
@ -50,12 +52,24 @@ type UseModelsInfoArgs = [
sortOrder?: string,
excludeAutoRouters?: boolean,
modelName?: string,
accessGroup?: string,
wildcardOnly?: boolean,
];
vi.mock("../../hooks/models/useModels", () => ({
useModelsInfo: (...args: UseModelsInfoArgs) => {
const [page, size, search, , teamId, sortBy, sortOrder, , modelName] = args;
const call: ModelsInfoArgs = { page, size, search, teamId, sortBy, sortOrder, modelName };
const [page, size, search, , teamId, sortBy, sortOrder, , modelName, accessGroup, wildcardOnly] = args;
const call: ModelsInfoArgs = {
page,
size,
search,
teamId,
sortBy,
sortOrder,
modelName,
accessGroup,
wildcardOnly,
};
modelsInfoCalls.push(call);
return { ...modelsInfoResult, refetch: mockRefetch };
},
@ -254,13 +268,38 @@ describe("AllModelsTab", () => {
});
});
it("filters the fetched page down to the selected model group", () => {
setModelsInfo([makeRow(), { ...makeRow(), model_name: "claude-opus" }], 2);
it("renders every row the server returned for the selected model group so rows match the footer total", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "claude-opus" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="claude-opus" />);
const table = screen.getByRole("table");
expect(within(table).getByText("claude-opus")).toBeInTheDocument();
expect(within(table).queryByText("gpt-4")).not.toBeInTheDocument();
expect(within(table).getByText("gpt-4")).toBeInTheDocument();
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-2 of 2");
});
it("asks the server for wildcard deployments instead of hiding rows client-side", () => {
setModelsInfo([makeRow(), { ...makeRow({ model_info: { id: "model-2" } }), model_name: "openai/*" }], 2);
render(<AllModelsTab {...defaultProps} selectedModelGroup="wildcard" />);
expect(lastModelsInfoCall().wildcardOnly).toBe(true);
expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument();
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-2 of 2");
});
it("asks the server for the selected access group instead of hiding rows client-side", async () => {
const user = userEvent.setup();
render(<AllModelsTab {...defaultProps} />);
expect(lastModelsInfoCall().wildcardOnly).toBe(false);
await user.click(screen.getByTestId("datatable-filters-trigger"));
await user.click(await screen.findByPlaceholderText("Filter by Model Access Group"));
await user.click(await screen.findByRole("option", { name: "sales-team" }));
await user.click(screen.getByTestId("filter-drawer-apply"));
await waitFor(() => expect(lastModelsInfoCall().accessGroup).toBe("sales-team"));
expect(within(screen.getByRole("table")).getByText("gpt-4")).toBeInTheDocument();
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-1 of 1");
});
it("asks the server for the exact selected model group so deployments beyond the first page are found", () => {

32
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx
View file

@ -86,6 +86,11 @@ const AllModelsTab = ({
selectedModelGroup !== ALL_MODEL_GROUPS_VALUE &&
selectedModelGroup !== WILDCARD_MODEL_GROUP_VALUE;
const modelNameForQuery = isConcreteModelGroup ? selectedModelGroup ?? undefined : undefined;
const accessGroupForQuery =
selectedModelAccessGroupFilter && selectedModelAccessGroupFilter !== ALL_MODEL_GROUPS_VALUE
? selectedModelAccessGroupFilter
: undefined;
const wildcardOnlyForQuery = selectedModelGroup === WILDCARD_MODEL_GROUP_VALUE;
const sortBy = useMemo(() => {
if (sorting.length === 0) return undefined;
@ -114,6 +119,8 @@ const AllModelsTab = ({
// lists and manages them. Excluded server-side so total_count stays honest.
true,
modelNameForQuery,
accessGroupForQuery,
wildcardOnlyForQuery,
);
const isLoading = isLoadingModelsInfo || isLoadingModelCostMap;
@ -129,32 +136,11 @@ const AllModelsTab = ({
[modelCostMapData],
);
const modelData = useMemo(() => {
const modelData = useMemo<{ data: ModelData[] }>(() => {
if (!rawModelData) return { data: [] };
return transformModelData(rawModelData, getProviderFromModel);
}, [rawModelData, getProviderFromModel]);
const filteredData = useMemo<ModelData[]>(() => {
if (!modelData || !modelData.data || modelData.data.length === 0) {
return [];
}
return modelData.data.filter((model: ModelData) => {
const modelNameMatch =
selectedModelGroup === ALL_MODEL_GROUPS_VALUE ||
model.model_name === selectedModelGroup ||
!selectedModelGroup ||
(selectedModelGroup === WILDCARD_MODEL_GROUP_VALUE && model.model_name?.includes("*"));
const accessGroupMatch =
selectedModelAccessGroupFilter === ALL_MODEL_GROUPS_VALUE ||
model.model_info["access_groups"]?.includes(selectedModelAccessGroupFilter ?? "") ||
!selectedModelAccessGroupFilter;
return modelNameMatch && accessGroupMatch;
});
}, [modelData, selectedModelGroup, selectedModelAccessGroupFilter]);
const columnFilters = useMemo<ColumnFiltersState>(
() =>
[
@ -270,7 +256,7 @@ const AllModelsTab = ({
<div className="w-full">
<div className="flex flex-col gap-3">
<AllModelsTable
data={filteredData}
data={modelData.data}
rowCount={rawModelData?.total_count ?? 0}
isLoading={isLoading}
isRefreshing={isFetchingModelsInfo}

1
ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AccessGroupBudgetsPanel.tsx
View file

@ -84,6 +84,7 @@ export default function AccessGroupBudgetsPanel() {
<DataTable
data={accessGroups ?? []}
paginationMode="client"
columns={columns}
getRowId={(group) => group.access_group}
sortingMode="client"

17
ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.test.tsx
View file

@ -192,6 +192,23 @@ describe("OrganizationsTable", () => {
expect(screen.queryByText("ShouldNotShow")).not.toBeInTheDocument();
});
it("pages long lists client-side with the shared size selector and footer", async () => {
const user = userEvent.setup();
const organizations = Array.from({ length: 30 }, (_, index) =>
makeOrganization({ organization_id: `org-${index}`, organization_alias: `Org ${index}` }),
);
render(<OrganizationsTable {...baseProps} organizations={organizations} />);
expect(screen.getAllByRole("row")).toHaveLength(26);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 30");
await user.click(screen.getByTestId("pagination-page-size"));
await user.click(await screen.findByRole("option", { name: "50" }));
expect(screen.getAllByRole("row")).toHaveLength(31);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-30 of 30");
});
it("uses a search-aware empty state", () => {
const { rerender } = render(<OrganizationsTable {...baseProps} searchActive={false} organizations={[]} />);
expect(screen.getByText("No organizations yet")).toBeInTheDocument();

1
ui/litellm-dashboard/src/app/(dashboard)/organizations/_components/OrganizationsTable.tsx
View file

@ -59,6 +59,7 @@ const OrganizationsTable: React.FC<OrganizationsTableProps> = ({
return (
<DataTable
data={organizations}
paginationMode="client"
columns={columns}
getRowId={(organization, index) => organization.organization_id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/policies/_components/AttachmentTable.tsx
View file

@ -50,6 +50,7 @@ const AttachmentTable: React.FC<AttachmentTableProps> = ({
return (
<DataTable
data={attachments}
paginationMode="client"
columns={columns}
getRowId={(row) => row.attachment_id}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/policies/_components/PolicyTable.tsx
View file

@ -71,6 +71,7 @@ const PolicyTable: React.FC<PolicyTableProps> = ({
return (
<DataTable
data={rows}
paginationMode="client"
columns={columns}
getRowId={(row) => `${row.primaryPolicy.definition_location ?? "db"}:${row.policy_name}`}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/prompts/_components/PromptTable.tsx
View file

@ -73,6 +73,7 @@ const PromptTable: React.FC<PromptTableProps> = ({
return (
<DataTable
data={promptsList}
paginationMode="client"
columns={columns}
getRowId={(prompt, index) =>
prompt.prompt_id ? `${prompt.prompt_id}::${prompt.environment || "development"}` : String(index)

1
ui/litellm-dashboard/src/app/(dashboard)/search-tools/_components/SearchToolTable.tsx
View file

@ -50,6 +50,7 @@ const SearchToolTable: React.FC<SearchToolTableProps> = ({
return (
<DataTable
data={searchTools}
paginationMode="client"
columns={columns}
getRowId={(tool, index) => searchToolKey(tool) || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/skills/_components/PluginTable.tsx
View file

@ -42,6 +42,7 @@ const PluginTable: React.FC<PluginTableProps> = ({ pluginsList, isLoading, onDel
return (
<DataTable
data={pluginsList}
paginationMode="client"
columns={columns}
getRowId={(plugin, index) => plugin.id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx
View file

@ -39,6 +39,7 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
return (
<DataTable
data={data}
paginationMode="client"
columns={columns}
getRowId={(tag, index) => tag.name || String(index)}
fillHeight

1
ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/IndexesTable.tsx
View file

@ -46,6 +46,7 @@ const IndexesTable: React.FC<IndexesTableProps> = ({
return (
<DataTable
data={data}
paginationMode="client"
columns={columns}
getRowId={(row, index) => row.id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreTable.tsx
View file

@ -41,6 +41,7 @@ const VectorStoreTable: React.FC<VectorStoreTableProps> = ({ data, onView, onEdi
return (
<DataTable
data={data}
paginationMode="client"
columns={columns}
getRowId={(vectorStore, index) => vectorStore.vector_store_id || String(index)}
sortingMode="client"

3
ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx
View file

@ -474,6 +474,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
{/* Model Table */}
<DataTable
data={filteredData}
paginationMode="client"
columns={modelColumns}
getRowId={(model, index) => model.model_group || String(index)}
sortingMode="client"
@ -540,6 +541,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
{/* Agent Table */}
<DataTable
data={filteredAgentData}
paginationMode="client"
columns={agentColumns}
getRowId={(agent, index) => agent.agent_id || agent.name || String(index)}
sortingMode="client"
@ -581,6 +583,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
{/* MCP Server Table */}
<DataTable
data={mcpHubData || []}
paginationMode="client"
columns={mcpColumns}
getRowId={(server, index) => server.server_id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx
View file

@ -162,6 +162,7 @@ const SkillHubDashboard: React.FC<SkillHubDashboardProps> = ({
</div>
<DataTable
data={filteredSkills}
paginationMode="client"
columns={columns}
getRowId={(skill, index) => skill.id || String(index)}
sortingMode="client"

48
ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.test.tsx
View file

@ -1,4 +1,5 @@
import { screen } from "@testing-library/react";
import { fireEvent, screen } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { vi, it, expect, beforeEach, MockedFunction } from "vitest";
import { renderWithProviders } from "../../../tests/test-utils";
import DeletedTeamsPage from "./DeletedTeamsPage";
@ -31,7 +32,7 @@ beforeEach(() => {
vi.clearAllMocks();
mockUseDeletedTeams.mockReturnValue({
data: [mockDeletedTeam],
data: { teams: [mockDeletedTeam], total: 1 },
isLoading: false,
} as unknown as ReturnType<typeof useDeletedTeams>);
});
@ -42,6 +43,49 @@ it("should render DeletedTeamsPage component", () => {
expect(screen.getByText("Test Team")).toBeInTheDocument();
});
it("requests the first page of 25 deleted teams and shows the server total in the footer", () => {
mockUseDeletedTeams.mockReturnValue({
data: { teams: [mockDeletedTeam], total: 137 },
isLoading: false,
} as unknown as ReturnType<typeof useDeletedTeams>);
renderWithProviders(<DeletedTeamsPage />);
expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(1, 25);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 137");
expect(screen.getByTestId("pagination-next")).toBeEnabled();
});
it("requests the next page from the server when Next is clicked", () => {
mockUseDeletedTeams.mockReturnValue({
data: { teams: [mockDeletedTeam], total: 137 },
isLoading: false,
} as unknown as ReturnType<typeof useDeletedTeams>);
renderWithProviders(<DeletedTeamsPage />);
fireEvent.click(screen.getByTestId("pagination-next"));
expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(2, 25);
});
it("offers the shared page sizes and refetches with the selected one", async () => {
const user = userEvent.setup();
mockUseDeletedTeams.mockReturnValue({
data: { teams: [mockDeletedTeam], total: 137 },
isLoading: false,
} as unknown as ReturnType<typeof useDeletedTeams>);
renderWithProviders(<DeletedTeamsPage />);
await user.click(screen.getByTestId("pagination-page-size"));
const options = await screen.findAllByRole("option");
expect(options.map((option) => option.textContent)).toEqual(["25", "50", "100"]);
await user.click(screen.getByRole("option", { name: "100" }));
expect(mockUseDeletedTeams).toHaveBeenLastCalledWith(1, 100);
});
it("should show the enterprise notice for a non-premium user", () => {
renderWithProviders(<DeletedTeamsPage />);

17
ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsPage.tsx
View file

@ -1,13 +1,20 @@
"use client";
import { PaginationState } from "@tanstack/react-table";
import { Info } from "lucide-react";
import { useState } from "react";
import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert";
import { DEFAULT_PAGE_SIZE_OPTIONS } from "@/components/shared/DataTable";
import { useDeletedTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
import { DeletedTeamsTable } from "./DeletedTeamsTable/DeletedTeamsTable";
export default function DeletedTeamsPage() {
const { premiumUser } = useAuthorized();
const { data: teamsData, isLoading } = useDeletedTeams(1, 100);
const [pagination, setPagination] = useState<PaginationState>({
pageIndex: 0,
pageSize: DEFAULT_PAGE_SIZE_OPTIONS[0],
});
const { data: teamsData, isLoading } = useDeletedTeams(pagination.pageIndex + 1, pagination.pageSize);
return (
<div className="flex flex-col gap-4">
@ -20,7 +27,13 @@ export default function DeletedTeamsPage() {
</AlertDescription>
</Alert>
)}
<DeletedTeamsTable teams={teamsData || []} isLoading={isLoading} />
<DeletedTeamsTable
teams={teamsData?.teams ?? []}
isLoading={isLoading}
pagination={pagination}
onPaginationChange={setPagination}
rowCount={teamsData?.total ?? 0}
/>
</div>
);
}

32
ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.test.tsx
View file

@ -22,12 +22,19 @@ const makeDeletedTeam = (overrides: Partial<DeletedTeam> = {}): DeletedTeam => (
...overrides,
});
const paginationProps = {
pagination: { pageIndex: 0, pageSize: 25 },
onPaginationChange: vi.fn(),
};
beforeEach(() => {
vi.clearAllMocks();
});
it("should display team information", () => {
renderWithProviders(<DeletedTeamsTable teams={[makeDeletedTeam()]} isLoading={false} />);
renderWithProviders(
<DeletedTeamsTable teams={[makeDeletedTeam()]} isLoading={false} rowCount={1} {...paginationProps} />,
);
expect(screen.getByText("Test Team")).toBeInTheDocument();
expect(screen.getByText("team-1")).toBeInTheDocument();
@ -39,7 +46,7 @@ it("should sort teams by deleted_at descending by default", () => {
makeDeletedTeam({ team_id: "team-old", team_alias: "older-team", deleted_at: "2024-01-01T10:00:00Z" }),
makeDeletedTeam({ team_id: "team-new", team_alias: "newer-team", deleted_at: "2024-06-01T10:00:00Z" }),
];
renderWithProviders(<DeletedTeamsTable teams={teams} isLoading={false} />);
renderWithProviders(<DeletedTeamsTable teams={teams} isLoading={false} rowCount={2} {...paginationProps} />);
const rows = screen.getAllByRole("row").slice(1);
expect(within(rows[0]).getByText("newer-team")).toBeInTheDocument();
@ -47,13 +54,30 @@ it("should sort teams by deleted_at descending by default", () => {
});
it("should show skeleton rows when loading", () => {
renderWithProviders(<DeletedTeamsTable teams={[]} isLoading />);
renderWithProviders(<DeletedTeamsTable teams={[]} isLoading rowCount={0} {...paginationProps} />);
expect(screen.getAllByTestId("skeleton-row").length).toBeGreaterThan(0);
});
it("should show the empty state when there are no deleted teams", () => {
renderWithProviders(<DeletedTeamsTable teams={[]} isLoading={false} />);
renderWithProviders(<DeletedTeamsTable teams={[]} isLoading={false} rowCount={0} {...paginationProps} />);
expect(screen.getByText("No deleted teams found")).toBeInTheDocument();
});
it("renders the shared pagination footer with the server row count", () => {
renderWithProviders(
<DeletedTeamsTable
teams={[makeDeletedTeam()]}
isLoading={false}
rowCount={137}
pagination={{ pageIndex: 2, pageSize: 50 }}
onPaginationChange={vi.fn()}
/>,
);
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 101-137 of 137");
expect(screen.getByTestId("pagination-page-size")).toHaveTextContent("50");
expect(screen.getByTestId("pagination-prev")).toBeEnabled();
expect(screen.getByTestId("pagination-next")).toBeDisabled();
});

17
ui/litellm-dashboard/src/components/DeletedTeamsPage/DeletedTeamsTable/DeletedTeamsTable.tsx
View file

@ -1,6 +1,6 @@
"use client";
import { SortingState } from "@tanstack/react-table";
import { OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
import { Inbox } from "lucide-react";
import { useMemo, useState } from "react";
@ -12,6 +12,9 @@ import { getDeletedTeamsTableColumns } from "./DeletedTeamsTableColumns";
interface DeletedTeamsTableProps {
teams: DeletedTeam[];
isLoading: boolean;
pagination: PaginationState;
onPaginationChange: OnChangeFn<PaginationState>;
rowCount: number;
}
const DEFAULT_SORTING: SortingState = [{ id: "deleted_at", desc: true }];
@ -28,7 +31,13 @@ function EmptyState() {
);
}
export function DeletedTeamsTable({ teams, isLoading }: DeletedTeamsTableProps) {
export function DeletedTeamsTable({
teams,
isLoading,
pagination,
onPaginationChange,
rowCount,
}: DeletedTeamsTableProps) {
const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
const columns = useMemo(() => getDeletedTeamsTableColumns(), []);
@ -41,6 +50,10 @@ export function DeletedTeamsTable({ teams, isLoading }: DeletedTeamsTableProps)
sortingMode="client"
sorting={sorting}
onSortingChange={setSorting}
paginationMode="server"
pagination={pagination}
onPaginationChange={onPaginationChange}
rowCount={rowCount}
isLoading={isLoading}
loadingMessage="Loading deleted teams…"
noDataMessage={<EmptyState />}

1
ui/litellm-dashboard/src/components/PassThroughSettings/PassThroughEndpointsTable.tsx
View file

@ -41,6 +41,7 @@ export function PassThroughEndpointsTable({
return (
<DataTable
data={endpoints}
paginationMode="client"
columns={columns}
getRowId={(endpoint, index) => endpoint.id || endpoint.path || String(index)}
isLoading={isLoading}

2
ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
View file

@ -1170,7 +1170,7 @@ describe("buildComplexityRouterConfig stall escalation", () => {
});
it("emits the toggle and both knobs when it is on", () => {
const params = {
const params: BuildComplexityRouterConfigParams = {
...baseParams,
stallEscalationEnabled: true,
stallEscalationWindow: 8,

1
ui/litellm-dashboard/src/components/model_add/CredentialsTable.tsx
View file

@ -48,6 +48,7 @@ const CredentialsTable: React.FC<CredentialsTableProps> = ({
return (
<DataTable
data={credentials}
paginationMode="client"
columns={columns}
getRowId={(credential, index) => credential.credential_name || String(index)}
sortingMode="client"

8
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -1724,6 +1724,8 @@ export const modelInfoCall = async (
sortOrder?: string,
excludeAutoRouters?: boolean,
modelName?: string,
accessGroup?: string,
wildcardOnly?: boolean,
) => {
/**
* Get all models on proxy
@ -1755,6 +1757,12 @@ export const modelInfoCall = async (
if (excludeAutoRouters) {
params.append("exclude_auto_routers", "true");
}
if (accessGroup && accessGroup.trim()) {
params.append("access_group", accessGroup.trim());
}
if (wildcardOnly) {
params.append("wildcard_only", "true");
}
if (params.toString()) {
url += `?${params.toString()}`;
}

2
ui/litellm-dashboard/src/components/public_model_hub.tsx
View file

@ -554,6 +554,7 @@ const PublicModelHub: React.FC<PublicModelHubProps> = ({ accessToken, isEmbedded
<DataTable
data={filteredAgentData}
paginationMode="client"
columns={agentColumns}
getRowId={(agent, index) => agent.name || String(index)}
sortingMode="client"
@ -620,6 +621,7 @@ const PublicModelHub: React.FC<PublicModelHubProps> = ({ accessToken, isEmbedded
<DataTable
data={filteredMcpData}
paginationMode="client"
columns={mcpColumns}
getRowId={(server, index) => server.server_id || String(index)}
sortingMode="client"

1
ui/litellm-dashboard/src/components/routing_groups/RoutingGroupsTable.tsx
View file

@ -64,6 +64,7 @@ const RoutingGroupsTable: React.FC<RoutingGroupsTableProps> = ({
return (
<DataTable
data={groups}
paginationMode="client"
columns={columns}
getRowId={(group) => group.group_name}
sortingMode="client"

1
ui/litellm-dashboard/src/components/team/AvailableTeamsTable.tsx
View file

@ -46,6 +46,7 @@ const AvailableTeamsTable: React.FC<AvailableTeamsTableProps> = ({ teams, isLoad
return (
<DataTable
data={teams}
paginationMode="client"
columns={columns}
getRowId={(team, index) => team.team_id || String(index)}
sortingMode="client"

52
ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx
View file

@ -148,13 +148,61 @@ describe("RequestLogsPanel", () => {
});
describe("server-grouped session pagination (#38060)", () => {
it("requests session-grouped pages of 10 rows by default without a cursor", async () => {
it("requests session-grouped pages of 25 rows by default without a cursor", async () => {
renderPanel();
await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
expect(lastCall()?.params?.group_by_session).toBe(true);
expect(lastCall()?.params?.session_cursor).toBeUndefined();
expect(lastCall()?.page_size).toBe(10);
expect(lastCall()?.page_size).toBe(25);
});
it("offers the same page sizes as the other tables", async () => {
const user = userEvent.setup();
respondWith([logEntry({ request_id: "req-a" })]);
renderPanel();
await waitFor(() => expect(row("req-a")).not.toBeNull());
await user.click(screen.getByTestId("pagination-page-size"));
const options = await screen.findAllByRole("option");
expect(options.map((option) => option.textContent)).toEqual(["25", "50", "100"]);
});
it("counts the rendered rows in the footer instead of the server's session total", async () => {
const lastPage = {
data: [logEntry({ request_id: "req-a" }), logEntry({ request_id: "req-b" }), logEntry({ request_id: "req-c" })],
total: 40,
page: 1,
page_size: 25,
total_pages: 2,
next_session_cursor: null,
has_more: false,
};
vi.mocked(uiSpendLogsCall).mockResolvedValue(lastPage);
renderPanel();
await waitFor(() => expect(row("req-a")).not.toBeNull());
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-3 of 3");
expect(screen.getByTestId("pagination-next")).toBeDisabled();
});
it("keeps Next enabled from the server total while more session pages remain", async () => {
const firstPage = {
data: Array.from({ length: 25 }, (_, index) => logEntry({ request_id: `req-${index}` })),
total: 80,
page: 1,
page_size: 25,
total_pages: 4,
next_session_cursor: "2026-07-07 09:50:13|key-1|sess-1",
has_more: true,
};
vi.mocked(uiSpendLogsCall).mockResolvedValue(firstPage);
renderPanel();
await waitFor(() => expect(row("req-0")).not.toBeNull());
expect(screen.getByTestId("pagination-range")).toHaveTextContent("Showing 1-25 of 80");
expect(screen.getByTestId("pagination-next")).toBeEnabled();
});
it("renders every row the server returns without client-side collapsing", async () => {

10
ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx
View file

@ -6,13 +6,13 @@ import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } fr
import moment from "moment";
import { useCallback, useEffect, useMemo, useState } from "react";
import { DEFAULT_PAGE_SIZE_OPTIONS } from "@/components/shared/DataTable";
import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells";
import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
import type { KeyResponse } from "../key_team_helpers/key_list";
import { keyInfoV1Call, uiSpendLogsCall } from "../networking";
import KeyInfoView from "../templates/key_info_view";
import type { LogEntry } from "./columns";
import { LOGS_PAGE_SIZE_OPTIONS } from "./constants";
import {
DEFAULT_LOGS_SORTING,
formatLogsWindow,
@ -26,7 +26,7 @@ import { LogDetailsDrawer } from "./LogDetailsDrawer";
import { LiveTailBanner, LogsTableToolbar } from "./LogsTableToolbar";
import { RequestLogsTable } from "./RequestLogsTable";
const PAGE_SIZE = LOGS_PAGE_SIZE_OPTIONS[0];
const PAGE_SIZE = DEFAULT_PAGE_SIZE_OPTIONS[0];
const DEFAULT_INTERVAL = { value: 24, unit: "hours" };
interface RequestLogsPanelProps {
@ -166,6 +166,10 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
const isDrawerOpen = displayLog !== null || displaySessionId !== null;
const rows: LogEntry[] = filteredLogs.data;
const rowsThroughThisPage = pagination.pageIndex * pagination.pageSize + rows.length;
const isLastPage =
filteredLogs.has_more === false || (filteredLogs.has_more === undefined && rows.length < pagination.pageSize);
const rowCount = isLastPage ? rowsThroughThisPage : Math.max(filteredLogs.total, rowsThroughThisPage);
const handleSearchChange = useCallback((value: string) => {
setColumnFilters((previous) => {
@ -290,7 +294,7 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
<RequestLogsTable
data={rows}
rowCount={filteredLogs.total}
rowCount={rowCount}
isLoading={logsQuery.isLoading}
isRefreshing={logsQuery.isFetching}
pagination={pagination}

2
ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
View file

@ -8,7 +8,6 @@ import { DataTable, DataTableFilterDrawer, DataTableToolbar } from "@/components
import type { Team } from "../key_team_helpers/key_list";
import type { LogEntry } from "./columns";
import { LOGS_PAGE_SIZE_OPTIONS } from "./constants";
import { LOG_FILTER_LABELS, type LogsWindow } from "./log_filter_logic";
import { RequestLogsFilters } from "./RequestLogsFilters";
import { getRequestLogsTableColumns } from "./RequestLogsTableColumns";
@ -93,7 +92,6 @@ export function RequestLogsTable({
paginationMode="server"
pagination={pagination}
onPaginationChange={onPaginationChange}
pageSizeOptions={LOGS_PAGE_SIZE_OPTIONS}
rowCount={rowCount}
filterMode="server"
columnFilters={columnFilters}

Some files were not shown because too many files have changed in this diff Show more