diff --git a/Dockerfile b/Dockerfile index b3ee85e9ed1..29a085a4ef9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -101,6 +101,12 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root +# The base image only configures Chainguard's authenticated apk repo, which +# requires an enterprise subscription. Add the public Wolfi repo so `apk add` +# also works for anyone installing extra packages into a running container. +# https://github.com/BerriAI/litellm/issues/33518 +RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories + # node (without npm) is required by the prisma CLI at runtime RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile diff --git a/litellm/constants.py b/litellm/constants.py index 9a50797f517..1bd977dd9a9 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -136,6 +136,7 @@ MCP_CLIENT_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_CLIENT_TIMEOUT", "60.0" MCP_TOOL_LISTING_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_TOOL_LISTING_TIMEOUT", "30.0")) MCP_METADATA_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_METADATA_TIMEOUT", "10.0")) MCP_HEALTH_CHECK_TIMEOUT: Final = float(os.getenv("LITELLM_MCP_HEALTH_CHECK_TIMEOUT", "10.0")) +MCP_TOOL_LISTING_MAX_PAGES: Final = 1000 # Allowlist of commands permitted for MCP stdio transport. # Prevents arbitrary command execution via /mcp-rest/test/* endpoints or server creation. @@ -1589,6 +1590,7 @@ KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job" MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job" +USER_SPEND_ALERTS_JOB_ID: Final = "user_spend_alerts_job" PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job" SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report" SLACK_MODEL_DEPRECATION_LOCK_ID: Final = "slack_model_deprecation_warning" diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index f0a1bff8fdc..ea81e323da4 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -7,6 +7,7 @@ import base64 import os from collections.abc import Awaitable, Callable, Generator from datetime import timedelta +from functools import partial from importlib import metadata from typing import Any, Final, TypeVar @@ -47,7 +48,8 @@ from mcp.types import Tool as MCPTool from pydantic import AnyUrl from litellm._logging import verbose_logger -from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_NPM_CACHE_DIR, MCP_TOOL_LISTING_TIMEOUT +from litellm.experimental_mcp_client.tools import list_tools_with_pagination from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.types.llms.custom_http import VerifyTypes from litellm.types.mcp import ( @@ -603,17 +605,19 @@ class MCPClient: """ verbose_logger.debug("MCP client listing tools from %s", self.server_url or "stdio") - async def _list_tools_operation(session: ClientSession): - return await session.list_tools() - try: - result: Final = await self.run_with_session(_list_tools_operation, quiet_on_error=raise_on_error) - tool_count: Final = len(result.tools) - tool_names: Final = [tool.name for tool in result.tools] + # A per-server timeout above the global default extends the whole-walk deadline + listing_deadline: Final = max(self.timeout, MCP_TOOL_LISTING_TIMEOUT) + tools: Final = await self.run_with_session( + partial(list_tools_with_pagination, listing_deadline=listing_deadline), + quiet_on_error=raise_on_error, + ) + tool_count: Final = len(tools) + tool_names: Final = tuple(tool.name for tool in tools) verbose_logger.info( "MCP client listed %s tools from %s: %s", tool_count, self.server_url or "stdio", tool_names ) - return result.tools + return tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") raise diff --git a/litellm/experimental_mcp_client/tools.py b/litellm/experimental_mcp_client/tools.py index 30d50e2a74b..51d2139ef3b 100644 --- a/litellm/experimental_mcp_client/tools.py +++ b/litellm/experimental_mcp_client/tools.py @@ -1,14 +1,22 @@ import json from typing import Final, Literal +import anyio from mcp import ClientSession from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import CallToolResult as MCPCallToolResult +from mcp.types import PaginatedRequestParams from mcp.types import Tool as MCPTool from openai.types.chat import ChatCompletionToolParam from openai.types.responses.function_tool_param import FunctionToolParam from openai.types.shared_params.function_definition import FunctionDefinition +from litellm._logging import verbose_logger +from litellm.constants import ( + MCP_CLIENT_TIMEOUT, + MCP_TOOL_LISTING_MAX_PAGES, + MCP_TOOL_LISTING_TIMEOUT, +) from litellm.types.llms.anthropic import AnthropicMessagesTool from litellm.types.utils import ChatCompletionMessageToolCall @@ -90,6 +98,64 @@ def transform_mcp_tool_to_anthropic_tool(mcp_tool: MCPTool) -> AnthropicMessages ) +async def list_tools_with_pagination( + session: ClientSession, listing_deadline: float | None = None +) -> list[MCPTool]: # mutable-ok: list return contract + """Collect tools from every tools/list page by following nextCursor. + + Stops and returns the tools collected so far when the upstream repeats a + cursor, the page cap is reached, or the whole-walk deadline expires, so a + buggy or slow upstream yields a partial catalog instead of an error. + listing_deadline overrides the default whole-walk deadline; callers with a + per-server timeout above the global default pass it through here. + """ + tools: Final[list[MCPTool]] = [] # mutable-ok: accumulates each page's tools + seen_cursors: Final[set[str]] = set() # mutable-ok: guards against cursor loops + cursor: str | None = None # rebind-ok: advances to each page's nextCursor + # The per-request session read timeout restarts on every page, so a multi-page + # walk needs its own overall deadline. max() keeps the pre-pagination guarantee + # that a single page slower than the listing timeout but within the client + # timeout still succeeds. + effective_deadline: Final = ( + listing_deadline if listing_deadline is not None else max(MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT) + ) + + with anyio.move_on_after(effective_deadline): + for _ in range(MCP_TOOL_LISTING_MAX_PAGES): + result = ( + await session.list_tools() + if cursor is None + else await session.list_tools(params=PaginatedRequestParams(cursor=cursor)) + ) + tools.extend(result.tools) + + next_cursor = getattr(result, "nextCursor", None) + if not isinstance(next_cursor, str) or not next_cursor: + return tools + if next_cursor in seen_cursors: + verbose_logger.warning( + "MCP server repeated a tools/list cursor while listing tools; returning %s tools collected so far", + len(tools), + ) + return tools + seen_cursors.add(next_cursor) + cursor = next_cursor + + verbose_logger.warning( + "MCP server tools/list pagination exceeded the maximum of %s pages; returning %s tools collected so far", + MCP_TOOL_LISTING_MAX_PAGES, + len(tools), + ) + return tools + + verbose_logger.warning( + "MCP server tools/list pagination exceeded the %s second listing deadline; returning %s tools collected so far", + effective_deadline, + len(tools), + ) + return tools + + async def load_mcp_tools( session: ClientSession, format: Literal["mcp", "openai"] = "mcp" ) -> list[MCPTool] | list[ChatCompletionToolParam]: @@ -103,10 +169,12 @@ async def load_mcp_tools( If format is set to "openai", the tools are converted to OpenAI API compatible tools. """ - tools: Final = await session.list_tools() + tools: Final = await list_tools_with_pagination(session) if format == "openai": - return [transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools.tools] - return tools.tools + return [ # mutable-ok: public API returns a list + transform_mcp_tool_to_openai_tool(mcp_tool=tool) for tool in tools + ] + return tools ######################################################## diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index c137164ecdb..748ef938cea 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -68,6 +68,7 @@ from .utils import process_slack_alerting_variables if TYPE_CHECKING: from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager + from litellm.proxy.utils import PrismaClient from litellm.router import Router as _Router Router = _Router @@ -1944,6 +1945,69 @@ Model Info: except Exception as e: verbose_proxy_logger.exception("Error sending weekly spend report %s", e) + async def send_user_spend_alerts(self, prisma_client: "PrismaClient | None" = None) -> None: + """Check per-user daily/monthly spend thresholds and spend anomalies, alerting once per user per period.""" + if self.alerting is None or "slack" not in self.alerting: + return + + thresholds_enabled: Final = AlertType.user_spend_thresholds in self.alert_types + anomalies_enabled: Final = AlertType.user_spend_anomalies in self.alert_types + if not thresholds_enabled and not anomalies_enabled: + return + + if prisma_client is None: + from litellm.proxy.proxy_server import prisma_client as global_prisma_client + + prisma_client = global_prisma_client # rebind-ok: fall back to the proxy's global client + if prisma_client is None: + return + + from litellm.integrations.SlackAlerting.user_spend_alerts import ( + evaluate_user_spend, + fetch_user_spend_rows, + ) + + try: + today: Final = datetime.datetime.now(datetime.timezone.utc).date() + rows: Final = await fetch_user_spend_rows( + prisma_client=prisma_client, + today=today, + baseline_days=self.alerting_args.spend_anomaly_baseline_days, + ) + all_events: Final = tuple( + event + for row in rows + for event in evaluate_user_spend( + row=row, + args=self.alerting_args, + today=today, + thresholds_enabled=thresholds_enabled, + anomalies_enabled=anomalies_enabled, + ) + ) + cached_flags: Final = await asyncio.gather( + *(self.internal_usage_cache.async_get_cache(key=event.cache_key) for event in all_events) + ) + new_events: Final = tuple(event for event, cached in zip(all_events, cached_flags) if not cached) + for alert_type in (AlertType.user_spend_thresholds, AlertType.user_spend_anomalies): + typed_events = tuple(event for event in new_events if event.alert_type == alert_type) + if not typed_events: + continue + await self.send_alert( + message="\n\n".join(event.message for event in typed_events), + level="High", + alert_type=alert_type, + alerting_metadata={}, # mutable-ok: send_alert takes a dict payload + ) + for event in typed_events: + await self.internal_usage_cache.async_set_cache( + key=event.cache_key, + value="SENT", + ttl=event.cache_ttl, + ) + except Exception as e: # noqa: BLE001 # background job must not crash the scheduler + verbose_proxy_logger.exception("Error sending user spend alerts: %s", e) + async def send_fallback_stats_from_prometheus(self): """ Helper to send fallback statistics from prometheus server -> to slack diff --git a/litellm/integrations/SlackAlerting/user_spend_alerts.py b/litellm/integrations/SlackAlerting/user_spend_alerts.py new file mode 100644 index 00000000000..38794735c1b --- /dev/null +++ b/litellm/integrations/SlackAlerting/user_spend_alerts.py @@ -0,0 +1,139 @@ +"""Per-user daily/monthly spend threshold alerts and spend anomaly detection.""" + +import datetime +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final, Literal + +from pydantic import TypeAdapter + +from litellm.constants import HOURS_IN_A_DAY +from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +DAY_SECONDS: Final = HOURS_IN_A_DAY * 60 * 60 +MONTHLY_ALERT_TTL_SECONDS: Final = 32 * DAY_SECONDS + +USER_SPEND_QUERY: Final = """ +SELECT + user_id, + COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float AS daily_spend, + COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0)::float AS monthly_spend, + COALESCE(SUM(spend) FILTER (WHERE date >= $3 AND date < $1), 0)::float AS baseline_spend +FROM "LiteLLM_DailyUserSpend" +WHERE date >= LEAST($2, $3) AND user_id IS NOT NULL +GROUP BY user_id +HAVING COALESCE(SUM(spend) FILTER (WHERE date >= $2), 0) > 0 +""" + + +@dataclass(frozen=True, slots=True) +class UserSpendRow: + user_id: str + daily_spend: float + monthly_spend: float + baseline_spend: float + + +@dataclass(frozen=True, slots=True) +class UserSpendAlertEvent: + kind: Literal["daily_threshold", "monthly_threshold", "anomaly"] + alert_type: AlertType + message: str + cache_key: str + cache_ttl: int + + +USER_SPEND_ROWS_ADAPTER: Final = TypeAdapter(tuple[UserSpendRow, ...]) + + +async def fetch_user_spend_rows( + prisma_client: "PrismaClient", + today: datetime.date, + baseline_days: int, +) -> tuple[UserSpendRow, ...]: + today_str: Final = today.strftime("%Y-%m-%d") + month_start_str: Final = today.replace(day=1).strftime("%Y-%m-%d") + baseline_start_str: Final = (today - datetime.timedelta(days=max(baseline_days, 1))).strftime("%Y-%m-%d") + raw: Final = await prisma_client.db.query_raw(USER_SPEND_QUERY, today_str, month_start_str, baseline_start_str) + return USER_SPEND_ROWS_ADAPTER.validate_python(raw) + + +def _daily_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: + threshold: Final = args.daily_spend_per_user_threshold + if threshold is None or row.daily_spend < threshold: + return None + return UserSpendAlertEvent( + kind="daily_threshold", + alert_type=AlertType.user_spend_thresholds, + message=( + f"User Daily Spend Threshold Crossed:\n" + f"User: `{row.user_id}`\n" + f"Spend Today: `${row.daily_spend:.2f}`\n" + f"Daily Threshold: `${threshold:.2f}`" + ), + cache_key=f"user_spend_alert_daily_{row.user_id}_{today_str}", + cache_ttl=DAY_SECONDS, + ) + + +def _monthly_threshold_event(row: UserSpendRow, args: SlackAlertingArgs, month_str: str) -> UserSpendAlertEvent | None: + threshold: Final = args.monthly_spend_per_user_threshold + if threshold is None or row.monthly_spend < threshold: + return None + return UserSpendAlertEvent( + kind="monthly_threshold", + alert_type=AlertType.user_spend_thresholds, + message=( + f"User Monthly Spend Threshold Crossed:\n" + f"User: `{row.user_id}`\n" + f"Spend This Month: `${row.monthly_spend:.2f}`\n" + f"Monthly Threshold: `${threshold:.2f}`" + ), + cache_key=f"user_spend_alert_monthly_{row.user_id}_{month_str}", + cache_ttl=MONTHLY_ALERT_TTL_SECONDS, + ) + + +def _anomaly_event(row: UserSpendRow, args: SlackAlertingArgs, today_str: str) -> UserSpendAlertEvent | None: + if row.daily_spend < args.spend_anomaly_min_spend: + return None + baseline_daily_avg: Final = row.baseline_spend / args.spend_anomaly_baseline_days + if row.baseline_spend > 0 and row.daily_spend <= args.spend_anomaly_multiplier * baseline_daily_avg: + return None + return UserSpendAlertEvent( + kind="anomaly", + alert_type=AlertType.user_spend_anomalies, + message=( + f"User Spend Anomaly Detected:\n" + f"User: `{row.user_id}`\n" + f"Spend Today: `${row.daily_spend:.2f}`\n" + f"Daily Average (last {args.spend_anomaly_baseline_days} days): `${baseline_daily_avg:.2f}`\n" + f"Trigger: spend above `{args.spend_anomaly_multiplier}x` the daily average " + f"(minimum `${args.spend_anomaly_min_spend:.2f}`)" + ), + cache_key=f"user_spend_alert_anomaly_{row.user_id}_{today_str}", + cache_ttl=DAY_SECONDS, + ) + + +def evaluate_user_spend( + row: UserSpendRow, + args: SlackAlertingArgs, + today: datetime.date, + thresholds_enabled: bool, + anomalies_enabled: bool, +) -> tuple[UserSpendAlertEvent, ...]: + today_str: Final = today.strftime("%Y-%m-%d") + month_str: Final = today.strftime("%Y-%m") + threshold_events: Final = ( + ( + _daily_threshold_event(row=row, args=args, today_str=today_str), + _monthly_threshold_event(row=row, args=args, month_str=month_str), + ) + if thresholds_enabled + else () + ) + anomaly_events: Final = (_anomaly_event(row=row, args=args, today_str=today_str),) if anomalies_enabled else () + return tuple(event for event in (*threshold_events, *anomaly_events) if event is not None) diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 207e024ce0b..ce51fb19970 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -127,7 +127,7 @@ def handle_anthropic_text_model_custom_llm_provider( return model, custom_llm_provider -def declared_authenticating_provider(model: str, custom_llm_provider: str | None = None) -> str | None: +def declared_authenticating_provider(model: str | None, custom_llm_provider: str | None = None) -> str | None: """The authenticating provider this pair already names, or None. get_llm_provider runs the OAuth device flow for github_copilot and chatgpt, because their @@ -135,7 +135,7 @@ def declared_authenticating_provider(model: str, custom_llm_provider: str | None and for a declared pair the resolver's answer is the declaration itself, so metadata callers adopt the declaration instead of resolving. """ - declared: Final = custom_llm_provider or model.split("/", 1)[0] + declared: Final = custom_llm_provider or (model.split("/", 1)[0] if model and "/" in model else None) return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index eb223c2f988..9a6fb11f978 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2959,13 +2959,25 @@ class Logging(LiteLLMLoggingBaseClass): "Model=%s not found in completion cost map. Setting 'response_cost' to None", self.model ) self.model_call_details["response_cost"] = None + except Exception: # noqa: BLE001 # cost calculation must never block later callbacks (slot release) + verbose_logger.exception( + "Error calculating streaming response cost for model=%s. Setting 'response_cost' to None", + self.model, + ) + self.model_call_details["response_cost"] = None self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - complete_streaming_response, start_time, end_time - ) + try: + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + complete_streaming_response, start_time, end_time + ) + except Exception: # noqa: BLE001 # payload build must never block later callbacks (slot release) + verbose_logger.exception( + "LiteLLM.LoggingError: [Non-Blocking] Exception building the standard logging payload " + "for a streaming response; callbacks still run without it" + ) # print standard logging payload if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: @@ -3005,32 +3017,39 @@ class Logging(LiteLLMLoggingBaseClass): ## LOGGING HOOK ## for callback in callbacks: - if isinstance(callback, CustomGuardrail): - from litellm.types.guardrails import GuardrailEventHooks + try: + if isinstance(callback, CustomGuardrail): + from litellm.types.guardrails import GuardrailEventHooks - if ( - callback.should_run_guardrail( - data=self.model_call_details, - event_type=GuardrailEventHooks.logging_only, + if ( + callback.should_run_guardrail( + data=self.model_call_details, + event_type=GuardrailEventHooks.logging_only, + ) + is not True + ): + continue + + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, ) - is not True - ): - continue - - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, - ) - elif isinstance(callback, CustomLogger): - result = redact_message_input_output_from_custom_logger( - result=result, litellm_logging_obj=self, custom_logger=callback - ) - self.model_call_details, result = await callback.async_logging_hook( - kwargs=self.model_call_details, - result=result, - call_type=self.call_type, + elif isinstance(callback, CustomLogger): + result = redact_message_input_output_from_custom_logger( + result=result, litellm_logging_obj=self, custom_logger=callback + ) + self.model_call_details, result = await callback.async_logging_hook( + kwargs=self.model_call_details, + result=result, + call_type=self.call_type, + ) + except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release) + verbose_logger.error( + "LiteLLM.LoggingError: [Non-Blocking] Exception occurred in async_logging_hook %s", + traceback.format_exc(), ) + self._handle_callback_failure(callback=callback) self.has_run_logging(event_type="async_success") diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 2fca19dc8d1..ff46440ff5c 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1281,6 +1281,19 @@ def flatten_top_level_schema_combinators(schema: Mapping[str, object]) -> Mappin return _flatten_schema_against_root(schema, schema, frozenset(), 0, {}) # mutable-ok: fresh per-call $ref memo +def tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: + function: Final = tool.get("function") + if not isinstance(function, dict): + return tool + parameters: Final = function.get("parameters") + if not isinstance(parameters, dict): + return tool + flattened: Final = flatten_top_level_schema_combinators(parameters) + if flattened is parameters: + return tool + return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts + + def _get_image_mime_type_from_url(url: str) -> str | None: """ Get mime type for common image URLs diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 45c7825344b..66e36dab2ba 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -86,22 +86,41 @@ def _decoded_sse_data_line(line: bytes) -> object | None: return None -def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: +def _anthropic_event_payload(chunk: object, event_type: str) -> Mapping[str, object] | None: if isinstance(chunk, dict): - return chunk if chunk.get("type") == "error" else None + return chunk if chunk.get("type") == event_type else None if isinstance(chunk, (bytes, bytearray)): decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines()) return next( ( candidate for candidate in decoded_lines - if isinstance(candidate, dict) and candidate.get("type") == "error" + if isinstance(candidate, dict) and candidate.get("type") == event_type ), None, ) return None +def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None: + return _anthropic_event_payload(chunk, "error") + + +def parse_anthropic_refusal_stop_details(chunk: object) -> Mapping[str, object] | None: + """ + Return the ``stop_details`` object of an Anthropic SSE ``message_delta`` + chunk whose delta carries ``stop_reason: "refusal"`` (a safeguard refusal: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other chunk, a plain refusal without ``stop_details`` included. + """ + payload: Final = _anthropic_event_payload(chunk, "message_delta") + delta: Final = payload.get("delta") if payload is not None else None + if not isinstance(delta, dict) or delta.get("stop_reason") != "refusal": + return None + stop_details: Final = delta.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None: """Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None.""" payload: Final = _anthropic_error_event_payload(chunk) diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py index 02d82887dde..9deff950724 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py @@ -1,11 +1,40 @@ +from collections.abc import Mapping from functools import lru_cache -from typing import Any, Final, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Final, cast, get_type_hints from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) +if TYPE_CHECKING: + from litellm.exceptions import ContentPolicyViolationError + + +def get_safeguard_refusal_stop_details(response: object) -> Mapping[str, Any] | None: + """ + Return the ``stop_details`` of an Anthropic Messages response refused by a + safeguard (``stop_reason: "refusal"`` carrying ``stop_details``: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback), + or None for any other response, a plain refusal without ``stop_details`` included. + """ + if not isinstance(response, dict) or response.get("stop_reason") != "refusal": + return None + stop_details: Final = response.get("stop_details") + return stop_details if isinstance(stop_details, dict) else None + + +def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "ContentPolicyViolationError": + """The exception a safeguard-refused Anthropic response converts into so the + content-policy fallback chain can re-dispatch it.""" + from litellm.exceptions import ContentPolicyViolationError + + return ContentPolicyViolationError( + message=f"Anthropic safeguard refusal (category: {stop_details.get('category')}).", + model=model, + llm_provider="anthropic", + ) + @lru_cache(maxsize=1) def _anthropic_messages_optional_param_keys() -> frozenset[str]: @@ -100,14 +129,12 @@ def mock_response( model=model, ) return AnthropicMessagesResponse( - **{ - "content": [{"text": mock_response, "type": "text"}], - "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", - "model": "claude-sonnet-4-20250514", - "role": "assistant", - "stop_reason": "end_turn", - "stop_sequence": None, - "type": "message", - "usage": {"input_tokens": 2095, "output_tokens": 503}, - } + content=[{"text": mock_response, "type": "text"}], + id="msg_013Zva2CMHLNnXjNJJKqJ2EF", + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage={"input_tokens": 2095, "output_tokens": 503}, ) diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 2df4ab731ab..0ac0662205a 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -1,3 +1,5 @@ +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final from httpx._models import Headers, Response @@ -6,6 +8,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, hoist_images_from_tool_messages, + tool_with_flattened_parameters, ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_azure_openai_messages, @@ -32,6 +35,19 @@ else: LoggingClass = Any +_NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) + + +def flattened_tools_update(optional_params: Mapping[str, object]) -> Mapping[str, object]: + tools: Final = optional_params.get("tools") + if not isinstance(tools, list): + return _NO_TOOLS_UPDATE + flattened: Final = [ # mutable-ok: request tools are a JSON list + tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + ] + return MappingProxyType({"tools": flattened}) + + class AzureOpenAIConfig(BaseConfig): """ Reference: https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#chat-completions @@ -261,6 +277,7 @@ class AzureOpenAIConfig(BaseConfig): "model": model, "messages": azure_messages, **optional_params, + **flattened_tools_update(optional_params), } def transform_response( diff --git a/litellm/llms/azure/chat/o_series_transformation.py b/litellm/llms/azure/chat/o_series_transformation.py index 6cbd91bab5d..246bf69cb5f 100644 --- a/litellm/llms/azure/chat/o_series_transformation.py +++ b/litellm/llms/azure/chat/o_series_transformation.py @@ -20,6 +20,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.utils import get_model_info, supports_reasoning from ...openai.chat.o_series_transformation import OpenAIOSeriesConfig +from .gpt_transformation import flattened_tools_update class AzureOpenAIO1Config(OpenAIOSeriesConfig): @@ -108,4 +109,8 @@ class AzureOpenAIO1Config(OpenAIOSeriesConfig): headers: dict, ) -> dict: model = model.replace("o_series/", "") # handle o_series/my-random-deployment-name - return super().transform_request(model, messages, optional_params, litellm_params, headers) + flattened_params: Final = { # mutable-ok: transform_request's contract takes a plain JSON params dict + **optional_params, + **flattened_tools_update(optional_params), + } + return super().transform_request(model, messages, flattened_params, litellm_params, headers) diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index f3e41c054ef..9afc6331d96 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -20,9 +20,9 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo ) from litellm.litellm_core_utils.prompt_templates.common_utils import ( drop_tool_reference_parts_from_tool_messages, - flatten_top_level_schema_combinators, get_tool_call_names, hoist_images_from_tool_messages, + tool_with_flattened_parameters, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, @@ -70,19 +70,6 @@ else: _NO_TOOLS_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) -def _tool_with_flattened_parameters(tool: Mapping[str, object]) -> Mapping[str, object]: - function: Final = tool.get("function") - if not isinstance(function, dict): - return tool - parameters: Final = function.get("parameters") - if not isinstance(parameters, dict): - return tool - flattened: Final = flatten_top_level_schema_combinators(parameters) - if flattened is parameters: - return tool - return {**tool, "function": {**function, "parameters": flattened}} # mutable-ok: request tools are JSON dicts - - class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): """ Reference: https://platform.openai.com/docs/api-reference/chat/create @@ -466,7 +453,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): ): return _NO_TOOLS_UPDATE flattened: Final = [ # mutable-ok: request tools are a JSON list - _tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools + tool_with_flattened_parameters(tool) if isinstance(tool, dict) else tool for tool in tools ] return MappingProxyType({"tools": flattened}) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 24dbd7c08d0..ec70fe0d795 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -82,6 +82,10 @@ class ResponsesStreamChunk(TypedDict, total=False): type: ReadOnly[str] text: ReadOnly[str] + delta: ReadOnly[str] + item_id: ReadOnly[str] + output_index: ReadOnly[int] + content_index: ReadOnly[int] def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: @@ -658,8 +662,32 @@ class OpenAIResponsesHandler(BaseTranslation): def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. + + ``response.output_text.done`` events carry the whole part in ``text``, while + ``response.output_text.delta`` events carry fragments in ``delta``. A stream + that dies before its done event (``response.failed`` / ``response.incomplete``) + has text only in deltas, so per content part the done text wins when present + and the joined deltas fill in otherwise, never both. """ - return "".join([response.get("text", "") for response in responses_so_far]) + keyed_events: Final = tuple( + ( + (event.get("item_id"), event.get("output_index"), event.get("content_index")), + event.get("text"), + event.get("delta"), + ) + for event in responses_so_far + if isinstance(event.get("text"), str) or isinstance(event.get("delta"), str) + ) + + def part_text(part_key: tuple[object, object, object]) -> str: + done_texts: Final = tuple( + text for key, text, _ in keyed_events if key == part_key and isinstance(text, str) + ) + if done_texts: + return done_texts[-1] + return "".join(delta for key, _, delta in keyed_events if key == part_key and isinstance(delta, str)) + + return "".join(part_text(key) for key in dict.fromkeys(key for key, _, _ in keyed_events)) def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 3efb6429326..2b89dba0e4f 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -4,10 +4,12 @@ from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal +import anyio import httpx from fastapi import APIRouter, Depends, HTTPException, Query, Request, status from litellm._logging import verbose_logger +from litellm.constants import MCP_CLIENT_TIMEOUT, MCP_TOOL_LISTING_TIMEOUT from litellm.exceptions import ( BlockedPiiEntityError, GuardrailRaisedException, @@ -86,8 +88,6 @@ def _connection_error_message(exc: BaseException) -> str: if MCP_AVAILABLE: - from mcp.types import Tool as MCPTool - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( _UPSTREAM_OAUTH_DISCOVERY_AUTH_TYPES, global_mcp_server_manager, @@ -1173,6 +1173,7 @@ if MCP_AVAILABLE: transport=request.transport, auth_type=request.auth_type, mcp_info=request.mcp_info, + timeout=request.timeout, command=request.command, args=request.args, env=request.env, @@ -1402,11 +1403,28 @@ if MCP_AVAILABLE: oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers) async def _list_tools_operation(client): - async def _list_tools_session_operation(session): - return await session.list_tools() - - list_tools_response: Final = await client.run_with_session(_list_tools_session_operation) - list_tools_result: Final[list[MCPTool]] = list_tools_response.tools + # Bound the whole pagination walk: without this the preview is limited only by the + # per-request timeout times the page cap. max() keeps the pre-pagination guarantee + # that a single slow page within the client timeout still succeeds, and a + # per-server timeout above the global default extends the deadline with it. + listing_deadline: Final = max( + getattr(client, "timeout", MCP_CLIENT_TIMEOUT) or MCP_CLIENT_TIMEOUT, + MCP_TOOL_LISTING_TIMEOUT, + ) + list_tools_result = None # rebind-ok: set inside the timeout scope below + with anyio.move_on_after(listing_deadline): + list_tools_result = await client.list_tools(raise_on_error=True) # rebind-ok: fills the init above + if list_tools_result is None: + verbose_logger.warning( + "MCP tools/list preview timed out after %s seconds while paginating upstream tools", + listing_deadline, + ) + return { # mutable-ok: error response payload + "status": "error", + "error": True, + "message": f"Timed out listing tools after {listing_deadline} seconds. " + "The MCP server may be responding slowly or paginating excessively.", + } model_dumped_tools: Final[list[dict]] = [tool.model_dump() for tool in list_tools_result] return { "tools": model_dumped_tools, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e0a2097919b..18714256a8f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1216,6 +1216,13 @@ class GenerateKeyRequest(KeyRequestBase): organization_id: str | None = None project_id: str | None = None + @field_validator("team_id", mode="before") + @classmethod + def treat_cleared_team_id_as_unset(cls, v: object) -> object: + if v == "": + return None + return v + class GenerateKeyResponse(KeyRequestBase): key: str diff --git a/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py new file mode 100644 index 00000000000..7da5e5099fc --- /dev/null +++ b/litellm/proxy/anthropic_endpoints/streaming_model_restamp.py @@ -0,0 +1,174 @@ +""" +Restamp the public ``model`` on the Anthropic Messages ``message_start`` event, the only +stream event carrying a model, so streamed responses report the requested model like +non-streaming ones do. + +Chunks reach the serializer either as already-encoded SSE frames (``bytes``/``str``, the +provider passthrough path) or as event dicts (fake-stream and agentic paths). +""" + +import json +import re +from collections.abc import Mapping +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +_MESSAGE_START_EVENT: Final = "message_start" +_MESSAGE_START_MARKER: Final = b"message_start" +_SSE_DATA_FIELD: Final = "data:" +_SSE_FRAME_END_PATTERN: Final = re.compile(rb"\r\n\r\n|\r\r|\n\n") +_MAX_HELD_BYTES: Final = 65536 +_PING_MARKERS: Final = (b"event: ping", b'"type": "ping"', b'"type":"ping"') + +_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _restamped_event(event: Mapping[str, object], requested_model: str) -> Mapping[str, object] | None: + message: Final = event.get("message") + if event.get("type") != _MESSAGE_START_EVENT or not isinstance(message, dict): + return None + if message.get("model") == requested_model: + return None + return {**event, "message": {**message, "model": requested_model}} # mutable-ok: SSE payload, re-serialized as is + + +def _restamped_data_line(line: str, requested_model: str) -> str | None: + stripped: Final = line.strip() + if not stripped.startswith(_SSE_DATA_FIELD): + return None + payload: Final = stripped[len(_SSE_DATA_FIELD) :].strip() + if not payload or payload == "[DONE]": + return None + try: + event: Final = _EVENT_ADAPTER.validate_json(payload) + except ValidationError: + return None + restamped: Final = _restamped_event(event, requested_model) + if restamped is None: + return None + terminator: Final = line[len(line.rstrip("\r\n")) :] + return f"data: {json.dumps(restamped, separators=(',', ':'))}{terminator}" + + +def _restamped_frame(frame: str, requested_model: str) -> str | None: + lines: Final = frame.splitlines(keepends=True) + restamped: Final = tuple(_restamped_data_line(line, requested_model) for line in lines) + if all(line is None for line in restamped): + return None + return "".join(new if new is not None else old for new, old in zip(restamped, lines)) + + +def restamp_anthropic_stream_chunk_model(chunk: object, requested_model: str) -> object: + """ + Return ``chunk`` with the ``message_start`` model replaced by ``requested_model``. + + Chunks that carry no model are returned unchanged. + """ + if isinstance(chunk, dict): + try: + event: Final = _EVENT_ADAPTER.validate_python(chunk) + except ValidationError: + return chunk + return _restamped_event(event, requested_model) or chunk + + if isinstance(chunk, (bytes, bytearray)): + if _MESSAGE_START_EVENT.encode() not in chunk: + return chunk + restamped_bytes: Final = _restamped_frame(chunk.decode("utf-8", errors="ignore"), requested_model) + return chunk if restamped_bytes is None else restamped_bytes.encode("utf-8") + + if isinstance(chunk, str): + if _MESSAGE_START_EVENT not in chunk: + return chunk + restamped_text: Final = _restamped_frame(chunk, requested_model) + return chunk if restamped_text is None else restamped_text + + return chunk + + +def _is_ping_frame(frame: bytes) -> bool: + return any(marker in frame for marker in _PING_MARKERS) + + +class AnthropicStreamModelRestamper: + """ + Per-stream restamper for the encoded passthrough path, where chunks are raw + transport reads: the ``message_start`` SSE frame can arrive split across + chunks or coalesced with later frames. Complete frames (``\\n\\n``, + ``\\r\\n\\r\\n``, or ``\\r\\r`` terminated) are emitted as their terminator + closes them and an incomplete tail is held until it completes, so the + restamp never misses a torn frame; ``flush`` returns whatever is still held + when the stream ends so no bytes are swallowed. Once ``message_start`` has + been handled, or the first real event proves the stream carries none, every + later chunk passes through untouched. + """ + + def __init__(self, requested_model: str) -> None: + self._requested_model: Final = requested_model + self._held = b"" + self._armed = True + + def process(self, chunk: object) -> object: + if not self._armed: + return chunk + if isinstance(chunk, (bytes, bytearray)): + return self._process_encoded(bytes(chunk)) + if isinstance(chunk, str): + return self._process_encoded(chunk.encode("utf-8")) + restamped: Final = restamp_anthropic_stream_chunk_model(chunk, self._requested_model) + if isinstance(chunk, dict) and chunk.get("type") not in (None, "ping"): + self._armed = False + return restamped + + def flush(self) -> bytes: + held: Final = self._held + self._held = b"" + self._armed = False + if not held: + return b"" + restamped: Final = restamp_anthropic_stream_chunk_model(held, self._requested_model) + return restamped if isinstance(restamped, bytes) else held + + def _process_encoded(self, data: bytes) -> bytes: + combined: Final = self._held + data + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(combined)) + if not boundaries: + if len(combined) > _MAX_HELD_BYTES: + self._held = b"" + self._armed = False + return combined + self._held = combined + return b"" + emitted: Final = self._restamped_closed_block(combined[: boundaries[-1]]) + tail: Final = combined[boundaries[-1] :] + if not self._armed: + self._held = b"" + return emitted + tail + self._held = tail + return emitted + + def _restamped_closed_block(self, closed: bytes) -> bytes: + boundaries: Final = tuple(match.end() for match in _SSE_FRAME_END_PATTERN.finditer(closed)) + frames: Final = tuple(closed[start:end] for start, end in zip((0, *boundaries[:-1]), boundaries)) + decider: Final = next( + ( + index + for index, frame in enumerate(frames) + if _MESSAGE_START_MARKER in frame or (b"data:" in frame and not _is_ping_frame(frame)) + ), + None, + ) + if decider is None: + return closed + self._armed = False + if _MESSAGE_START_MARKER not in frames[decider]: + return closed + restamped_text: Final = _restamped_frame( + frames[decider].decode("utf-8", errors="ignore"), self._requested_model + ) + if restamped_text is None: + return closed + return b"".join( + restamped_text.encode("utf-8") if index == decider else frame for index, frame in enumerate(frames) + ) diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 45e05d353fb..c591cbabee1 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -9,6 +9,7 @@ import click import requests from .auth import context_secret_vault, get_stored_api_key, login +from .cmd_quoting import quote_for_cmd ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" @@ -151,31 +152,9 @@ def verify_proxy_key( _WINDOWS_SHIM_SUFFIXES: Final[frozenset[str]] = frozenset({".cmd", ".bat"}) -_CMD_PERCENT_GUARD: Final = "%%cd:~,%" _CMD_LINE_BREAKS: Final = ("\r", "\n") -def _double_trailing_backslashes(segment: str) -> str: - bare: Final = segment.rstrip("\\") - return bare + "\\" * 2 * (len(segment) - len(bare)) - - -def _quote_for_cmd(token: str) -> str: - """Quote one token so both parsers that read it see the original text. - - Follows the algorithm the Rust standard library settled on for batch files - after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a - quoted string on a lone `"` and so wants an embedded one doubled, and the - shim's own interpreter, which re-splits `%*` under C runtime rules where a - backslash escapes the quote that follows it, so every backslash run standing - before a quote is doubled. Quoting cannot stop cmd expanding `%VAR%`, so each - `%` is prefixed with `%%cd:~,`: the zero-length substring of the always - defined `cd` expands to nothing and leaves no `%` pair for cmd to match. - """ - escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"')) - return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"' - - def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: """Build what CreateProcess runs, routing batch shims through cmd.exe. @@ -202,7 +181,7 @@ def _windows_command(path: str, args: Sequence[str]) -> str | tuple[str, ...]: f"Cannot pass an argument containing a line break to `{os.path.basename(path)}` on " "Windows: cmd.exe ends the command line there, so the agent would silently lose it." ) - inner: Final = " ".join(_quote_for_cmd(token) for token in (path, *rest)) + inner: Final = " ".join(quote_for_cmd(token) for token in (path, *rest)) return f'cmd.exe /d /e:on /v:off /s /c "{inner}"' diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index ea1fa019c83..46af641636e 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -8,6 +8,7 @@ live here rather than in either command module. import shlex import shutil +import sys from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -17,6 +18,8 @@ from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.private_json import write_private_json +from .cmd_quoting import quote_for_cmd + ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" @@ -87,9 +90,12 @@ def merge_claude_settings( return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} -def resolve_api_key_helper(base_url: str) -> str: +def resolve_api_key_helper(base_url: str, platform: str = sys.platform) -> str: """Build the shell command Claude Code should run for its apiKeyHelper. + Claude Code hands the string to the system shell, `sh` on POSIX and cmd.exe + on Windows, so every token is quoted for the shell that will read it. + Resolves `lite` to an absolute path so the helper works regardless of the PATH visible to whatever subprocess Claude Code spawns it from. Passing --base-url explicitly (rather than relying on the bare invocation Claude @@ -106,7 +112,8 @@ def resolve_api_key_helper(base_url: str) -> str: raise ClaudeSettingsError( "Could not find `lite` on your PATH. Claude Code's apiKeyHelper needs an absolute path to it." ) - return f"{shlex.quote(lite_path)} --base-url {shlex.quote(base_url)} auth print-token" + quote: Final = quote_for_cmd if platform.startswith("win") else shlex.quote + return " ".join(quote(token) for token in (lite_path, "--base-url", base_url, "auth", "print-token")) def write_claude_settings(base_url: str, settings_path: Path, owners: Sequence[SettingsFileOwner]) -> None: diff --git a/litellm/proxy/client/cli/commands/cmd_quoting.py b/litellm/proxy/client/cli/commands/cmd_quoting.py new file mode 100644 index 00000000000..efd6d584527 --- /dev/null +++ b/litellm/proxy/client/cli/commands/cmd_quoting.py @@ -0,0 +1,26 @@ +"""Quoting for command lines that cmd.exe reads before handing them to a program.""" + +from typing import Final + +_CMD_PERCENT_GUARD: Final = "%%cd:~,%" + + +def _double_trailing_backslashes(segment: str) -> str: + bare: Final = segment.rstrip("\\") + return bare + "\\" * 2 * (len(segment) - len(bare)) + + +def quote_for_cmd(token: str) -> str: + """Quote one token so both parsers that read it see the original text. + + Follows the algorithm the Rust standard library settled on for batch files + after CVE-2024-24576. Two parsers see this token: cmd.exe, which ends a + quoted string on a lone `"` and so wants an embedded one doubled, and the + program's own C runtime argv split, where a backslash escapes the quote that + follows it, so every backslash run standing before a quote is doubled. + Quoting cannot stop cmd expanding `%VAR%`, so each `%` is prefixed with + `%%cd:~,`: the zero-length substring of the always defined `cd` expands to + nothing and leaves no `%` pair for cmd to match. + """ + escaped: Final = '""'.join(_double_trailing_backslashes(part) for part in token.split('"')) + return '"' + escaped.replace("%", _CMD_PERCENT_GUARD) + '"' diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1d860b02875..05ddef822f1 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -176,6 +176,9 @@ if TYPE_CHECKING: ProxyConfig = _ProxyConfig else: ProxyConfig = Any +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, +) from litellm.proxy.litellm_pre_call_utils import ( add_litellm_data_to_request, refresh_proxy_server_request_body_snapshot, @@ -2490,6 +2493,9 @@ class ProxyBaseLLMRequestProcessing: request_data=self.data, proxy_logging_obj=proxy_logging_obj, request=request, + restamp_model=( + None if _should_return_raw_model_name(self.data) else requested_model_from_client + ), ) return await create_response( generator=wrap_sse_stream_with_keepalive_pings( @@ -3442,6 +3448,16 @@ class ProxyBaseLLMRequestProcessing: else: return chunk + @staticmethod + def _sse_chunk_serializer(restamper: AnthropicStreamModelRestamper | None) -> StreamChunkSerializer: + if restamper is None: + return ProxyBaseLLMRequestProcessing.return_sse_chunk + + def serialize(chunk: object) -> str: + return ProxyBaseLLMRequestProcessing.return_sse_chunk(restamper.process(chunk)) + + return serialize + @staticmethod async def _finalize_streaming_generator_cleanup( request: Request | None, @@ -3502,11 +3518,16 @@ class ProxyBaseLLMRequestProcessing: serialize_chunk: StreamChunkSerializer, serialize_error: StreamErrorSerializer, request: Request | None = None, + flush_tail: Callable[[], bytes] | None = None, ) -> AsyncGenerator[str, None]: """ Shared streaming data generator: runs proxy iterator hook, per-chunk hook, cost injection, then yields chunks via serialize_chunk; on exception runs failure hook and yields via serialize_error. Use for SSE or NDJSON. + + ``flush_tail`` runs once after the upstream iterator completes cleanly and + its non-empty result is yielded, so a serializer that buffers bytes across + chunks can emit anything still held at end of stream. """ verbose_proxy_logger.debug("inside generator") # Resolve per-stream (not per-chunk) whether the heavy per-chunk path @@ -3569,6 +3590,9 @@ class ProxyBaseLLMRequestProcessing: # so it must not suppress that refund. delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES yield serialize_chunk(chunk) + held_tail: Final = flush_tail() if flush_tail is not None else b"" + if held_tail: + yield serialize_chunk(held_tail) stream_completed = True except (asyncio.CancelledError, GeneratorExit): # Client disconnected mid-stream. CancelledError / GeneratorExit @@ -3579,8 +3603,7 @@ class ProxyBaseLLMRequestProcessing: # billing and release exactly once. This is the outermost generator # Starlette closes on disconnect, so the nested iterator hook (which # only sees GeneratorExit on GC) cannot own the refund. - if not stream_completed: - client_disconnected = True + client_disconnected = not stream_completed if not delivered_chunk and not _withheld_provider_output(response): from litellm.proxy.spend_tracking.budget_reservation import ( release_budget_reservation_on_cancel, @@ -3634,6 +3657,7 @@ class ProxyBaseLLMRequestProcessing: request_data: dict, proxy_logging_obj: ProxyLogging, request: Request | None = None, + restamp_model: str | None = None, ) -> AsyncGenerator[str, None]: """ Anthropic /messages and Google /generateContent streaming data generator require SSE events. @@ -3642,17 +3666,23 @@ class ProxyBaseLLMRequestProcessing: SSE serializers directly (rather than re-wrapping it in another ``async for: yield`` trampoline), so a streamed chunk traverses one fewer async-generator layer / coroutine resume on the hot path. + + ``restamp_model`` publishes that name on the Anthropic ``message_start`` + event in place of the provider's model, matching what the non-streaming + response reports. """ + restamper: Final = AnthropicStreamModelRestamper(restamp_model) if restamp_model else None return ProxyBaseLLMRequestProcessing.async_streaming_data_generator( response=response, user_api_key_dict=user_api_key_dict, request_data=request_data, proxy_logging_obj=proxy_logging_obj, - serialize_chunk=ProxyBaseLLMRequestProcessing.return_sse_chunk, + serialize_chunk=ProxyBaseLLMRequestProcessing._sse_chunk_serializer(restamper), serialize_error=lambda proxy_exc: ( f"{STREAM_SSE_DATA_PREFIX}{json.dumps({'error': proxy_exc.to_dict()})}\n\n" ), request=request, + flush_tail=None if restamper is None else restamper.flush, ) @overload diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0237d82a0d9..30526d30dc5 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -31,6 +31,7 @@ from litellm.caching import DualCache from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys from litellm.litellm_core_utils.litellm_logging import ( _get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name @@ -215,6 +216,16 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]: return redacted if isinstance(redacted, list) else assessments +_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses}) + + +def _is_responses_api_route(request_route: str | None) -> bool: + if request_route is None: + return False + call_types: Final = get_call_types_for_route(request_route) + return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types) + + class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL. @@ -2709,6 +2720,24 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): yield streamed_chunk return + # Responses-API events are neither chat-completions chunks nor raw + # Anthropic SSE, so the assembly below cannot scan them; the unified + # guardrail's translation layer can, with buffering semantics kept. + if _is_responses_api_route(user_api_key_dict.request_route): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=True, + ): + yield translated_chunk + return + # Import here to avoid circular imports from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70bf8fd0554..77a80ea0052 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -39,7 +39,7 @@ from typing import ( import anyio import websockets import websockets.exceptions -from pydantic import BaseModel, Json, JsonValue +from pydantic import BaseModel, Json, JsonValue, ValidationError from typing_extensions import NotRequired, ReadOnly, assert_never from litellm._uuid import uuid @@ -253,6 +253,7 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + USER_SPEND_ALERTS_JOB_ID, WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError @@ -9866,6 +9867,35 @@ class ProxyStartupEvent: replace_existing=True, ) + slack_alerting_args: Final = proxy_logging_obj.slack_alerting_instance.alerting_args + user_spend_check_interval: Final = ( + slack_alerting_args.user_spend_check_interval + if isinstance(slack_alerting_args, SlackAlertingArgs) # pyright: ignore[reportUnnecessaryIsInstance] # tests inject a mock slack_alerting_instance + else SlackAlertingArgs().user_spend_check_interval + ) + + async def _scheduled_user_spend_alerts() -> None: + if ( + await pod_lock_manager.acquire_lock( + cronjob_id=USER_SPEND_ALERTS_JOB_ID, + ttl=max(user_spend_check_interval - 60, 60), + allow_reentrant=False, + ) + is False + ): + return + await proxy_logging_obj.slack_alerting_instance.send_user_spend_alerts() + + scheduler.add_job( + _scheduled_user_spend_alerts, + "interval", + seconds=user_spend_check_interval, + next_run_time=datetime.now(timezone.utc) + timedelta(seconds=10 + random.randint(0, 60)), + id=USER_SPEND_ALERTS_JOB_ID, + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + if os.getenv("PROMETHEUS_URL"): from zoneinfo import ZoneInfo @@ -12481,11 +12511,21 @@ async def supported_openai_params(model: str): --header 'Authorization: Bearer sk-1234' ``` """ + from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider + + global llm_router try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) + resolved_models: Final = llm_router.resolved_litellm_models(model) if llm_router is not None else () + target_model: Final = resolved_models[0] if resolved_models else model + declared_provider: Final = declared_authenticating_provider(target_model) + litellm_model, custom_llm_provider = ( + (target_model.removeprefix(f"{declared_provider}/"), declared_provider) + if declared_provider is not None + else litellm.get_llm_provider(model=target_model)[:2] + ) return { "supported_openai_params": litellm.get_supported_openai_params( - model=model, custom_llm_provider=custom_llm_provider + model=litellm_model, custom_llm_provider=custom_llm_provider ) } except Exception: @@ -14962,17 +15002,25 @@ async def alerting_settings( alerting_args_dict = {} alerting_values = None - allowed_args: Final = { - "slack_alerting": {"type": "Boolean"}, - "daily_report_frequency": {"type": "Integer"}, - "report_check_interval": {"type": "Integer"}, - "budget_alert_ttl": {"type": "Integer"}, - "outage_alert_ttl": {"type": "Integer"}, - "region_outage_alert_ttl": {"type": "Integer"}, - "minor_outage_alert_threshold": {"type": "Integer"}, - "major_outage_alert_threshold": {"type": "Integer"}, - "max_outage_alert_list_size": {"type": "Integer"}, - } + allowed_args: Final = MappingProxyType( + { + "slack_alerting": "Boolean", + "daily_report_frequency": "Integer", + "report_check_interval": "Integer", + "budget_alert_ttl": "Integer", + "outage_alert_ttl": "Integer", + "region_outage_alert_ttl": "Integer", + "minor_outage_alert_threshold": "Integer", + "major_outage_alert_threshold": "Integer", + "max_outage_alert_list_size": "Integer", + "daily_spend_per_user_threshold": "Float", + "monthly_spend_per_user_threshold": "Float", + "spend_anomaly_multiplier": "Float", + "spend_anomaly_baseline_days": "Integer", + "spend_anomaly_min_spend": "Float", + "user_spend_check_interval": "Integer", + } + ) _slack_alerting: Final[SlackAlerting] = proxy_logging_obj.slack_alerting_instance _slack_alerting_args_dict: Final = _slack_alerting.alerting_args.model_dump() @@ -14987,7 +15035,7 @@ async def alerting_settings( _response_obj = ConfigList( field_name="slack_alerting", - field_type=allowed_args["slack_alerting"]["type"], + field_type=allowed_args["slack_alerting"], field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.", field_value=is_slack_enabled, stored_in_db=True if alerting_values is not None else False, @@ -15006,7 +15054,7 @@ async def alerting_settings( _response_obj = ConfigList( field_name=field_name, - field_type=allowed_args[field_name]["type"], + field_type=allowed_args[field_name], field_description=field_info.description or "", field_value=_slack_alerting_args_dict.get(field_name, None), stored_in_db=_stored_in_db, @@ -16434,6 +16482,16 @@ async def update_config_general_settings( detail={"error": f"Invalid type of field value={type(data.field_value)} passed in."}, ) + if data.field_name == "alerting_args": + try: + SlackAlertingArgs.model_validate(data.field_value) + except ValidationError as e: + errors: Final = "; ".join(f"{'.'.join(str(loc) for loc in err['loc'])}: {err['msg']}" for err in e.errors()) + raise HTTPException( + status_code=400, + detail={"error": f"Invalid alerting_args: {errors}"}, + ) + ## get general settings from db db_general_settings: Final = await _config_param_table(prisma_client).find_first( where={"param_name": "general_settings"} diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index f4d8fd7d906..91d2ece7a51 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -120,13 +120,15 @@ async def _apply_over_budget_reservation_policy( applied_entries: list[dict[str, float | str]], reservation_cost: float, current_spend: float, + fail_closed_budget_enforcement: bool = False, ) -> float: """ Decide what to do when a counter is over budget, and return the reservation cost to carry into the next counter. Three outcomes: an over-budget key that opted into throttling releases its own reservation (the rate limiter slows it) and keeps the cost; a partially-remaining budget resizes the reservation - down to what is left; anything else hard-blocks by raising. + down to what is left, unless strict enforcement is on, because the known + estimate already does not fit; anything else hard-blocks by raising. """ if _key_reservation_should_release_for_throttle(counter.counter_key, valid_token): await _release_applied_entries_best_effort(entries=[entry], default_reserved_cost=reservation_cost) @@ -134,21 +136,36 @@ async def _apply_over_budget_reservation_policy( return reservation_cost remaining_before_reservation: Final = counter.max_budget - (current_spend - reservation_cost) - if remaining_before_reservation > 1e-12: - await _resize_applied_reservation( - entries=applied_entries, - current_reserved_cost=reservation_cost, - new_reserved_cost=remaining_before_reservation, + if remaining_before_reservation <= 1e-12: + _raise_counter_budget_exceeded(counter=counter, current_cost=current_spend) + if fail_closed_budget_enforcement and current_spend - counter.max_budget > 1e-12: + _raise_counter_budget_exceeded( + counter=counter, + current_cost=current_spend - reservation_cost, + estimated_cost=reservation_cost, ) - return remaining_before_reservation + await _resize_applied_reservation( + entries=applied_entries, + current_reserved_cost=reservation_cost, + new_reserved_cost=remaining_before_reservation, + ) + return remaining_before_reservation + +def _raise_counter_budget_exceeded( + counter: _BudgetCounter, + current_cost: float, + estimated_cost: float | None = None, +) -> NoReturn: + estimate_detail: Final = "" if estimated_cost is None else f"Estimated request cost: {estimated_cost}, " raise litellm.BudgetExceededError( - current_cost=current_spend, + current_cost=current_cost, max_budget=counter.max_budget, message=( "Budget has been exceeded! " f"{counter.entity_type}={counter.entity_id} " - f"Current cost: {current_spend}, " + f"Current cost: {current_cost}, " + f"{estimate_detail}" f"Max budget: {counter.max_budget}" ), entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type), @@ -258,6 +275,7 @@ async def reserve_budget_for_request( applied_entries=applied_entries, reservation_cost=reservation_cost, current_spend=current_spend, + fail_closed_budget_enforcement=fail_closed_budget_enforcement, ) continue except Exception: diff --git a/litellm/router.py b/litellm/router.py index 471a1116f44..6e4405ebfef 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -143,7 +143,11 @@ from litellm.router_utils.cooldown_handlers import ( from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _check_non_standard_fallback_format, - get_fallback_model_group, + clear_pre_routing_selection, + fallback_lookup_groups, + get_fallback_model_group_for_lookup_groups, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) from litellm.router_utils.get_retry_from_policy import ( @@ -4918,6 +4922,19 @@ class Router: ) response = await response + if self._should_raise_anthropic_refusal_error( + model=model, + original_generic_function=original_generic_function, + response=response, + kwargs=kwargs, + ): + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, + ) + + refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape + raise safeguard_refusal_error(model=model, stop_details=refusal_details) + self.success_calls[model_name] += 1 verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) @@ -4964,6 +4981,11 @@ class Router: # fallback to the original reference for any non-picklable value. # The original_generic_function is preserved so the per-attempt # helper knows which underlying API to call on fallback. + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -4973,6 +4995,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): return await self._aresponses_streaming_iterator( response=response, @@ -5030,6 +5060,10 @@ class Router: from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( aclose_if_supported, parse_anthropic_error_event, + parse_anthropic_refusal_stop_details, + ) + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + safeguard_refusal_error, ) source_iterator: Final = response @@ -5068,13 +5102,35 @@ class Router: continue if _anthropic_stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): has_generated_content = True # rebind-ok: real content seen, or the buffer cap was hit - error_event = parse_anthropic_error_event(chunk) + # A transport can split one SSE data line across byte chunks, so pre-content + # detection parses the accumulated buffer plus the current chunk, never the + # chunk alone; the buffer is already capped, which bounds this window too. + parse_window = ( # rebind-ok: freshly computed each iteration, never carried over + b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime + else chunk + ) + error_event = parse_anthropic_error_event(parse_window) retriable_pending_error = ( # rebind-ok: freshly computed each iteration, never carried over not has_generated_content and error_event is not None and _is_retriable_anthropic_status(error_event[2]) and not _anthropic_stream_error_is_gateway_verdict(chunk) ) + refusal_stop_details = ( # rebind-ok: freshly computed each iteration, never carried over + parse_anthropic_refusal_stop_details(parse_window) + if not has_generated_content and error_event is None + else None + ) + if refusal_stop_details is not None and self._has_content_policy_fallback(model, initial_kwargs): + refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) + raise MidStreamFallbackError( + message=refusal_error.message, + model=model, + llm_provider="anthropic", + original_exception=refusal_error, + is_pre_first_chunk=True, + ) if not has_generated_content and not retriable_pending_error and error_event is None: buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) continue @@ -5186,8 +5242,13 @@ class Router: kwargs=initial_kwargs, metadata_variable_name="litellm_metadata", ) + # The content-policy dispatch branch matches on the trigger's own type, so a refusal's + # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. + fallback_trigger: Final[Exception] = ( + e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success - e=e, + e=fallback_trigger, disable_fallbacks=False, fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -5243,6 +5304,11 @@ class Router: # share, leaking primary-deployment metadata into the mid-stream # fallback request. safe_deep_copy avoids deep-copying the full # kwargs (which can hold non-deepcopyable logging handles/clients). + # The pre-routing hook stamps its tier selection into this bucket during the primary + # attempt; seeding it before the snapshot gives both the live kwargs and the copy a + # bucket, so the post-call carry-over below always has somewhere to read and write. + kwargs.setdefault("litellm_metadata", {}) # mutable-ok: shared bucket # rebind-ok: stamp must be readable here + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() # mutable-ok: mutated below before re-entry if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) @@ -5252,6 +5318,14 @@ class Router: response: Final = await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + # The snapshot predates the pre-routing hook, so the tier it stamped into the live kwargs + # is carried over write-or-clear: a stale or caller-supplied selection left in the copy + # would key the mid-stream fallback lookup off a tier this attempt never routed to. + clear_pre_routing_selection(fallback_kwargs) + live_pre_routing_selection: Final = get_pre_routing_selection(kwargs) + if live_pre_routing_selection is not None: + record_pre_routing_selection(fallback_kwargs, live_pre_routing_selection) + if kwargs.get("stream") and hasattr(response, "__aiter__"): return await self._aanthropic_messages_streaming_iterator( response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator @@ -6807,6 +6881,9 @@ class Router: original_exception: Final = e fallback_model_group = None original_model_group: Final[str | None] = kwargs.get("model") + # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier + # behind the router name, and fallbacks are configured per tier, not per router. + lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) fallback_failure_exception_str = "" if disable_fallbacks is True or original_model_group is None: @@ -6851,15 +6928,15 @@ class Router: ] # Get external fallbacks — handle both standard and non-standard formats external_fallback_group: list | None = None - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: if _check_non_standard_fallback_format(fallbacks=fallbacks): # Non-standard formats (e.g. ["claude-3-haiku"] or # [{"model": "...", "messages": [...]}]) are passed through directly external_fallback_group = fallbacks else: - external_fallback_group, generic_idx = get_fallback_model_group( + external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) if external_fallback_group is None and generic_idx is not None: external_fallback_group = fallbacks[generic_idx]["*"] @@ -6917,9 +6994,9 @@ class Router: if isinstance(e, litellm.ContextWindowExceededError): if context_window_fallbacks is not None: context_window_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=context_window_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if context_window_fallback_model_group is None: @@ -6950,9 +7027,9 @@ class Router: elif isinstance(e, litellm.ContentPolicyViolationError): if content_policy_fallbacks is not None: content_policy_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_from_fallbacks( + self._get_fallback_model_group_for_lookup_groups( fallbacks=content_policy_fallbacks, - model_group=model_group, + lookup_groups=lookup_groups, ) ) if content_policy_fallback_model_group is None: @@ -6979,14 +7056,14 @@ class Router: if litellm.expose_router_debug_in_errors: e.message += f"\n{error_message}" - if fallbacks is not None and model_group is not None: + if fallbacks is not None and lookup_groups: verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks)) ( fallback_model_group, generic_fallback_idx, - ) = get_fallback_model_group( + ) = get_fallback_model_group_for_lookup_groups( fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - model_group=cast(str, model_group), + lookup_groups=lookup_groups, ) ## if none, check for generic fallback if fallback_model_group is None and generic_fallback_idx is not None: @@ -6995,12 +7072,12 @@ class Router: if fallback_model_group is None: masked_fallbacks: Final = mask_sensitive_structure(fallbacks) verbose_router_logger.info( - "No fallback model group found for original model_group=%s. Fallbacks=%s", - model_group, + "No fallback model group found for lookup_groups=%s. Fallbacks=%s", + " -> ".join(lookup_groups), masked_fallbacks, ) if hasattr(original_exception, "message") and litellm.expose_router_debug_in_errors: - original_exception.message += f"No fallback model group found for original model_group={model_group}. Fallbacks={masked_fallbacks}" + original_exception.message += f"No fallback model group found for lookup_groups={' -> '.join(lookup_groups)}. Fallbacks={masked_fallbacks}" raise original_exception input_kwargs.update( @@ -7046,6 +7123,7 @@ class Router: If it fails after num_retries, fall back to another model group """ model_group: Final[str | None] = kwargs.get("model") + clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): _fallback_metadata_key: Final = _get_router_metadata_variable_name( function_name=getattr(kwargs.get("original_function"), "__name__", None) @@ -7471,6 +7549,24 @@ class Router: break return fallback_model_group + def _get_fallback_model_group_for_lookup_groups( + self, + fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the sibling resolver's contract + lookup_groups: tuple[str, ...], + ) -> list[str] | None: # mutable-ok: mirrors the sibling resolver's contract + """First lookup group whose exact-key chain resolves (tier first, then requested group).""" + return next( + ( + resolved + for resolved in ( + self._get_fallback_model_group_from_fallbacks(fallbacks=fallbacks, model_group=group) + for group in lookup_groups + ) + if resolved is not None + ), + None, + ) + def _get_first_default_fallback(self) -> str | None: """ Returns the first model from the default_fallbacks list, if it exists. @@ -7886,6 +7982,31 @@ class Router: return True return False + def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: + """ + Whether a content-policy fallback would resolve for this request, keyed the same way + async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook + selected wins over the requested group. Raising without this returning True would turn + a deliverable response into an error the fallback chain cannot recover from. + """ + content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + if content_policy_fallbacks is not None: + return ( + self._get_fallback_model_group_for_lookup_groups( + fallbacks=content_policy_fallbacks, + lookup_groups=fallback_lookup_groups(kwargs, model_group), + ) + is not None + ) + if self._has_default_fallbacks(): + return True + verbose_router_logger.debug( + "No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s", + model_group, + content_policy_fallbacks, + ) + return False + def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: """ Determines if a content policy error should be raised. @@ -7898,27 +8019,26 @@ class Router: if response.choices[0].finish_reason != "content_filter": return False - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self.content_policy_fallbacks) + return self._has_content_policy_fallback(model, kwargs) - ### ONLY RAISE ERROR IF CP FALLBACK AVAILABLE ### - if content_policy_fallbacks is not None: - fallback_model_group = None - for item in content_policy_fallbacks: # [{"gpt-3.5-turbo": ["gpt-4"]}] - if list(item.keys())[0] == model: - fallback_model_group = item[model] - break - - if fallback_model_group is not None: - return True - elif self._has_default_fallbacks(): # default fallbacks set - return True - - verbose_router_logger.debug( - "Content Policy Error occurred. No available fallbacks. Returning original response. model=%s, content_policy_fallbacks=%s", - model, - content_policy_fallbacks, + def _should_raise_anthropic_refusal_error( + self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, Any] + ) -> bool: + """ + The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard + refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only + when a content-policy fallback is configured; a plain refusal without stop_details, or + any response with nothing configured, is returned to the client unchanged. + """ + from litellm.llms.anthropic.experimental_pass_through.messages.utils import ( + get_safeguard_refusal_stop_details, ) - return False + + if getattr(original_generic_function, "__name__", "") != "anthropic_messages": + return False + if get_safeguard_refusal_stop_details(response) is None: + return False + return self._has_content_policy_fallback(model, kwargs) def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): _all_deployments: list = [] @@ -12087,6 +12207,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs @@ -12202,6 +12323,7 @@ class Router: if pre_routing_hook_response is not None: model = pre_routing_hook_response.model messages = pre_routing_hook_response.messages + record_pre_routing_selection(request_kwargs, model) if pre_routing_hook_response.litellm_params: accepted_tier_params: Final = self._tier_params_the_target_accepts( model, pre_routing_hook_response.litellm_params, request_kwargs diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 924574537f3..3d37ca216a7 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -214,6 +214,91 @@ def _check_stripped_model_group(model_group: str, fallback_key: str) -> bool: return False +PRE_ROUTING_SELECTED_MODEL_KEY: Final = "pre_routing_selected_model" +_ROUTER_METADATA_BUCKETS: Final = ("metadata", "litellm_metadata") + + +def record_pre_routing_selection(request_kwargs: Mapping[str, Any] | None, selected_model: str) -> None: + """ + Remember which model a pre-routing hook picked, so fallback lookup can key off it. + + Fallback resolution runs on an outer kwargs dict that ``**kwargs`` already copied, so + writing the model there is invisible by the time routing picks a tier. The metadata + buckets are nested dicts shared by reference across those copies, which is how the + router already carries values back up. + + The write goes through the proxy-internal bucket resolver, never into both buckets: + on /v1/messages the top-level ``metadata`` dict is the provider's own request field, + so a blanket write would forward the tier stamp upstream. + """ + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if request_kwargs is None: + return + bucket: Final = request_kwargs.get(get_metadata_variable_name_from_kwargs(request_kwargs)) + if isinstance(bucket, dict): + bucket[PRE_ROUTING_SELECTED_MODEL_KEY] = selected_model + + +def clear_pre_routing_selection(request_kwargs: Mapping[str, object] | None) -> None: + """ + Drop any selection the router did not make itself on this hop. + + The buckets carry whatever the caller sent, so an inbound value is the caller + choosing a fallback chain rather than the router choosing a tier. A fallback hop + also inherits the previous hop's selection, which would key its own failure off + the tier that already failed. Clearing at the start of every hop leaves only a + value the pre-routing hook wrote while routing that hop. + """ + if request_kwargs is None: + return + for bucket in (request_kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS): + if isinstance(bucket, dict) and PRE_ROUTING_SELECTED_MODEL_KEY in bucket: + del bucket[PRE_ROUTING_SELECTED_MODEL_KEY] + + +def get_pre_routing_selection(kwargs: Mapping[str, Any]) -> str | None: + """The model a pre-routing hook selected for this request, if one did.""" + buckets: Final = (kwargs.get(name) for name in _ROUTER_METADATA_BUCKETS) + selections: Final = (bucket.get(PRE_ROUTING_SELECTED_MODEL_KEY) for bucket in buckets if isinstance(bucket, dict)) + return next((selected for selected in selections if isinstance(selected, str) and selected), None) + + +def fallback_lookup_groups(kwargs: Mapping[str, Any], model_group: str | None) -> tuple[str, ...]: + """ + Ordered keys for resolving a fallback chain: the tier a pre-routing hook selected wins, + and the requested group still resolves when no tier-keyed chain exists, so configs keyed + on the router name (the documented contract) keep working behind auto-routers. + """ + ordered: Final = (get_pre_routing_selection(kwargs), model_group) + return tuple(dict.fromkeys(group for group in ordered if group)) + + +def _resolved_a_specific_chain( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + result: tuple[list[str] | None, int | None], # mutable-ok: mirrors get_fallback_model_group's contract +) -> bool: + resolved, generic_idx = result + if resolved is None: + return False + return generic_idx is None or resolved is not fallbacks[generic_idx]["*"] + + +def get_fallback_model_group_for_lookup_groups( + fallbacks: list[Any], # mutable-ok: mirrors get_fallback_model_group's contract + lookup_groups: tuple[str, ...], +) -> tuple[list[str] | None, int | None]: # mutable-ok: mirrors get_fallback_model_group's contract + """ + First lookup group with a specifically-keyed chain wins; the generic "*" chain applies + only after every group missed, so a catch-all cannot shadow a later group's own chain. + """ + results: Final = tuple(get_fallback_model_group(fallbacks=fallbacks, model_group=group) for group in lookup_groups) + specific: Final = next((result for result in results if _resolved_a_specific_chain(fallbacks, result)), None) + if specific is not None: + return specific + return next((result for result in results if result[0] is not None), (None, None)) + + def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[list[str] | None, int | None]: """ Returns: @@ -412,6 +497,7 @@ async def run_async_fallback( # LOGGING kwargs = litellm_router.log_retry(kwargs=kwargs, e=original_exception) verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(mg)) + kwargs.pop("_target_order", None) # rebind-ok: next hop must not inherit the previous order target if isinstance(mg, str): kwargs["model"] = mg elif isinstance(mg, dict): diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 0d5ef01bc04..0775e0a4039 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -8,7 +8,7 @@ from re import Match from typing import Final from litellm._logging import verbose_router_logger -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider, get_llm_provider class PatternUtils: @@ -204,7 +204,7 @@ class PatternMatchRouter: return litellm_deployment_litellm_model - def get_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict] | None: + def get_pattern(self, model: str | None, custom_llm_provider: str | None = None) -> list[dict] | None: """ Check if a pattern exists for the given model and custom llm provider @@ -215,18 +215,17 @@ class PatternMatchRouter: Returns: bool: True if pattern exists, False otherwise """ - if custom_llm_provider is None: - try: - ( - _, - custom_llm_provider, - _, - _, - ) = get_llm_provider(model=model) - except Exception: - # get_llm_provider raises exception when provider is unknown - pass - return self.route(model) or self.route(f"{custom_llm_provider}/{model}") + provider: Final = ( + custom_llm_provider or declared_authenticating_provider(model) or self._resolved_provider(model) + ) + return self.route(model) or self.route(f"{provider}/{model}") + + @staticmethod + def _resolved_provider(model: str | None) -> str | None: + try: + return get_llm_provider(model=model)[1] if model else None + except Exception: # noqa: BLE001 # get_llm_provider raises when the provider is unknown; the name then routes as-is + return None def get_deployments_by_pattern(self, model: str, custom_llm_provider: str | None = None) -> list[dict]: """ diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 7fb90ab89de..b1e9dbdefa8 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -427,6 +427,8 @@ class DeploymentAffinityCheck(CustomLogger): """ request_kwargs = request_kwargs or {} typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) + if request_kwargs.get("_target_order") is not None: + return typed_healthy_deployments ( enable_user_key, diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 6e8406b2ec7..0788c8db710 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -58,6 +58,9 @@ class PromptCachingDeploymentCheck(CustomLogger): request_kwargs: dict | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: + if request_kwargs is not None and request_kwargs.get("_target_order") is not None: + return healthy_deployments + if messages is not None and is_prompt_caching_valid_prompt( messages=messages, model=model, diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index b1b7bc3541a..64c0c530e9b 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -91,6 +91,40 @@ class SlackAlertingArgs(LiteLLMPydanticObjectBase): default=False, description="If true, the alerting payload will be printed to the console.", ) + daily_spend_per_user_threshold: float | None = Field( + default=None, + gt=0, + allow_inf_nan=False, + description="Alert when a user's spend for the current day (UTC) crosses this USD amount. Off by default.", + ) + monthly_spend_per_user_threshold: float | None = Field( + default=None, + gt=0, + allow_inf_nan=False, + description="Alert when a user's spend for the current calendar month (UTC) crosses this USD amount. Off by default.", + ) + spend_anomaly_multiplier: float = Field( + default=3.0, + gt=0, + allow_inf_nan=False, + description="Flag a user's spend as anomalous when today's spend exceeds this multiple of their trailing daily average.", + ) + spend_anomaly_baseline_days: int = Field( + default=7, + ge=1, + description="Number of trailing days used to compute a user's daily average spend for anomaly detection.", + ) + spend_anomaly_min_spend: float = Field( + default=10.0, + gt=0, + allow_inf_nan=False, + description="Minimum spend (USD) a user must reach today before an anomaly alert can fire. Reduces false positives.", + ) + user_spend_check_interval: int = Field( + default=3600, + ge=60, + description="How often (in seconds) to check per-user spend thresholds and anomalies. Default is hourly.", + ) class DeploymentMetrics(LiteLLMPydanticObjectBase): @@ -138,6 +172,8 @@ class AlertType(str, Enum): budget_alerts = "budget_alerts" spend_reports = "spend_reports" failed_tracking_spend = "failed_tracking_spend" + user_spend_thresholds = "user_spend_thresholds" + user_spend_anomalies = "user_spend_anomalies" # Database alerts db_exceptions = "db_exceptions" @@ -182,6 +218,7 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [ AlertType.budget_alerts, AlertType.spend_reports, AlertType.failed_tracking_spend, + AlertType.user_spend_thresholds, # Database alerts AlertType.db_exceptions, # Report alerts diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 42ca3fd6d4b..4fe1dafc73b 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -78,6 +78,16 @@ class AnthropicUsage(TypedDict, total=False): server_tool_use: NotRequired[ReadOnly[ServerToolUsage]] +class AnthropicStopDetails(TypedDict, total=False): + """ + Safeguard verdict accompanying a `stop_reason: "refusal"` response: + https://platform.claude.com/docs/en/build-with-claude/refusals-and-fallback + """ + + category: ReadOnly[str | None] + explanation: ReadOnly[str | None] + + class AnthropicMessagesResponse(TypedDict, total=False): """ Anthropic Messages API Response: https://docs.anthropic.com/en/api/messages @@ -90,7 +100,8 @@ class AnthropicMessagesResponse(TypedDict, total=False): id: str model: str | None # This represents the Model type from Anthropic role: Literal["assistant"] | None - stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use"] | None + stop_reason: Literal["end_turn", "max_tokens", "stop_sequence", "tool_use", "refusal"] | None + stop_details: NotRequired[ReadOnly[AnthropicStopDetails | None]] stop_sequence: str | None type: Literal["message"] | None usage: AnthropicUsage | None diff --git a/litellm/utils.py b/litellm/utils.py index 2fdc4a45148..252b6756937 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4876,11 +4876,7 @@ def _get_deployment_order(deployment: dict | Any) -> int | None: def _get_order_filtered_deployments(healthy_deployments: list[dict], target_order: int | None = None) -> list: if target_order is not None: - filtered: Final = [d for d in healthy_deployments if _get_deployment_order(d) == target_order] - if filtered: - return filtered - # target_order doesn't match any deployment (e.g., external fallback model) — return all - return healthy_deployments + return [d for d in healthy_deployments if _get_deployment_order(d) == target_order] # Default: pick min order group _valid_orders: Final[list[int]] = [ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index c541c035db7..a5e00799519 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -81,6 +81,7 @@ ignored_function_names = [ "_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name) "_invalidate_access_groups_cache", # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name) "has_buffered_provider_output", # Property, so its reads in test_router.py are never an ast.Call + "_resolved_provider", # Tested via get_pattern in test_pattern_match_deployments.py (file lacks "router" in name) ] diff --git a/tests/e2e/junit_properties.py b/tests/e2e/junit_properties.py index e4f59f5c4d2..c5971c5362c 100644 --- a/tests/e2e/junit_properties.py +++ b/tests/e2e/junit_properties.py @@ -2,10 +2,16 @@ The e2e suite ships results to Loki/Grafana from a standard pytest JUnit report (`--junitxml=e2e-report.xml`), not a bespoke log line. JUnit already records -outcome, duration, and node id for every ``; the only signals it cannot -derive on its own are the normalized suite package and the coverage-registry cell -ids a test covers. Those ride along as JUnit `` entries via each item's -`user_properties`, attached in `conftest.py::pytest_collection_modifyitems`. +outcome, duration, and node id for every ``; the signals it cannot +derive on its own are the normalized suite package, the coverage-registry cell +ids a test covers, and where the test's source lives. Those ride along as JUnit +`` entries via each item's `user_properties`, attached in +`conftest.py::pytest_collection_modifyitems`. + +`source` is a property rather than the `file=` / `line=` attributes pytest used +to write, because the `xunit2` family this suite runs on drops those, and +switching families would change the XML for every consumer of it -- the +Buildkite Test Engine upload and the Loki pipeline included. """ from __future__ import annotations @@ -14,22 +20,61 @@ from collections.abc import Iterable import pytest +# Hardcoded because the runner image copies tests/e2e/ to /app/e2e, so nothing +# at runtime names this suite's place in the repo. test_junit_properties.py +# fails from a checkout if it moves. +SUITE_ROOT = "tests/e2e" + + +def suite_parts(path_part: str) -> tuple[str, ...]: + """Path components of a suite file relative to tests/e2e, however it ran. + + Pytest paths are rootdir-relative, and rootdir moves with the invocation: a + repo-root run gives `tests/e2e/logging/test_x.py`, a suite-cwd run (the + runner image) gives `logging/test_x.py`. Both collapse to the same tuple. + """ + raw = tuple(p for p in path_part.replace("\\", "/").split("/") if p and p != ".") + return raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + def package_from_nodeid(nodeid: str) -> str: - """Top-level suite package under tests/e2e/, or 'root' for top-level files. - - Pytest nodeids are relative to the invocation cwd. Repo-root runs look like - `tests/e2e/logging/...`; suite-cwd runs look like `logging/...`. Strip the - `tests/e2e` prefix so package is the suite dir either way. - """ - path_part = nodeid.split("::", 1)[0].replace("\\", "/") - raw = tuple(p for p in path_part.split("/") if p and p != ".") - parts = raw[2:] if len(raw) >= 3 and raw[0] == "tests" and raw[1] == "e2e" else raw + """Top-level suite package under tests/e2e/, or 'root' for top-level files.""" + parts = suite_parts(nodeid.split("::", 1)[0]) if len(parts) <= 1: return "root" return parts[0] +def source_from_location(path: str, lineno: int | None) -> str: + """Repo-relative `path:line` for a test, or '' when nothing is linkable. + + `pytest.Item.location` gives a rootdir-relative path and a ZERO-based line. + The path is re-rooted at SUITE_ROOT so consumers need not know how pytest was + started, and the line is emitted ONE-based to match editors, tracebacks and + code hosts. A decorated test anchors at its first decorator, which is where + pytest reports it. + + Empty rather than a guess for anything unlinkable: no line, a path reaching + upward, or a path carrying a colon, which is both how an absolute Windows + path arrives and a character `path:line` has no way to represent. + """ + if lineno is None: + return "" + normalized = path.replace("\\", "/") + if normalized.startswith("/") or ":" in normalized or ".." in normalized.split("/"): + return "" + parts = suite_parts(normalized) + if not parts: + return "" + return f"{'/'.join((SUITE_ROOT, *parts))}:{lineno + 1}" + + +def source_from_item(item: pytest.Item) -> str: + """Read the repo-relative `path:line` off a pytest Item's reported location.""" + path, lineno, _ = item.location + return source_from_location(path, lineno) + + def dedupe_covers(marker_args: Iterable[tuple[object, ...]]) -> tuple[str, ...]: """Flatten @pytest.mark.covers arg lists into unique, order-preserving cell ids, dropping anything that is not a non-empty string.""" @@ -43,10 +88,12 @@ def covers_from_item(item: pytest.Item) -> tuple[str, ...]: def result_properties(item: pytest.Item) -> tuple[tuple[str, str], ...]: """The custom signals a standard reporter cannot derive: the normalized suite - package and the comma-joined coverage-registry cell ids this test covers.""" + package, the comma-joined coverage-registry cell ids this test covers, and the + repo-relative `path:line` its source sits at.""" return ( ("package", package_from_nodeid(item.nodeid)), ("covers", ",".join(covers_from_item(item))), + ("source", source_from_item(item)), ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 967144dfb16..79d9e011f7e 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -806,6 +806,7 @@ class LiteLLMParamsBody(BaseModel): mock_response: str | None = None timeout: float | None = None tpm: int | None = None + weight: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] @@ -820,6 +821,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails_policy: dict[str, int] | None = None class ModelNewBody(BaseModel): diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index e342aa363ca..5822058003c 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -19,6 +19,8 @@ from models import ( ChatMessage, ChatResponse, LiteLLMParamsBody, + ModelInfoBody, + ModelNewBody, ReliabilityChatBody, RouterSettingsOverride, ) @@ -26,6 +28,18 @@ from models import ( REAL_MODEL = "openai/gpt-5.5" REAL_KEY = "os.environ/OPENAI_API_KEY" +# The smallest-context chat model OpenAI still serves (16385 tokens). A prompt +# past that limit comes back as a real `context_length_exceeded` 400, which is +# what litellm maps to ContextWindowExceededError. +SMALL_CONTEXT_MODEL = "openai/gpt-3.5-turbo" +SMALL_CONTEXT_LIMIT_TOKENS = 16385 + + +def oversized_prompt(marker: str) -> str: + """A prompt comfortably past SMALL_CONTEXT_MODEL's context limit, so the + provider refuses it on length rather than answering a truncated version.""" + return f"{marker} " + ("token " * (SMALL_CONTEXT_LIMIT_TOKENS + 4000)) + def create_bad_base_deployment(proxy: ProxyClient, name: str) -> str: """Register a deployment pointing at an unreachable base, so every call to it @@ -40,6 +54,38 @@ def create_timeout_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model(name, LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001)) +def create_small_context_deployment(proxy: ProxyClient, name: str) -> str: + """Register a deployment on the smallest-context model OpenAI still serves, so an + oversized prompt earns a real context-window refusal from the provider.""" + return proxy.create_model(name, LiteLLMParamsBody(model=SMALL_CONTEXT_MODEL, api_key=REAL_KEY)) + + +def create_always_timing_out_deployment(proxy: ProxyClient, name: str) -> str: + """The always-picked half of a retry pair: a 1ms deadline the backend always + exceeds, all of the model group's shuffle weight, and a cooldown policy that + benches it on its first Timeout so the retry cannot land on it again.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, timeout=0.001, weight=1), + model_info=ModelInfoBody(allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + ) + ) + + +def create_zero_weight_backup_deployment(proxy: ProxyClient, name: str) -> str: + """The other half of a retry pair: healthy, but weight 0, so the weighted shuffle + never opens on it. It is reachable only once its sibling is benched and the + weighted pick falls through to a uniform one over what is left.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody(model=REAL_MODEL, api_key=REAL_KEY, weight=0), + model_info=ModelInfoBody(), + ) + ) + + def chat_override( proxy: ProxyClient, key: str, diff --git a/tests/e2e/router/test_reliability_fallbacks_e2e.py b/tests/e2e/router/test_reliability_fallbacks_e2e.py index d3ce62f8f95..8cece41ce2d 100644 --- a/tests/e2e/router/test_reliability_fallbacks_e2e.py +++ b/tests/e2e/router/test_reliability_fallbacks_e2e.py @@ -9,6 +9,10 @@ in the x-litellm-attempted-fallbacks header. Empty content is accepted only when `finish_reason == "length"` and the response billed completion tokens, since gpt-5.5 counts reasoning against max_tokens and can consume the whole budget before emitting any text; a fallback that produced nothing at all still fails. + +The context-window case is a different reroute from a plain failure: the provider +refuses the prompt on length, and `context_window_fallbacks` is the setting that +reroutes it, not `fallbacks`. """ from __future__ import annotations @@ -25,8 +29,10 @@ from reliability_support import ( completion_tokens_of, content_of, create_bad_base_deployment, + create_small_context_deployment, create_timeout_deployment, finish_reason_of, + oversized_prompt, reasoning_tokens_of, ) @@ -82,3 +88,17 @@ class TestReliabilityFallbacks: override=RouterSettingsOverride(fallbacks=[{primary: ["gpt-5.5"]}]), ) _assert_served_by_fallback(resp) + + @pytest.mark.covers("reliability.fallback.context_window.routes_to_fallback") + def test_context_window_routes_to_fallback( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + primary = f"reliability-ctxfail-{unique_marker()}" + model_id = create_small_context_deployment(client.proxy, primary) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + resp = chat_override( + client.proxy, scoped_key, primary, oversized_prompt(unique_marker()), + override=RouterSettingsOverride(context_window_fallbacks=[{primary: ["gpt-5.5"]}]), + ) + _assert_served_by_fallback(resp) diff --git a/tests/e2e/router/test_reliability_retries_e2e.py b/tests/e2e/router/test_reliability_retries_e2e.py new file mode 100644 index 00000000000..5441412935c --- /dev/null +++ b/tests/e2e/router/test_reliability_retries_e2e.py @@ -0,0 +1,73 @@ +"""Live e2e: a request that fails on its first deployment is retried inside its own +model group and still comes back a completion. + +The model group is a pair: an always-timing-out deployment that holds all of the +group's shuffle weight, and a healthy backup at weight 0. The weighted pick always +opens on the timing-out one, its first Timeout benches it (an +`allowed_fails_policy` of `TimeoutErrorAllowedFails: 0`), and the retry falls +through to the only deployment left. So the customer sees a completion and the +proxy reports that it took a retry to get there, with no random first pick in the +middle of it. +""" + +from __future__ import annotations + +import pytest + +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from lifecycle import ResourceManager +from models import RouterSettingsOverride +from reliability_support import ( + chat_override, + completion_tokens_of, + content_of, + create_always_timing_out_deployment, + create_zero_weight_backup_deployment, + finish_reason_of, +) + +pytestmark = pytest.mark.e2e + + +class TestReliabilityRetries: + @pytest.mark.covers("reliability.retry.timeout.succeeds_within_retries") + def test_timeout_on_first_deployment_succeeds_on_retry( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-retry-{unique_marker()}" + timing_out = create_always_timing_out_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(timing_out)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + resp = chat_override( + client.proxy, + scoped_key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=2), + ) + + assert resp.status_code == 200, ( + f"the retry should have landed on the healthy backup, got {resp.status_code}: {resp.body[:300]}" + ) + + attempted = resp.headers.get("x-litellm-attempted-retries") + assert attempted is not None, "response is missing the x-litellm-attempted-retries header" + assert int(attempted) >= 1, ( + f"x-litellm-attempted-retries is {attempted!r}; a 200 with no retry means the request never " + "opened on the timing-out deployment, so this proves nothing about retries" + ) + + content = content_of(resp) + finish_reason = finish_reason_of(resp) + completion_tokens = completion_tokens_of(resp) or 0 + assert isinstance(content, str), ( + f"the retry should have returned a completion body, got content {content!r} (body={resp.body[:300]})" + ) + assert content or (finish_reason == "length" and completion_tokens > 0), ( + f"the retry returned empty content with finish_reason={finish_reason!r}, " + f"completion_tokens={completion_tokens}; empty content is only acceptable when the budget " + f"was spent on non-visible reasoning (body={resp.body[:300]})" + ) diff --git a/tests/e2e/test_junit_properties.py b/tests/e2e/test_junit_properties.py new file mode 100644 index 00000000000..c0596177cc1 --- /dev/null +++ b/tests/e2e/test_junit_properties.py @@ -0,0 +1,146 @@ +"""Harness coverage for the custom JUnit properties. + +No proxy and no ``e2e`` marker. Pins the two normalizations that have to agree +about where a suite file lives -- ``package_from_nodeid`` (strip the suite root) +and ``source_from_location`` (re-root at it) -- across both ways the suite is +launched, plus the one-based line offset and the refusal to emit a path that +escapes the suite. The consumers of these properties are the Loki/Grafana +rollups and, for ``source``, the status page's per-test links to GitHub. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from junit_properties import ( + SUITE_ROOT, + attach_result_properties, + dedupe_covers, + package_from_nodeid, + result_properties, + source_from_location, + suite_parts, +) + + +class FakeMarker: + def __init__(self, name: str, *args: object) -> None: + self.name = name + self.args = args + + +class FakeItem: + """The three attributes junit_properties reads off a pytest Item.""" + + def __init__( + self, nodeid: str, location: tuple[str, int | None, str], markers: tuple[FakeMarker, ...] = () + ) -> None: + self.nodeid = nodeid + self.location = location + self.user_properties: list[tuple[str, str]] = [] + self._markers = markers + + def iter_markers(self, name: str): + return (marker for marker in self._markers if marker.name == name) + + +def repo_root() -> Path | None: + """The litellm checkout above this file, or None when there isn't one.""" + return next((p for p in Path(__file__).resolve().parents if (p / ".git").exists()), None) + + +class TestSuiteParts: + @pytest.mark.parametrize( + "path", + ["logging/test_x.py", "tests/e2e/logging/test_x.py", "./logging/test_x.py", "tests\\e2e\\logging\\test_x.py"], + ) + def test_both_invocation_shapes_collapse_to_the_same_components(self, path: str) -> None: + """A repo-root run and a suite-cwd run report the same file differently; + every downstream signal has to see one spelling.""" + assert suite_parts(path) == ("logging", "test_x.py") + + def test_top_level_suite_file_keeps_its_single_component(self) -> None: + assert suite_parts("tests/e2e/test_fixture_mode.py") == ("test_fixture_mode.py",) + + +class TestPackageFromNodeid: + @pytest.mark.parametrize( + ("nodeid", "expected"), + [ + ("logging/test_x.py::TestFoo::test_bar", "logging"), + ("tests/e2e/logging/test_x.py::TestFoo::test_bar", "logging"), + ("quota_management/spend_tracking/test_x.py::test_bar", "quota_management"), + ("test_fixture_mode.py::TestParseFixtureMode::test_known_values_normalize", "root"), + ("tests/e2e/test_fixture_mode.py::test_bar", "root"), + ], + ) + def test_package_is_the_first_dir_under_the_suite_root(self, nodeid: str, expected: str) -> None: + assert package_from_nodeid(nodeid) == expected + + +class TestSourceFromLocation: + @pytest.mark.parametrize("path", ["a2a/test_a2a_agent_e2e.py", "tests/e2e/a2a/test_a2a_agent_e2e.py"]) + def test_path_is_repo_relative_however_pytest_was_started(self, path: str) -> None: + assert source_from_location(path, 40) == "tests/e2e/a2a/test_a2a_agent_e2e.py:41" + + def test_line_is_emitted_one_based(self) -> None: + """pytest.Item.location counts from 0; editors, tracebacks and GitHub's + #L anchor all count from 1, and an off-by-one lands on the decorator.""" + assert source_from_location("a2a/test_x.py", 0) == "tests/e2e/a2a/test_x.py:1" + + def test_top_level_suite_file_sits_directly_under_the_suite_root(self) -> None: + assert source_from_location("test_fixture_mode.py", 39) == "tests/e2e/test_fixture_mode.py:40" + + @pytest.mark.parametrize( + ("path", "lineno"), + [ + ("a2a/test_x.py", None), + ("/app/e2e/a2a/test_x.py", 40), + ("C:\\app\\e2e\\a2a\\test_x.py", 40), + ("../conftest.py", 40), + ("", 40), + ], + ) + def test_nothing_linkable_yields_empty_rather_than_a_guess(self, path: str, lineno: int | None) -> None: + """A colon is rejected on two counts: it is how a Windows absolute path + arrives, and `path:line` cannot represent one in the path half.""" + assert source_from_location(path, lineno) == "" + + +class TestResultProperties: + def test_every_test_carries_package_covers_and_source(self) -> None: + item = FakeItem( + "logging/test_x.py::TestFoo::test_bar", + ("logging/test_x.py", 40, "TestFoo.test_bar"), + (FakeMarker("covers", "LOG-1", "LOG-2"),), + ) + assert result_properties(item) == ( + ("package", "logging"), + ("covers", "LOG-1,LOG-2"), + ("source", "tests/e2e/logging/test_x.py:41"), + ) + + def test_attach_is_idempotent(self) -> None: + """Collection can run the hook more than once; a second pass must not + double the entries in the report.""" + item = FakeItem("logging/test_x.py::test_bar", ("logging/test_x.py", 40, "test_bar")) + attach_result_properties(item) + attach_result_properties(item) + assert [name for name, _ in item.user_properties] == ["package", "covers", "source"] + + +class TestSuiteRoot: + def test_suite_root_names_this_file_s_real_home(self) -> None: + """SUITE_ROOT is hardcoded because the runner image has no repo to read it + from. Where there IS a checkout, prove the constant still points at us -- + otherwise a moved tests/e2e/ ships links that 404.""" + root = repo_root() + if root is None: + pytest.skip("no checkout above this file (the runner image copies tests/e2e/ to /app/e2e)") + assert (root / SUITE_ROOT / Path(__file__).name).resolve() == Path(__file__).resolve() + + +class TestDedupeCovers: + def test_ids_are_unique_order_preserving_and_non_empty_strings(self) -> None: + assert dedupe_covers([("A", "B"), ("B", ""), ("C", 7)]) == ("A", "B", "C") diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index aadaadd510e..6438525706a 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -11,7 +11,9 @@ from unittest.mock import AsyncMock, MagicMock, patch, ANY import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient from litellm.types.mcp import MCPAuth, MCPTransport -from mcp.types import Tool as MCPTool, CallToolResult as MCPCallToolResult +from mcp.types import CallToolResult as MCPCallToolResult +from mcp.types import ListToolsResult, PaginatedRequestParams +from mcp.types import Tool as MCPTool def test_mcp_client_uses_configurable_default_timeout(): @@ -185,6 +187,80 @@ class TestMCPClientUnitTests: mock_session_instance.initialize.assert_called_once() mock_session_instance.list_tools.assert_called_once() + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + @patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + async def test_list_tools_follows_next_cursor_until_exhausted( + self, + mock_session_class, + mock_transport, + ): + """Test listing tools follows MCP pagination cursors until exhausted.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + first_page_tools = [ + MCPTool(name=f"tool_{idx}", description=f"Tool {idx}", inputSchema={}) for idx in range(100) + ] + second_page_tool = MCPTool( + name="tool_100", + description="Tool 100", + inputSchema={}, + ) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult(tools=first_page_tools, nextCursor="page-2"), + ListToolsResult(tools=[second_page_tool]), + ] + + client = MCPClient("http://example.com") + result = await client.list_tools() + + assert result == [*first_page_tools, second_page_tool] + assert mock_session_instance.list_tools.call_count == 2 + second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + + @pytest.mark.asyncio + @patch.object(mcp_client_module, "streamable_http_client") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + @patch.object(mcp_client_module, "ClientSession") # test-quality-ok: exercises MCPClient wiring; the walk itself is covered sessionless in test_tools.py + async def test_list_tools_swallows_mid_walk_error_without_raise_on_error( + self, + mock_session_class, + mock_transport, + ): + """Test a mid-walk failure returns [] when raise_on_error is False.""" + mock_transport_ctx = AsyncMock() + mock_transport.return_value = mock_transport_ctx + mock_transport_instance = MagicMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=mock_transport_instance) + + mock_session_ctx = AsyncMock() + mock_session_class.return_value = mock_session_ctx + mock_session_instance = AsyncMock() + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="Tool 0", inputSchema={})], + nextCursor="page-2", + ), + RuntimeError("transient upstream failure"), + ] + + client = MCPClient("http://example.com") + result = await client.list_tools() + + assert result == [] + assert mock_session_instance.list_tools.call_count == 2 + @pytest.mark.asyncio @patch.object(mcp_client_module, "streamable_http_client") @patch.object(mcp_client_module, "ClientSession") diff --git a/tests/router_unit_tests/test_router_anthropic_messages_fallback.py b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py new file mode 100644 index 00000000000..0c4d1dfc21e --- /dev/null +++ b/tests/router_unit_tests/test_router_anthropic_messages_fallback.py @@ -0,0 +1,402 @@ +""" +Unit tests for safeguard-refusal fallback on the /v1/messages router surface. + +An Anthropic safeguard refusal is an HTTP 200 whose body carries +stop_reason "refusal" plus a stop_details object; the router converts it +into a ContentPolicyViolationError so the content-policy fallback chain +runs, but only when a matching fallback is configured. A plain refusal +without stop_details, or any refusal with nothing configured, must reach +the client byte-identical. + +The upstream is faked at the HTTP boundary by intercepting the third-party +transport (httpx.AsyncClient.send), so requests run litellm's real +transformation, allowlist, and streaming pipeline end to end. +""" + +import json +from typing import Any, AsyncIterator +from unittest.mock import patch + +import httpx +import pytest + +from litellm import Router +from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + record_pre_routing_selection, +) + +REFUSAL_RESPONSE: dict[str, Any] = { + "id": "msg_refusal", + "type": "message", + "role": "assistant", + "model": "claude-fable-5", + "content": [], + "stop_reason": "refusal", + "stop_sequence": None, + "stop_details": {"category": "cyber", "explanation": "flagged"}, + "usage": {"input_tokens": 25, "output_tokens": 1}, +} + +PLAIN_REFUSAL_RESPONSE: dict[str, Any] = {k: v for k, v in REFUSAL_RESPONSE.items() if k != "stop_details"} + +OK_RESPONSE: dict[str, Any] = { + "id": "msg_ok", + "type": "message", + "role": "assistant", + "model": "claude-opus-5", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 25, "output_tokens": 2}, +} + + +def _sse(event: str, data: dict[str, Any]) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +REFUSAL_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**REFUSAL_RESPONSE, "stop_reason": None}}), + _sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "refusal", "stop_details": {"category": "cyber"}}, + "usage": {"output_tokens": 1}, + }, + ), + _sse("message_stop", {"type": "message_stop"}), +) + +OK_STREAM_FRAMES: tuple[bytes, ...] = ( + _sse("message_start", {"type": "message_start", "message": {**OK_RESPONSE, "stop_reason": None}}), + _sse( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello"}}, + ), + _sse("message_stop", {"type": "message_stop"}), +) + + +def _split_frames_mid_data_line(frames: tuple[bytes, ...]) -> tuple[bytes, ...]: + """Split each frame's data line in half, modeling a transport chunk boundary.""" + return tuple(part for frame in frames for part in (frame[: len(frame) // 2], frame[len(frame) // 2 :])) + + +class _FrameStream(httpx.AsyncByteStream): + def __init__(self, frames: tuple[bytes, ...]) -> None: + self._frames = frames + + async def __aiter__(self) -> AsyncIterator[bytes]: + for frame in self._frames: + yield frame + + async def aclose(self) -> None: + return None + + +class FakeAnthropicUpstream: + """Intercepts the third-party transport (httpx.AsyncClient.send): refuses on fable + models, answers on others. The router deliberately does not forward caller-injected + clients, so the transport is the seam that exercises the real litellm pipeline.""" + + def __init__( + self, + refusal_body: dict[str, Any] = REFUSAL_RESPONSE, + refusal_frames: tuple[bytes, ...] = REFUSAL_STREAM_FRAMES, + ) -> None: + self.refusal_body = refusal_body + self.refusal_frames = refusal_frames + self.calls: list[str] = [] + self.bodies: list[dict[str, Any]] = [] + + async def send(self, request: httpx.Request, **kwargs: Any) -> httpx.Response: + body = json.loads(request.content or b"{}") + model = body.get("model", "") + self.calls.append(model) + self.bodies.append(body) + refuses = "fable" in model + if body.get("stream"): + frames = self.refusal_frames if refuses else OK_STREAM_FRAMES + return httpx.Response( + 200, + stream=_FrameStream(frames), + headers={"content-type": "text/event-stream"}, + request=request, + ) + return httpx.Response(200, json=self.refusal_body if refuses else OK_RESPONSE, request=request) + + def install(self): + async def _send(_client: httpx.AsyncClient, request: httpx.Request, **kwargs: Any) -> httpx.Response: + return await self.send(request, **kwargs) + + return patch("httpx.AsyncClient.send", new=_send) + + +FABLE_TIER = { + "model_name": "fable-tier", + "litellm_params": {"model": "anthropic/claude-fable-5", "api_key": "sk-test"}, +} +OPUS_TARGET = { + "model_name": "opus-target", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "sk-test"}, +} + + +def _router(content_policy_fallbacks: list | None) -> Router: + return Router(model_list=[FABLE_TIER, OPUS_TARGET], content_policy_fallbacks=content_policy_fallbacks) + + +async def _collect(stream: AsyncIterator[bytes]) -> bytes: + return b"".join([chunk async for chunk in stream]) + + +@pytest.mark.asyncio +async def test_non_streaming_refusal_with_fallback_row_returns_fallback_response(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "end_turn" + assert response["id"] == "msg_ok" + assert len(fake.calls) == 2 + assert "claude-opus-5" in fake.calls[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content_policy_fallbacks, upstream_body", + [ + (None, REFUSAL_RESPONSE), + ([{"unrelated-group": ["opus-target"]}], REFUSAL_RESPONSE), + ([{"fable-tier": ["opus-target"]}], PLAIN_REFUSAL_RESPONSE), + ], + ids=["nothing-configured", "row-for-other-group", "refusal-without-stop-details"], +) +async def test_non_streaming_refusal_passes_through_untouched(content_policy_fallbacks, upstream_body): + fake = FakeAnthropicUpstream(refusal_body=upstream_body) + router = _router(content_policy_fallbacks=content_policy_fallbacks) + + with fake.install(): + response = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, messages=[{"role": "user", "content": "hi"}] + ) + + assert response["stop_reason"] == "refusal" + assert response.get("stop_details") == upstream_body.get("stop_details") + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_with_fallback_row_streams_fallback_frames(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_split_across_chunks_still_falls_back(): + fake = FakeAnthropicUpstream(refusal_frames=_split_frames_mid_data_line(REFUSAL_STREAM_FRAMES)) + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_streaming_refusal_without_fallback_row_passes_frames_through(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=None) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert b"stop_details" in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata(): + """The pre-routing hook's tier stamp must reach the mid-stream fallback lookup even when the + request carries no metadata bucket at all (the snapshot is taken before the request runs).""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + stream = await router.aanthropic_messages( + model="smart-router", max_tokens=16, stream=True, messages=[{"role": "user", "content": "hi"}] + ) + body = await _collect(stream) + + assert b'"refusal"' not in body + assert b"text_delta" in body + assert len(fake.calls) == 2 + + +@pytest.mark.asyncio +async def test_caller_forged_tier_stamp_cannot_pick_the_streaming_fallback_chain(): + fake = FakeAnthropicUpstream() + router = _router(content_policy_fallbacks=[{"forged-tier": ["opus-target"]}]) + + with fake.install(): + stream = await router.aanthropic_messages( + model="fable-tier", + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "hi"}], + litellm_metadata={PRE_ROUTING_SELECTED_MODEL_KEY: "forged-tier"}, + ) + body = await _collect(stream) + + assert b'"stop_reason": "refusal"' in body + assert len(fake.calls) == 1 + + +@pytest.mark.asyncio +async def test_tier_stamp_never_reaches_provider_bound_metadata(): + """On /v1/messages the top-level metadata dict is Anthropic's own request field, so the + routed-tier stamp must never appear in any upstream body even when the client sends one.""" + fake = FakeAnthropicUpstream() + smart_router = { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "fable-tier", "MEDIUM": "fable-tier", "COMPLEX": "fable-tier"} + }, + "complexity_router_default_model": "fable-tier", + }, + "model_info": {"id": "router-1", "db_model": True}, + } + router = Router( + model_list=[FABLE_TIER, OPUS_TARGET, smart_router], + content_policy_fallbacks=[{"fable-tier": ["opus-target"]}], + ignore_invalid_deployments=True, + ) + + with fake.install(): + response = await router.aanthropic_messages( + model="smart-router", + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + metadata={"user_id": "u1"}, + ) + + assert response["stop_reason"] == "end_turn" + assert len(fake.bodies) == 2 + for body in fake.bodies: + assert body.get("metadata") == {"user_id": "u1"} + + +def test_record_pre_routing_selection_writes_only_the_internal_bucket(): + """The Anthropic request's own metadata field must never carry the tier stamp.""" + kwargs = {"metadata": {"user_id": "u1"}, "litellm_metadata": {}} + + record_pre_routing_selection(kwargs, "tier-x") + + assert kwargs["litellm_metadata"] == {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-x"} + assert kwargs["metadata"] == {"user_id": "u1"} + + +def test_refusal_gate_keys_on_pre_routing_tier_stamp(): + router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}]) + + def anthropic_messages(**kwargs: Any) -> None: + return None + + refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}} + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs=refusal_kwargs, + ) + is True + ) + assert ( + router._should_raise_anthropic_refusal_error( + model="router-group", + original_generic_function=anthropic_messages, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) + + +def test_has_content_policy_fallback_default_fallbacks_arm(): + router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}]) + + assert router._has_content_policy_fallback("any-group", {}) is True + assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False + + +def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested(): + router = _router(content_policy_fallbacks=None) + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}] + + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier1", "smart-router") + ) == ["backup-a"] + assert router._get_fallback_model_group_for_lookup_groups( + fallbacks=fallbacks, lookup_groups=("tier9", "smart-router") + ) == ["backup-b"] + assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None + + +def test_refusal_gate_ignores_other_generic_call_types(): + router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}]) + + def aresponses(**kwargs: Any) -> None: + return None + + assert ( + router._should_raise_anthropic_refusal_error( + model="fable-tier", + original_generic_function=aresponses, + response=dict(REFUSAL_RESPONSE), + kwargs={}, + ) + is False + ) diff --git a/tests/test_litellm/experimental_mcp_client/test_tools.py b/tests/test_litellm/experimental_mcp_client/test_tools.py index 89f67452f29..6645b06664d 100644 --- a/tests/test_litellm/experimental_mcp_client/test_tools.py +++ b/tests/test_litellm/experimental_mcp_client/test_tools.py @@ -8,11 +8,13 @@ from mcp.types import ( CallToolRequestParams, CallToolResult, ListToolsResult, + PaginatedRequestParams, TextContent, ) from mcp.types import Tool as MCPTool from litellm.experimental_mcp_client.tools import ( + list_tools_with_pagination, transform_mcp_tool_to_anthropic_tool, _get_function_arguments, _normalize_mcp_input_schema, @@ -106,6 +108,134 @@ async def test_load_mcp_tools_openai_format(mock_session, mock_list_tools_result mock_session.list_tools.assert_called_once() +@pytest.mark.asyncio() +async def test_load_mcp_tools_follows_pagination(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[ + MCPTool(name="tool_a", description="a", inputSchema={}), + MCPTool(name="tool_b", description="b", inputSchema={}), + ], + nextCursor="page-2", + ), + ListToolsResult(tools=[MCPTool(name="tool_c", description="c", inputSchema={})]), + ] + result = await load_mcp_tools(mock_session, format="mcp") + assert [tool.name for tool in result] == ["tool_a", "tool_b", "tool_c"] + assert mock_session.list_tools.call_count == 2 + second_call_params = mock_session.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_at_page_cap(mock_session, monkeypatch): + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_MAX_PAGES", 2) + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="page-2", + ), + ListToolsResult( + tools=[MCPTool(name="tool_1", description="1", inputSchema={})], + nextCursor="page-3", + ), + ListToolsResult(tools=[MCPTool(name="tool_2", description="2", inputSchema={})]), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + assert mock_session.list_tools.call_count == 2 + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_on_repeated_cursor(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="same-cursor", + ), + ListToolsResult( + tools=[MCPTool(name="tool_1", description="1", inputSchema={})], + nextCursor="same-cursor", + ), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + assert mock_session.list_tools.call_count == 2 + + +@pytest.mark.asyncio() +async def test_pagination_walk_treats_empty_cursor_as_terminal(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_0", description="0", inputSchema={})], + nextCursor="", + ), + ] + result = await list_tools_with_pagination(mock_session) + assert [tool.name for tool in result] == ["tool_0"] + mock_session.list_tools.assert_called_once() + + +@pytest.mark.asyncio() +async def test_pagination_walk_stops_at_whole_walk_deadline(mock_session, monkeypatch): + import anyio + + from litellm.experimental_mcp_client.tools import list_tools_with_pagination + + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.2) + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.2) + + async def slow_page(params=None): + await anyio.sleep(0.15) + idx = int(params.cursor) if params is not None else 0 + return ListToolsResult( + tools=[MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})], + nextCursor=str(idx + 1), + ) + + mock_session.list_tools = slow_page + result = await list_tools_with_pagination(mock_session) + + assert [tool.name for tool in result] == ["tool_0"] + + +@pytest.mark.asyncio() +async def test_pagination_walk_honors_explicit_deadline_over_globals(mock_session, monkeypatch): + import anyio + + from litellm.experimental_mcp_client.tools import list_tools_with_pagination + + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_CLIENT_TIMEOUT", 0.1) + monkeypatch.setattr("litellm.experimental_mcp_client.tools.MCP_TOOL_LISTING_TIMEOUT", 0.1) + + async def slow_page(params=None): + await anyio.sleep(0.15) + idx = int(params.cursor) if params is not None else 0 + tools = [MCPTool(name=f"tool_{idx}", description=str(idx), inputSchema={})] + if idx == 0: + return ListToolsResult(tools=tools, nextCursor="1") + return ListToolsResult(tools=tools) + + mock_session.list_tools = slow_page + result = await list_tools_with_pagination(mock_session, listing_deadline=2.0) + + assert [tool.name for tool in result] == ["tool_0", "tool_1"] + + +@pytest.mark.asyncio() +async def test_load_mcp_tools_openai_format_spans_pages(mock_session): + mock_session.list_tools.side_effect = [ + ListToolsResult( + tools=[MCPTool(name="tool_a", description="a", inputSchema={})], + nextCursor="page-2", + ), + ListToolsResult(tools=[MCPTool(name="tool_b", description="b", inputSchema={})]), + ] + result = await load_mcp_tools(mock_session, format="openai") + assert [t["function"]["name"] for t in result] == ["tool_a", "tool_b"] + + def test_get_function_arguments(): # Test with string arguments function = {"arguments": '{"test": "value"}'} diff --git a/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py b/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py new file mode 100644 index 00000000000..45e1acecec8 --- /dev/null +++ b/tests/test_litellm/integrations/SlackAlerting/test_user_spend_alerts.py @@ -0,0 +1,193 @@ +import datetime +from typing import Final +from unittest.mock import AsyncMock, patch + +import pytest +from pydantic import ValidationError + +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.SlackAlerting.user_spend_alerts import ( + UserSpendRow, + evaluate_user_spend, +) +from litellm.types.integrations.slack_alerting import ( + DEFAULT_ALERT_TYPES, + AlertType, + SlackAlertingArgs, +) + +TODAY: Final = datetime.date(2026, 8, 15) + + +def _row( + daily_spend: float = 0.0, + monthly_spend: float = 0.0, + baseline_spend: float = 0.0, +) -> UserSpendRow: + return UserSpendRow( + user_id="user-1", + daily_spend=daily_spend, + monthly_spend=monthly_spend, + baseline_spend=baseline_spend, + ) + + +def _evaluate(row: UserSpendRow, args: SlackAlertingArgs, thresholds: bool = True, anomalies: bool = True): + return evaluate_user_spend( + row=row, + args=args, + today=TODAY, + thresholds_enabled=thresholds, + anomalies_enabled=anomalies, + ) + + +def test_daily_threshold_crossed(): + args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0) + events: Final = _evaluate(_row(daily_spend=75.0, monthly_spend=75.0), args) + assert [e.kind for e in events] == ["daily_threshold"] + assert "`$75.00`" in events[0].message + assert "`$50.00`" in events[0].message + assert events[0].alert_type == AlertType.user_spend_thresholds + assert events[0].cache_key == "user_spend_alert_daily_user-1_2026-08-15" + + +def test_daily_threshold_not_crossed(): + args: Final = SlackAlertingArgs(daily_spend_per_user_threshold=50.0, spend_anomaly_min_spend=1000.0) + assert _evaluate(_row(daily_spend=49.99, monthly_spend=49.99), args) == () + + +def test_thresholds_unset_by_default(): + args: Final = SlackAlertingArgs(spend_anomaly_min_spend=1000.0) + assert _evaluate(_row(daily_spend=999.0, monthly_spend=999.0), args) == () + + +def test_monthly_threshold_crossed(): + args: Final = SlackAlertingArgs(monthly_spend_per_user_threshold=200.0, spend_anomaly_min_spend=1000.0) + events: Final = _evaluate(_row(daily_spend=5.0, monthly_spend=250.0), args) + assert [e.kind for e in events] == ["monthly_threshold"] + assert events[0].cache_key == "user_spend_alert_monthly_user-1_2026-08" + + +def test_thresholds_disabled_suppresses_threshold_events(): + args: Final = SlackAlertingArgs( + daily_spend_per_user_threshold=50.0, + monthly_spend_per_user_threshold=200.0, + spend_anomaly_min_spend=1000.0, + ) + assert _evaluate(_row(daily_spend=75.0, monthly_spend=250.0), args, thresholds=False) == () + + +def test_anomaly_detected_above_multiple_of_baseline(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + events: Final = _evaluate( + _row(daily_spend=70.0, monthly_spend=100.0, baseline_spend=70.0), args + ) + assert [e.kind for e in events] == ["anomaly"] + assert events[0].alert_type == AlertType.user_spend_anomalies + assert "`$10.00`" in events[0].message + assert events[0].cache_key == "user_spend_alert_anomaly_user-1_2026-08-15" + + +def test_no_anomaly_within_baseline_multiple(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert ( + _evaluate(_row(daily_spend=25.0, monthly_spend=100.0, baseline_spend=70.0), args) == () + ) + + +def test_no_anomaly_below_min_spend_floor(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert _evaluate(_row(daily_spend=9.0, monthly_spend=9.0, baseline_spend=0.1), args) == () + + +def test_anomaly_for_new_user_without_baseline(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + events: Final = _evaluate(_row(daily_spend=15.0, monthly_spend=15.0), args) + assert [e.kind for e in events] == ["anomaly"] + + +def test_sparse_baseline_averages_over_full_window(): + args: Final = SlackAlertingArgs( + spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0, spend_anomaly_baseline_days=7 + ) + events: Final = _evaluate(_row(daily_spend=13.0, monthly_spend=20.0, baseline_spend=7.0), args) + assert [e.kind for e in events] == ["anomaly"] + + +def test_anomalies_not_in_default_alert_types(): + assert AlertType.user_spend_anomalies not in DEFAULT_ALERT_TYPES + assert AlertType.user_spend_thresholds in DEFAULT_ALERT_TYPES + + +def test_invalid_config_rejected(): + with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"): + SlackAlertingArgs(daily_spend_per_user_threshold=0) + with pytest.raises(ValidationError, match="spend_anomaly_baseline_days"): + SlackAlertingArgs(spend_anomaly_baseline_days=0) + with pytest.raises(ValidationError, match="user_spend_check_interval"): + SlackAlertingArgs(user_spend_check_interval=10) + + +def test_non_finite_config_rejected(): + with pytest.raises(ValidationError, match="daily_spend_per_user_threshold"): + SlackAlertingArgs(daily_spend_per_user_threshold=float("inf")) + with pytest.raises(ValidationError, match="spend_anomaly_multiplier"): + SlackAlertingArgs(spend_anomaly_multiplier=float("nan")) + with pytest.raises(ValidationError, match="spend_anomaly_min_spend"): + SlackAlertingArgs(spend_anomaly_min_spend=float("inf")) + with pytest.raises(ValidationError, match="user_spend_check_interval"): + SlackAlertingArgs(user_spend_check_interval=float("inf")) + + +def test_anomalies_disabled_suppresses_anomaly_events(): + args: Final = SlackAlertingArgs(spend_anomaly_multiplier=3.0, spend_anomaly_min_spend=10.0) + assert _evaluate(_row(daily_spend=500.0, monthly_spend=500.0), args, anomalies=False) == () + + +@pytest.mark.asyncio +async def test_send_user_spend_alerts_sends_and_dedupes(): + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alerting_args={"daily_spend_per_user_threshold": 50.0, "spend_anomaly_min_spend": 1000.0}, + ) + mock_prisma: Final = AsyncMock() + mock_prisma.db.query_raw = AsyncMock( + return_value=[ + { + "user_id": "user-1", + "daily_spend": 75.0, + "monthly_spend": 75.0, + "baseline_spend": 0.0, + }, + { + "user_id": "user-2", + "daily_spend": 60.0, + "monthly_spend": 60.0, + "baseline_spend": 0.0, + }, + ] + ) + with patch.object(slack_alerting, "send_alert", new_callable=AsyncMock) as mock_send_alert: + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + assert mock_send_alert.call_count == 1 + sent_kwargs: Final = mock_send_alert.call_args.kwargs + assert sent_kwargs["alert_type"] == AlertType.user_spend_thresholds + assert "User Daily Spend Threshold Crossed" in sent_kwargs["message"] + assert "`user-1`" in sent_kwargs["message"] + assert "`user-2`" in sent_kwargs["message"] + + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + assert mock_send_alert.call_count == 1 + + +@pytest.mark.asyncio +async def test_send_user_spend_alerts_noop_when_alert_types_disabled(): + slack_alerting: Final = SlackAlerting( + alerting=["slack"], + alert_types=[AlertType.budget_alerts], + alerting_args={"daily_spend_per_user_threshold": 50.0}, + ) + mock_prisma: Final = AsyncMock() + await slack_alerting.send_user_spend_alerts(prisma_client=mock_prisma) + mock_prisma.db.query_raw.assert_not_called() diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index 9ab66d55f3f..c037f928593 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -1435,6 +1435,80 @@ class TestFlattenTopLevelSchemaCombinators: assert schema == snapshot +class TestToolWithFlattenedParameters: + def _anyof_tool(self): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def test_flattens_anyof_parameters_into_new_tool(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + tool = self._anyof_tool() + result = tool_with_flattened_parameters(tool) + + assert result is not tool + parameters = result["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert result["function"]["name"] == "automation_update" + assert tool == self._anyof_tool() + + def test_clean_parameters_return_the_same_tool_object(self): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + + assert tool_with_flattened_parameters(tool) is tool + + @pytest.mark.parametrize( + "tool", + [ + {"type": "function"}, + {"type": "function", "function": "not-a-dict"}, + {"type": "function", "function": {"name": "no_params"}}, + {"type": "function", "function": {"name": "bad_params", "parameters": "not-a-dict"}}, + ], + ) + def test_non_dict_function_or_parameters_return_the_same_tool_object(self, tool): + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + tool_with_flattened_parameters, + ) + + assert tool_with_flattened_parameters(tool) is tool + + class TestRequestContainsImageContent: """One detector for every dialect that reaches pre-routing hooks untranslated.""" diff --git a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py index cb4e72ab3ad..722818598af 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_supported_openai_params.py @@ -188,6 +188,8 @@ class TestDeclaredAuthenticatingProvider: ("gpt-4o", "github_copilot", "github_copilot"), ("openai/gpt-4o", None, None), ("gpt-4o", "openai", None), + ("github_copilot", None, None), + ("chatgpt", None, None), ], ) def test_names_only_the_providers_whose_resolution_authenticates(self, model, provider, expected): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index c7328adb0b3..366f61ded49 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5479,6 +5479,97 @@ def test_pre_call_redacts_and_masks_raw_request(logging_obj): assert "key=*****" in raw_api_base +def _streaming_logging_obj_with_callbacks(callbacks: list[CustomLogger]): + import datetime + + obj = LitellmLogging( + model="anthropic/claude-opus-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="slot-leak-test", + function_id="slot-leak-test", + ) + obj.model_call_details["litellm_params"] = {"metadata": {}} + return patch.object(obj, "get_combined_callback_list", return_value=callbacks), obj + + +def _assembled_stream_result(): + response = ModelResponse() + response.choices[0].message.content = "hello" + return response + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_logging_hook_failure(): + """Regression for leaked max_parallel_requests slots: a raising + async_logging_hook must not abort the success-callback loop that + releases the rate-limiter slot.""" + broken = CustomLogger() + broken.async_logging_hook = AsyncMock(side_effect=RuntimeError("broken stream payload")) + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([broken, releasing]) + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_cost_calculation_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_response_cost_calculator", side_effect=ValueError("bad usage block") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + assert logging_obj.model_call_details["response_cost"] is None + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_standard_logging_payload_failure(): + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([releasing]) + with patcher, patch.object( + logging_obj, "_build_standard_logging_payload", side_effect=ValueError("incomplete stream") + ): + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + assert logging_obj.model_call_details.get("standard_logging_object") is None + releasing.async_log_success_event.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failure(): + from litellm.integrations.custom_guardrail import CustomGuardrail + + skipping = CustomGuardrail(guardrail_name="skipping-guardrail") + skipping.should_run_guardrail = MagicMock(return_value=False) + skipping.async_logging_hook = AsyncMock() + raising = CustomGuardrail(guardrail_name="raising-guardrail") + raising.should_run_guardrail = MagicMock(return_value=True) + raising.async_logging_hook = AsyncMock(side_effect=RuntimeError("guardrail hook failed")) + releasing = CustomLogger() + releasing.async_log_success_event = AsyncMock() + + patcher, logging_obj = _streaming_logging_obj_with_callbacks([skipping, raising, releasing]) + with patcher: + await logging_obj.async_success_handler(result=_assembled_stream_result()) + + skipping.async_logging_hook.assert_not_awaited() + raising.async_logging_hook.assert_awaited_once() + releasing.async_log_success_event.assert_awaited_once() + + def _resolve(custom_llm_provider, litellm_params, optional_params, model): from litellm.litellm_core_utils.litellm_logging import ( _resolve_vertex_location_for_cost, diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py index 2cf7cd142d6..4e6b9ed0188 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_gpt_transformation.py @@ -11,6 +11,7 @@ sys.path.insert( import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import TOOL_RESULT_IMAGE_BOUNDARY +from litellm.llms.azure.chat.gpt_5_transformation import AzureOpenAIGPT5Config from litellm.llms.azure.chat.gpt_transformation import AzureOpenAIConfig from litellm.utils import get_optional_params @@ -195,3 +196,91 @@ def test_azure_gpt_5_takes_the_reasoning_path() -> None: assert "presence_penalty" not in mapped assert "logit_bias" not in mapped assert "reasoning_effort" in supported + + +class TestAzureToolSchemaCombinatorFlattening: + """ + Regression tests for LIT-6510: Azure's chat completions validator rejects + tool parameters carrying a top-level anyOf/oneOf/allOf for every model + family, so AzureOpenAIConfig.transform_request must flatten them. + """ + + @staticmethod + def _anyof_tool(): + return { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + + def _transform(self, config, model, tools): + return config.transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={"tools": tools}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + def test_transform_request_flattens_top_level_anyof(self): + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [self._anyof_tool()]) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert request["tools"][0]["function"]["name"] == "automation_update" + + def test_gpt5_config_flattens_via_shared_transform(self): + request = self._transform(AzureOpenAIGPT5Config(), "gpt-5.4-mini", [self._anyof_tool()]) + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + + def test_caller_tool_dict_is_not_mutated(self): + tool = self._anyof_tool() + self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + assert tool == self._anyof_tool() + + def test_clean_object_schema_passes_through_as_same_object(self): + tool = { + "type": "function", + "function": { + "name": "lookup", + "parameters": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + } + request = self._transform(AzureOpenAIConfig(), "gpt-4o", [tool]) + assert request["tools"][0] is tool + + def test_non_dict_tool_entries_pass_through_unchanged(self): + request = self._transform(AzureOpenAIConfig(), "gpt-4o", ["not-a-tool"]) + assert request["tools"] == ["not-a-tool"] + + def test_request_without_tools_is_unchanged(self): + request = AzureOpenAIConfig().transform_request( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + optional_params={"temperature": 0.2}, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + assert "tools" not in request + assert request["temperature"] == 0.2 diff --git a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py index fc7e94a77ba..202f81f1252 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_chat_o_series_transformation.py @@ -23,3 +23,48 @@ async def test_azure_chat_o_series_transformation(): ) print(response) assert response["model"] == "web-interface-o1-mini" + + +def test_azure_o_series_transform_request_flattens_top_level_anyof(): + """Regression test for LIT-6510: the o-series super() chain ends in + OpenAIGPTConfig, whose flatten gate skips provider 'azure', so + AzureOpenAIO1Config must flatten tool schema combinators itself.""" + tool = { + "type": "function", + "function": { + "name": "automation_update", + "description": "Update an automation", + "parameters": { + "type": "object", + "anyOf": [ + { + "properties": {"id": {"type": "string"}, "enabled": {"type": "boolean"}}, + "required": ["id", "enabled"], + }, + { + "properties": {"id": {"type": "string"}, "schedule": {"type": "string"}}, + "required": ["id", "schedule"], + }, + ], + "properties": {"id": {"type": "string"}}, + "required": ["id"], + }, + }, + } + optional_params = {"tools": [tool]} + + request = AzureOpenAIO1Config().transform_request( + model="o3-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={"custom_llm_provider": "azure"}, + headers={}, + ) + + parameters = request["tools"][0]["function"]["parameters"] + assert "anyOf" not in parameters + assert parameters["type"] == "object" + assert set(parameters["properties"]) == {"id", "enabled", "schedule"} + assert parameters["required"] == ["id"] + assert "anyOf" in tool["function"]["parameters"] + assert optional_params["tools"][0] is tool diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..ad392d4962d 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -828,6 +828,24 @@ class MockPassThroughGuardrail(CustomGuardrail): return inputs +class MockRecordingGuardrail(MockPassThroughGuardrail): + """Pass-through guardrail that records every apply_guardrail inputs payload""" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.seen_inputs: List[GenericGuardrailAPIInputs] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + self.seen_inputs.append(inputs) + return inputs + + class TestOpenAIResponsesHandlerStreamingOutputProcessing: """Test streaming output processing functionality""" @@ -1104,6 +1122,80 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: output_text = result[-1]["response"]["output"][0]["content"][0]["text"] assert output_text == original_text + @pytest.mark.asyncio + async def test_failed_stream_scans_delta_text(self): + """A stream ending in response.failed has text only in delta events; the + fallback scan must assemble and scan it instead of skipping on an empty string.""" + handler = OpenAIResponsesHandler() + guardrail = MockRecordingGuardrail(guardrail_name="test") + + responses_so_far = [ + {"type": "response.created", "response": {"id": "resp_123"}}, + {"type": "response.output_item.added", "item": {"type": "message", "id": "msg_123"}}, + { + "type": "response.output_text.delta", + "item_id": "msg_123", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_123", + "output_index": 0, + "content_index": 0, + "delta": " world", + }, + {"type": "response.failed", "response": {"id": "resp_123", "status": "failed"}}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result == responses_so_far + assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["Hello world"]] + + def test_get_streaming_string_so_far_prefers_done_text_over_deltas(self): + """The done event repeats the whole part, so deltas must not be double counted; + a part with no done event yet still contributes its joined deltas.""" + handler = OpenAIResponsesHandler() + + events = [ + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": "Hello", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "delta": " world", + }, + { + "type": "response.output_text.done", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "text": "Hello world", + }, + { + "type": "response.output_text.delta", + "item_id": "msg_2", + "output_index": 1, + "content_index": 0, + "delta": "; unfinished", + }, + ] + + assert handler.get_streaming_string_so_far(events) == "Hello world; unfinished" + class TestGetStructuredMessages: """Test the get_structured_messages method for Responses API handler.""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ef5631218f3..e36bef229f6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -214,6 +214,46 @@ class TestExecuteWithMcpClient: assert server.scopes == ["read", "write"] assert server.has_client_credentials is True + async def test_preview_forwards_per_server_timeout_to_client_factory(self, monkeypatch): + """The request's per-server timeout must reach the temporary MCPServer model: + the client factory reads ``server.timeout`` for both the per-request timeout + and the preview's whole-walk listing deadline.""" + captured: dict = {} + + def fake_build_stdio_env(server, raw_headers): + return None + + async def fake_create_client(*args, **kwargs): + captured["server"] = kwargs.get("server") + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_build_stdio_env", + fake_build_stdio_env, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + raising=False, + ) + + async def ok_operation(client): + return {"status": "ok"} + + payload = NewMCPServerRequest( + server_name="slow-catalog-server", + url="https://example.com", + timeout=120.5, + ) + + result = await rest_endpoints._execute_with_mcp_client(payload, ok_operation) + + assert result["status"] == "ok" + assert captured["server"].timeout == 120.5 + @pytest.mark.asyncio async def test_m2m_drops_incoming_oauth2_headers(self, monkeypatch): """For M2M OAuth servers the incoming Authorization header (which carries @@ -524,6 +564,131 @@ class TestTestToolsList: assert captured["oauth2_headers"] is None assert oauth_call_counter["count"] == 0 + async def test_preview_tools_list_times_out_on_slow_pagination(self, monkeypatch): + """A preview whose upstream paginates past the listing deadline returns a + timeout error instead of holding the request open.""" + monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False) + monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False) + + class SlowClient: + async def list_tools(self, raise_on_error=False): + await asyncio.sleep(1) + return [] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(SlowClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["status"] == "error" + assert result["error"] is True + assert "Timed out listing tools" in result["message"] + + async def test_preview_tools_list_succeeds_within_deadline(self, monkeypatch): + """The preview timeout scope passes a fast listing through untouched.""" + from mcp.types import Tool as MCPTool + + class QuickClient: + async def list_tools(self, raise_on_error=False): + return [MCPTool(name="quick_tool", description="q", inputSchema={})] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(QuickClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["error"] is None + assert result["message"] == "Successfully retrieved tools" + assert [tool["name"] for tool in result["tools"]] == ["quick_tool"] + + async def test_preview_tools_list_honors_per_server_timeout(self, monkeypatch): + """A per-server timeout above the global default extends the preview deadline.""" + monkeypatch.setattr(rest_endpoints, "MCP_CLIENT_TIMEOUT", 0.05, raising=False) + monkeypatch.setattr(rest_endpoints, "MCP_TOOL_LISTING_TIMEOUT", 0.05, raising=False) + + from mcp.types import Tool as MCPTool + + class SlowConfiguredClient: + timeout = 1.0 + + async def list_tools(self, raise_on_error=False): + await asyncio.sleep(0.2) + return [MCPTool(name="slow_tool", description="s", inputSchema={})] + + async def fake_execute( + request, + operation, + mcp_auth_header=None, + oauth2_headers=None, + raw_headers=None, + ): + return await operation(SlowConfiguredClient()) + + monkeypatch.setattr(rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False) + + from litellm.proxy._types import LitellmUserRoles + + request = _build_request() + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.api_key, + credentials={"auth_value": "secret-key"}, + ) + + result = await rest_endpoints.test_tools_list( + request, + payload, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert result["error"] is None + assert [tool["name"] for tool in result["tools"]] == ["slow_tool"] + async def test_extracts_oauth2_headers(self, monkeypatch): """Ensure oauth2 auth type pulls oauth headers and omits MCP auth header.""" @@ -786,9 +951,7 @@ class TestListToolsRestAPI: they do for a gateway session, never to the bare session key.""" from litellm.constants import UI_SESSION_TOKEN_TEAM_ID - session_auth = UserAPIKeyAuth( - team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" - ) + session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user") admitted_auth = UserAPIKeyAuth(user_id="grant-user", org_id="admitted-org") async def fake_reload(user_id): @@ -868,9 +1031,7 @@ class TestListToolsRestAPI: from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy._types import LiteLLM_ObjectPermissionTable - session_auth = UserAPIKeyAuth( - team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user" - ) + session_auth = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="grant-user", user_role="internal_user") scoped_auth = UserAPIKeyAuth( object_permission=LiteLLM_ObjectPermissionTable( object_permission_id="toolset-scope", @@ -952,6 +1113,123 @@ class TestListToolsRestAPI: assert scope_inputs == [session_auth] assert reload_calls == [] + async def test_single_server_response_includes_paginated_upstream_tools( + self, + monkeypatch, + ): + """The REST tools/list path should include tools beyond the upstream first page.""" + import litellm.experimental_mcp_client.client as mcp_client_module + from mcp.types import ListToolsResult, PaginatedRequestParams + from mcp.types import Tool as MCPTool + + from litellm.proxy._experimental.mcp_server.server import MCPServer + from litellm.types.mcp import MCPTransport + + async def fake_contexts(user_api_key_auth): + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return ["server-1"] + + stub_server = MCPServer( + server_id="server-1", + name="stub", + server_name="stub", + alias="stub", + url="https://example.com/mcp", + transport=MCPTransport.http, + mcp_info={"server_name": "stub"}, + ) + stub_server.available_on_public_internet = True + + mock_transport_ctx = AsyncMock() + mock_transport_ctx.__aenter__ = AsyncMock(return_value=(MagicMock(), MagicMock())) + mock_transport_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "streamable_http_client", + MagicMock(return_value=mock_transport_ctx), + raising=False, + ) + + mock_session_ctx = AsyncMock() + mock_session_instance = AsyncMock() + mock_session_instance.initialize = AsyncMock(return_value=None) + mock_session_instance.list_tools.side_effect = [ + ListToolsResult( + tools=[ + MCPTool( + name="first_page_tool", + description="First page tool", + inputSchema={}, + ) + ], + nextCursor="page-2", + ), + ListToolsResult( + tools=[ + MCPTool( + name="second_page_tool", + description="Second page tool", + inputSchema={}, + ) + ] + ), + ] + mock_session_ctx.__aenter__ = AsyncMock(return_value=mock_session_instance) + mock_session_ctx.__aexit__ = AsyncMock(return_value=None) + monkeypatch.setattr( + mcp_client_module, + "ClientSession", + MagicMock(return_value=mock_session_ctx), + raising=False, + ) + + monkeypatch.setattr( + rest_endpoints, + "build_effective_auth_contexts", + fake_contexts, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "filter_server_ids_by_ip_with_info", + lambda server_ids, client_ip: (server_ids, 0), + raising=False, + ) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_mcp_server_by_id", + lambda server_id: stub_server if server_id == "server-1" else None, + raising=False, + ) + + request = _build_request(path="/mcp-rest/tools/list", method="GET") + result = await rest_endpoints.list_tool_rest_api( + request, + server_id="server-1", + user_api_key_dict=UserAPIKeyAuth(), + ) + + assert set(result.keys()) == {"tools", "error", "message"} + assert [tool.name for tool in result["tools"]] == [ + "first_page_tool", + "second_page_tool", + ] + assert result["error"] is None + assert result["message"] == "Successfully retrieved tools" + + assert mock_session_instance.list_tools.call_count == 2 + second_call_params = mock_session_instance.list_tools.call_args_list[1].kwargs["params"] + assert isinstance(second_call_params, PaginatedRequestParams) + assert second_call_params.cursor == "page-2" + async def test_include_disabled_tools_is_admin_only(self, monkeypatch): """include_disabled_tools skips the allowlist filter only for PROXY_ADMIN; a non-admin passing it stays filtered so the REST endpoint can't be used @@ -3021,9 +3299,7 @@ class TestRestListToolsetFiltering: mock_manager = MagicMock() mock_manager.expand_tool_permissions = MagicMock(side_effect=lambda perms: perms or {}) - mock_manager.resolve_toolset_tool_permissions = AsyncMock( - return_value={"server-a": ["lookup_status"]} - ) + mock_manager.resolve_toolset_tool_permissions = AsyncMock(return_value={"server-a": ["lookup_status"]}) monkeypatch.setattr( rest_endpoints.global_mcp_server_manager, diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py new file mode 100644 index 00000000000..b7bc670c7f8 --- /dev/null +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_streaming_model_restamp.py @@ -0,0 +1,288 @@ +""" +Tests for restamping the public model on Anthropic Messages streaming chunks. +""" + +import json +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.proxy.anthropic_endpoints.streaming_model_restamp import ( + AnthropicStreamModelRestamper, + restamp_anthropic_stream_chunk_model, +) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + +def _message_start_frame(model: str, line_end: str = "\n") -> bytes: + payload = { + "type": "message_start", + "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": model, "content": []}, + } + return f"event: message_start{line_end}data: {json.dumps(payload)}{line_end}{line_end}".encode() + + +def _proxy_logging_obj_streaming(frames: list[bytes]) -> MagicMock: + async def _iterator_hook(**_kwargs): + for frame in frames: + yield frame + + proxy_logging_obj = MagicMock() + proxy_logging_obj.async_post_call_streaming_iterator_hook = _iterator_hook + proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["response"]) + return proxy_logging_obj + + +def _model_from_frame(frame: bytes | str) -> str: + text = frame.decode("utf-8") if isinstance(frame, bytes) else frame + data_line = next(line for line in text.split("\n") if line.startswith("data:")) + return json.loads(data_line[len("data:") :])["message"]["model"] + + +def test_restamps_sse_bytes_frame(): + restamped = restamp_anthropic_stream_chunk_model( + _message_start_frame("claude-haiku-4-5-20251001"), "claude-auto-1" + ) + + assert isinstance(restamped, bytes) + assert _model_from_frame(restamped) == "claude-auto-1" + assert b"event: message_start" in restamped + + +def test_restamps_event_dict(): + chunk = {"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}} + + restamped = restamp_anthropic_stream_chunk_model(chunk, "claude-auto-2") + + assert restamped == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-2"}} + assert chunk["message"]["model"] == "claude-sonnet-4-6" + + +@pytest.mark.parametrize( + "chunk", + [ + b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n', + {"type": "content_block_delta", "delta": {"text": "hi"}}, + {"type": "message_start", "message": "not-a-dict"}, + b"event: message_start\ndata: not-json\n\n", + b"data: [DONE]\n\n", + ], +) +def test_leaves_chunks_without_a_model_untouched(chunk): + assert restamp_anthropic_stream_chunk_model(chunk, "claude-auto-1") == chunk + + +@pytest.mark.asyncio +async def test_sse_generator_publishes_requested_model_on_message_start(): + """The message_start event reports the requested model, not the provider's.""" + delta_frame = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001"), delta_frame]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta_frame + + +@pytest.mark.asyncio +async def test_sse_generator_keeps_provider_model_when_restamping_is_off(): + proxy_logging_obj = _proxy_logging_obj_streaming([_message_start_frame("claude-haiku-4-5-20251001")]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-haiku-4-5-20251001" + + +def test_restamps_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_emits_coalesced_frames_with_only_message_start_rewritten(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + combined = _message_start_frame("claude-haiku-4-5-20251001") + delta + + emitted = restamper_output = AnthropicStreamModelRestamper("claude-auto-1").process(combined) + + assert isinstance(restamper_output, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(delta) + + +def test_ping_frames_keep_the_restamper_armed(): + ping = b'event: ping\ndata: {"type": "ping"}\n\n' + frame = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(ping) == ping + reassembled = restamper.process(frame[:10]) + reassembled += restamper.process(frame[10:]) + + assert _model_from_frame(reassembled) == "claude-auto-1" + + +def test_first_non_ping_event_disarms_the_restamper(): + delta = b'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\n\n' + late_message_start = _message_start_frame("claude-haiku-4-5-20251001") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(delta) == delta + assert restamper.process(late_message_start) == late_message_start + + +def test_oversized_unterminated_chunk_flushes_unmodified(): + blob = b"data: " + b"x" * 70000 + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(blob) == blob + frame = _message_start_frame("claude-haiku-4-5-20251001") + assert restamper.process(frame) == frame + + +def test_dict_message_start_disarms_after_restamp(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + first = restamper.process({"type": "message_start", "message": {"id": "msg_1", "model": "claude-sonnet-4-6"}}) + second = {"type": "message_start", "message": {"id": "msg_2", "model": "claude-sonnet-4-6"}} + + assert first == {"type": "message_start", "message": {"id": "msg_1", "model": "claude-auto-1"}} + assert restamper.process(second) == second + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_message_start_split_across_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001") + proxy_logging_obj = _proxy_logging_obj_streaming([frame[:30], frame[30:]]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert _model_from_frame(joined) == "claude-auto-1" + + +def test_restamps_crlf_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + assert emitted.endswith(b"\r\n\r\n") + assert restamper.process(delta) == delta + + +def test_restamps_cr_terminated_message_start_frame(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + emitted = restamper.process(frame) + + assert isinstance(emitted, bytes) + assert b'"model":"claude-auto-1"' in emitted + assert emitted.endswith(b"\r\r") + + +def test_restamps_crlf_message_start_split_across_transport_chunks(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + held = restamper.process(frame[:25]) + emitted = restamper.process(frame[25:]) + + assert held == b"" + assert isinstance(emitted, bytes) + assert _model_from_frame(emitted) == "claude-auto-1" + + +def test_flush_returns_restamped_held_tail(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + restamper = AnthropicStreamModelRestamper("claude-auto-1") + + assert restamper.process(unterminated) == b"" + flushed = restamper.flush() + + assert b'"model":"claude-auto-1"' in flushed + assert restamper.flush() == b"" + + +def test_flush_disarms_the_restamper(): + restamper = AnthropicStreamModelRestamper("claude-auto-1") + frame = _message_start_frame("claude-haiku-4-5-20251001") + + assert restamper.flush() == b"" + assert restamper.process(frame) == frame + + +@pytest.mark.asyncio +async def test_sse_generator_flushes_held_tail_at_end_of_stream(): + unterminated = _message_start_frame("claude-haiku-4-5-20251001")[:-2] + proxy_logging_obj = _proxy_logging_obj_streaming([unterminated]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + joined = b"".join(chunk if isinstance(chunk, bytes) else chunk.encode("utf-8") for chunk in chunks) + assert b'"model":"claude-auto-1"' in joined + + +@pytest.mark.asyncio +async def test_sse_generator_restamps_crlf_stream(): + frame = _message_start_frame("claude-haiku-4-5-20251001", line_end="\r\n") + delta = b'event: content_block_delta\r\ndata: {"type":"content_block_delta","delta":{"text":"hi"}}\r\n\r\n' + proxy_logging_obj = _proxy_logging_obj_streaming([frame, delta]) + + chunks = [ + chunk + async for chunk in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=MagicMock(), + user_api_key_dict=MagicMock(), + request_data={"model": "claude-auto-1"}, + proxy_logging_obj=proxy_logging_obj, + restamp_model="claude-auto-1", + ) + ] + + assert _model_from_frame(chunks[0]) == "claude-auto-1" + assert chunks[1] == delta diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index 898f9ab1ed7..e5f2a9d95bd 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -26,6 +26,64 @@ def _owners(*backup_paths): CLAUDE_SETTINGS_MODULE = "litellm.proxy.client.cli.commands.claude_settings" AUTH_MODULE = "litellm.proxy.client.cli.commands.auth" +WINDOWS_LITE_EXE = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" + +CMD_METACHARACTERS = frozenset("&|<>^()") +CMD_PERCENT_GUARD = "%%cd:~,%" + + +def _through_cmd_exe(command): + """The line cmd.exe hands to CreateProcess after reading the apiKeyHelper. + + A `"` toggles cmd's quote state and the metacharacters only act outside it. cmd expands + `%VAR%` even inside quotes, so every `%` has to arrive as the `%%cd:~,%` guard: the first + `%` has no variable name and stays literal, and `%cd:~,%` is a zero length substring of `cd`. + """ + assert not any(CMD_METACHARACTERS & set(run) for run in command.split('"')[::2]), command + assert command.count("%") == 3 * command.count(CMD_PERCENT_GUARD), command + return command.replace(CMD_PERCENT_GUARD, "%") + + +def _through_c_runtime(command_line): + """argv as the Microsoft C runtime builds it for the `lite` executable. + + Outside quotes whitespace ends an argument. A `"` toggles quoting, and inside quotes `""` + is a literal quote. Backslashes are literal unless they run up to a `"`, where each pair + is one backslash and an odd one left over makes the quote literal. + """ + argv = [] + current = None + quoted = False + i = 0 + while i < len(command_line): + ch = command_line[i] + if ch in " \t" and not quoted: + if current is not None: + argv.append(current) + current = None + i += 1 + continue + if current is None: + current = "" + if ch == "\\": + run = len(command_line[i:]) - len(command_line[i:].lstrip("\\")) + before_quote = command_line[i + run : i + run + 1] == '"' + current += "\\" * (run // 2 if before_quote else run) + if before_quote and run % 2: + current += '"' + i += 1 + i += run + elif ch == '"': + if quoted and command_line[i + 1 : i + 2] == '"': + current += '"' + i += 1 + else: + quoted = not quoted + i += 1 + else: + current += ch + i += 1 + return argv if current is None else [*argv, current] @pytest.fixture @@ -199,6 +257,36 @@ class TestApiKeyHelperIsActuallyInvocable: assert "Not authenticated for this server" in result.output + def _windows_argv(self, lite_exe, base_url): + with patch(f"{CLAUDE_SETTINGS_MODULE}.shutil.which", return_value=lite_exe): + helper = resolve_api_key_helper(base_url, platform="win32") + return _through_c_runtime(_through_cmd_exe(helper)) + + @pytest.mark.parametrize( + ("lite_exe", "base_url"), + [ + (WINDOWS_LITE_EXE, "http://localhost:4000"), + ("C:\\Program Files\\LiteLLM\\lite.EXE", "https://gateway.example.com/?a=1&b=2"), + ("C:\\Users\\u\\Scripts\\lite.EXE", "https://gateway.example.com/team%20a/%7Eproxy"), + ('C:\\odd "dir"\\lite.EXE', "http://localhost:4000/x\\"), + ], + ) + def test_the_windows_command_survives_cmd_exe_and_the_c_runtime(self, lite_exe, base_url): + assert self._windows_argv(lite_exe, base_url) == [lite_exe, "--base-url", base_url, "auth", "print-token"] + + def test_the_windows_command_carries_the_base_url_through_cmd_quoting(self): + stale = CliTokenRecord( + base_url="http://other-proxy.example.com", + key="sk-stale", + timestamp=time.time(), + ) + argv = self._windows_argv(WINDOWS_LITE_EXE, "http://localhost:4000") + with patch(f"{AUTH_MODULE}.load_cli_token", return_value=stale): + result = CliRunner().invoke(cli, argv[1:]) + + assert argv[0] == WINDOWS_LITE_EXE + assert "Not authenticated for this server" in result.output + class TestConflictingOwnersOfTheSettingsFile: """Both `lite up` and `lite autoroute up` restore a backup when they stop. diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 053c90c36b4..c78bdfa75b1 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -225,6 +225,32 @@ class TestResolveApiKeyHelper: with pytest.raises(ClaudeSettingsError, match="Could not find `lite`"): resolve_api_key_helper("http://localhost:4000") + def test_windows_quotes_for_cmd_exe_instead_of_posix_sh(self, monkeypatch): + """cmd.exe takes a single quote literally, so a POSIX-quoted backslashed path is unrunnable.""" + lite_exe = "C:\\Users\\u\\AppData\\Local\\Programs\\Python\\Python313\\Scripts\\lite.EXE" + monkeypatch.setattr(shutil, "which", lambda name: lite_exe) + + helper = resolve_api_key_helper("https://gateway.example.com", platform="win32") + + assert helper == f'"{lite_exe}" "--base-url" "https://gateway.example.com" "auth" "print-token"' + + def test_windows_keeps_a_spaced_path_and_a_metacharacter_url_as_single_tokens(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "C:\\Program Files\\LiteLLM\\lite.EXE") + + helper = resolve_api_key_helper("https://gateway.example.com/?a=1&b=2", platform="win32") + + assert helper == ( + '"C:\\Program Files\\LiteLLM\\lite.EXE" "--base-url" "https://gateway.example.com/?a=1&b=2" ' + '"auth" "print-token"' + ) + + def test_non_windows_platforms_keep_posix_quoting(self, monkeypatch): + monkeypatch.setattr(shutil, "which", lambda name: "/usr/local/bin/lite") + + helper = resolve_api_key_helper("http://example.com/path; rm -rf /", platform="darwin") + + assert helper == "/usr/local/bin/lite --base-url 'http://example.com/path; rm -rf /' auth print-token" + def _make_ctx(base_url): return click.Context(click.Command("test"), obj={"base_url": base_url}) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index bcda1b8b61d..8bc367b6b59 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5592,6 +5592,156 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" +def _responses_stream_events() -> list: + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseCompletedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + deltas = [ + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_lit6457", + output_index=0, + content_index=0, + delta=part, + ) + for part in ("Hello", " world") + ] + completed = ResponseCompletedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + response=ResponsesAPIResponse( + id="resp_lit6457", + created_at=1234567890, + model="gpt-4o", + object="response", + status="completed", + output=[ + { + "type": "message", + "id": "msg_lit6457", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Hello world"}], + } + ], + ), + ) + return [*deltas, completed] + + +@pytest.mark.asyncio +async def test_responses_api_stream_scans_output_and_replays_buffered_events(): + """Streamed /v1/responses events must be scanned via the unified translation + layer, not fed to stream_chunk_builder (which raises APIError on them).""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-responses-stream", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + stream_events = _responses_stream_events() + order = [] + yielded = [] + + async def record_scan(*args, **kwargs): + order.append("scan") + return {"action": "NONE", "assessments": [], "outputs": []} + + async def mock_stream(): + for event in stream_events: + yield event + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=mock_stream(), + request_data={"model": "gpt-4o", "input": "hi"}, + ): + order.append("chunk") + yielded.append(chunk) + + assert order == ["scan", "chunk", "chunk", "chunk"] + assert len(yielded) == len(stream_events) + assert all(emitted is original for emitted, original in zip(yielded, stream_events)) + + +def _responses_failed_stream_events() -> list: + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + ResponseFailedEvent, + ResponsesAPIResponse, + ResponsesAPIStreamEvents, + ) + + deltas = [ + OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_lit6457_failed", + output_index=0, + content_index=0, + delta=part, + ) + for part in ("Hello", " world") + ] + failed = ResponseFailedEvent( + type=ResponsesAPIStreamEvents.RESPONSE_FAILED, + response=ResponsesAPIResponse( + id="resp_lit6457_failed", + created_at=1234567890, + model="gpt-4o", + object="response", + status="failed", + output=[], + ), + ) + return [*deltas, failed] + + +@pytest.mark.asyncio +async def test_responses_api_failed_stream_scans_delta_text_before_replay(): + """A responses stream that dies mid-generation carries its text only in delta + events; the end-of-stream scan must still see that text instead of skipping + on an empty assembled string and replaying the buffer unmoderated.""" + guardrail = BedrockGuardrail( + guardrail_name="bedrock-responses-failed-stream", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + stream_events = _responses_failed_stream_events() + order = [] + scan_payloads = [] + yielded = [] + + async def record_scan(*args, **kwargs): + order.append("scan") + scan_payloads.append(str(args) + str(kwargs)) + return {"action": "NONE", "assessments": [], "outputs": []} + + async def mock_stream(): + for event in stream_events: + yield event + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/responses"), + response=mock_stream(), + request_data={"model": "gpt-4o", "input": "hi"}, + ): + order.append("chunk") + yielded.append(chunk) + + assert order == ["scan", "chunk", "chunk", "chunk"] + assert "Hello world" in scan_payloads[0] + assert len(yielded) == len(stream_events) + assert all(emitted is original for emitted, original in zip(yielded, stream_events)) + + @pytest.mark.asyncio async def test_apply_guardrail_debug_log_masks_signed_request_headers(): import logging diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index a81c6b4c656..7e2e680743f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -17489,3 +17489,49 @@ async def test_check_project_key_limits_still_rejects_real_model_outside_project assert exc_info.value.status_code == 400 assert "Model 'gpt-5.4-mini' not in project's allowed models" in exc_info.value.detail["error"] + + +def test_generate_key_request_blank_team_id_is_personal(): + """The UI Team-field clear submits team_id=""; it must count as no team (LIT-3925).""" + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _is_team_key, + ) + + cleared = GenerateKeyRequest(team_id="") + assert cleared.team_id is None + assert _is_team_key(data=cleared) is False + assert RegenerateKeyRequest(team_id="").team_id is None + assert GenerateKeyRequest(team_id="team-1").team_id == "team-1" + + +def test_key_generation_check_blank_team_id_uses_personal_permissions(monkeypatch): + """key_generation_check with team_id="" must take the personal-key path instead + of failing the team lookup with "Unable to find team object" (LIT-3925).""" + from litellm.proxy._types import KeyManagementRoutes + from litellm.proxy.management_endpoints.key_management_endpoints import ( + key_generation_check, + ) + + monkeypatch.setattr( + litellm, + "key_generation_settings", + { + "team_key_generation": {"allowed_team_member_roles": ["admin"]}, + "personal_key_generation": {"allowed_user_roles": ["proxy_admin", "internal_user"]}, + }, + ) + + assert ( + key_generation_check( + team_table=None, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + data=GenerateKeyRequest(key_alias="personal", team_id=""), + route=KeyManagementRoutes.KEY_GENERATE, + ) + is True + ) diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py index 35b5c72f92e..1e1436fcef8 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_utils.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_utils.py @@ -9,12 +9,13 @@ Pins (PR2): from __future__ import annotations import asyncio -from unittest.mock import AsyncMock, MagicMock +import json import pytest import litellm from litellm.proxy import proxy_server +from litellm.router_utils import pattern_match_deployments from .conftest import normalize # type: ignore[import-not-found] @@ -99,6 +100,7 @@ def test_token_counter_missing_input_returns_400( @pytest.fixture def patched_supported_params(monkeypatch): + monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr( litellm, "get_llm_provider", @@ -124,12 +126,104 @@ def test_supported_openai_params_happy_path(client, auth_as, patched_supported_p } +def test_supported_openai_params_resolves_router_alias(client, auth_as, monkeypatch): + """A router alias absent from the cost map resolves through the deployment's underlying model.""" + router = litellm.Router( + model_list=[ + { + "model_name": "claude-opus-4-6-cached", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with auth_as(): + response = client.get("/utils/supported_openai_params", params={"model": "claude-opus-4-6-cached"}) + + assert response.status_code == 200 + expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic") + assert response.json() == {"supported_openai_params": expected} + assert "max_tokens" in response.json()["supported_openai_params"] + + +def test_supported_openai_params_declared_prefix_alias_resolves_through_router(client, auth_as, monkeypatch): + """Regression: an alias whose name starts with an authenticating provider's prefix skipped + router resolution and answered with that provider's params instead of the deployment's.""" + router = litellm.Router( + model_list=[ + { + "model_name": "github_copilot/gpt-4o", + "litellm_params": {"model": "anthropic/claude-opus-4-6", "api_key": "sk-test"}, + } + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + with auth_as(): + response = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"}) + + assert response.status_code == 200 + expected = litellm.get_supported_openai_params(model="claude-opus-4-6", custom_llm_provider="anthropic") + assert response.json() == {"supported_openai_params": expected} + + +def test_supported_openai_params_never_runs_oauth_for_authenticating_providers(client, auth_as, monkeypatch, tmp_path): + """Regression: github_copilot/chatgpt names answer from their declaration; resolving them + through ``get_llm_provider`` would run the provider's OAuth device flow and block the event loop.""" + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "access-token").write_text("fake-access-token") + (tmp_path / "api-key.json").write_text( + json.dumps( + { + "token": "fake-api-key", + "expires_at": 4102444800, + "endpoints": {"api": "https://api.githubcopilot.com"}, + } + ) + ) + router = litellm.Router( + model_list=[ + { + "model_name": "copilot-alias", + "litellm_params": {"model": "github_copilot/gpt-4o"}, + }, + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*"}, + }, + ] + ) + monkeypatch.setattr(proxy_server, "llm_router", router) + + resolution_attempts: list[str] = [] + + def _oauth_tripwire(model, *args, **kwargs): + resolution_attempts.append(model) + raise AssertionError("get_llm_provider would run the OAuth device flow") + + monkeypatch.setattr(litellm, "get_llm_provider", _oauth_tripwire) + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire) + expected = litellm.get_supported_openai_params(model="gpt-4o", custom_llm_provider="github_copilot") + + with auth_as(): + via_alias = client.get("/utils/supported_openai_params", params={"model": "copilot-alias"}) + via_direct_name = client.get("/utils/supported_openai_params", params={"model": "github_copilot/gpt-4o"}) + + assert via_alias.status_code == 200 + assert via_alias.json() == {"supported_openai_params": expected} + assert via_direct_name.status_code == 200 + assert via_direct_name.json() == {"supported_openai_params": expected} + assert resolution_attempts == [] + + def test_supported_openai_params_invalid_model(client, auth_as, monkeypatch): """Pins ``GET /utils/supported_openai_params`` (error: unknown model).""" def _raise(model): raise Exception("unknown") + monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "get_llm_provider", _raise) with auth_as(): response = client.get("/utils/supported_openai_params", params={"model": "??"}) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 95067929ac1..b8fb6170d34 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -819,6 +819,94 @@ async def test_should_cap_known_estimate_to_remaining_budget( ) == pytest.approx(0.9) +@pytest.mark.asyncio +async def test_fail_closed_rejects_known_estimate_exceeding_remaining_budget( + spend_counter_state, +): + """LIT-5922: with strict enforcement on, a request whose known estimate does + not fit the remaining budget must be rejected before dispatch instead of + having its reservation shrunk to the headroom and admitted, and the counter + must be restored to the pre-request spend.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-known-estimate-fail-closed", + spend=0.9, + max_budget=1.0, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-known-estimate-fail-closed", + value=0.9, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.6, + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert exc_info.value.current_cost == pytest.approx(0.9) + assert exc_info.value.max_budget == pytest.approx(1.0) + assert "Current cost: 0.9, Estimated request cost: 0.6, Max budget: 1.0" in str(exc_info.value) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-known-estimate-fail-closed" + ) == pytest.approx(0.9) + + +@pytest.mark.asyncio +async def test_fail_closed_tolerates_float_noise_when_estimate_exactly_fits( + spend_counter_state, +): + """0.1 + 0.2 lands a hair above 0.3 in floating point. Strict enforcement + must treat that as fitting the budget, not reject it.""" + counter_cache, key_cache = spend_counter_state + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth( + token="key-budget-fail-closed-float-noise", + spend=0.1, + max_budget=0.3, + ) + counter_cache.in_memory_cache.set_cache( + key="spend:key:key-budget-fail-closed-float-noise", + value=0.1, + ) + + with patch( # test-quality-ok: reserve_budget_for_request takes no estimator, so pinning the estimate needs this attribute + "litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost", + return_value=0.2, + ): + reservation = await reserve_budget_for_request( + request_body=_request_body(), + route="/chat/completions", + llm_router=None, + valid_token=valid_token, + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=key_cache, + proxy_logging_obj=proxy_logging_obj, + fail_closed_budget_enforcement=True, + ) + + assert reservation is not None + assert reservation["reserved_cost"] == pytest.approx(0.2) + assert counter_cache.in_memory_cache.get_cache( + key="spend:key:key-budget-fail-closed-float-noise" + ) == pytest.approx(0.3) + + @pytest.mark.asyncio async def test_should_clamp_reservation_to_default_when_output_cap_missing( spend_counter_state, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 43280258153..91fca8f1e27 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -10445,6 +10445,75 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): assert before["some_api_key"] != "sk-stored-secret" +@pytest.mark.asyncio +async def test_update_config_field_rejects_out_of_range_alerting_args(monkeypatch): + """Out-of-range alerting_args must be rejected at save time. If they land in the + DB, SlackAlertingArgs raises during the config reload and alerting breaks.""" + from unittest.mock import MagicMock + + from fastapi import HTTPException + + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.proxy_server import update_config_general_settings + + monkeypatch.setattr(proxy_server_module, "prisma_client", MagicMock()) + + admin = UserAPIKeyAuth( + api_key="hashed-admin", + user_id="admin-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + with pytest.raises(HTTPException) as exc_info: + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="alerting_args", + field_value={ + "daily_spend_per_user_threshold": -5.0, + "user_spend_check_interval": 20, + }, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + assert exc_info.value.status_code == 400 + error_msg = exc_info.value.detail["error"] + assert "daily_spend_per_user_threshold" in error_msg + assert "user_spend_check_interval" in error_msg + + +@pytest.mark.asyncio +async def test_update_config_field_accepts_valid_alerting_args(monkeypatch): + import litellm.proxy.proxy_server as proxy_server_module + from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.proxy_server import update_config_general_settings + + fake = _fake_prisma_with_config({}) + monkeypatch.setattr(proxy_server_module, "prisma_client", fake) + monkeypatch.setattr(litellm, "store_audit_logs", False) + + admin = UserAPIKeyAuth( + api_key="hashed-admin", + user_id="admin-1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + await update_config_general_settings( + data=ConfigFieldUpdate( + field_name="alerting_args", + field_value={ + "daily_spend_per_user_threshold": 5.0, + "user_spend_check_interval": 60, + }, + config_type="general_settings", + ), + user_api_key_dict=admin, + ) + + written = json.loads(fake.db.litellm_config.upsert.call_args.kwargs["data"]["update"]["param_value"]) + assert written["alerting_args"]["daily_spend_per_user_threshold"] == 5.0 + + @pytest.mark.asyncio async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index b3a2bdda53c..60433921de6 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -598,6 +598,45 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh assert filtered == healthy_deployments +@pytest.mark.asyncio +async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): + user_key = "user-key-order-fallback" + stable_model_map_key = "claude-sonnet-4-5@20250929" + cache = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model_id": "deployment-1"}) + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=123, + enable_user_key_affinity=True, + enable_responses_api_affinity=False, + ) + healthy_deployments = [ + { + "model_name": stable_model_map_key, + "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": stable_model_map_key, + "litellm_params": { + "model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0" + }, + "model_info": {"id": "deployment-2"}, + }, + ] + + filtered = await callback.async_filter_deployments( + model="some-router-model-group", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={"_target_order": 2, "metadata": {"user_api_key_hash": user_key}}, + parent_otel_span=None, + ) + + assert filtered == healthy_deployments + cache.async_get_cache.assert_not_called() + + @pytest.mark.asyncio async def test_async_user_key_affinity_ttl_expiry_allows_reroute(): """ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index f54a1cfa284..79ae00e155c 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -150,6 +150,25 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum(): assert filtered == [deployments[1]] +@pytest.mark.asyncio +async def test_async_filter_deployments_does_not_pin_when_target_order_is_set(): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + messages = _messages(word_count=5000) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"_target_order": 2}, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is_lower(): """ diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 8336926c050..894b2d9e74f 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -11,7 +11,10 @@ from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, _trigger_cooldown_for_failed_deployment, fallback_attempt_key, + clear_pre_routing_selection, get_fallback_model_group, + get_pre_routing_selection, + record_pre_routing_selection, run_async_fallback, ) @@ -1090,3 +1093,119 @@ async def test_run_async_fallback_preserves_original_model_group_on_nested_fallb metadata = router.received_kwargs["metadata"] assert metadata["attempted_fallbacks"] == 2 assert metadata["original_model_group"] == "primary-model" + + +class TestPreRoutingSelectionCarriesToFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup kept using the router name, so the tier's configured chain never ran.""" + + def test_selection_is_recorded_in_the_metadata_bucket(self): + kwargs = {"model": "smart-router", "metadata": {}} + record_pre_routing_selection(kwargs, "tier1") + assert kwargs["metadata"]["pre_routing_selected_model"] == "tier1" + assert get_pre_routing_selection(kwargs) == "tier1" + + def test_selection_is_recorded_in_the_litellm_metadata_bucket(self): + kwargs = {"model": "smart-router", "litellm_metadata": {}} + record_pre_routing_selection(kwargs, "tier2") + assert get_pre_routing_selection(kwargs) == "tier2" + + def test_a_bucket_survives_the_kwargs_copy_that_fallbacks_run_on(self): + """The bucket is shared by reference, which is the whole reason this works.""" + outer = {"model": "smart-router", "metadata": {}} + inner = {**outer} + record_pre_routing_selection(inner, "tier1") + assert get_pre_routing_selection(outer) == "tier1" + + def test_no_selection_reads_as_none(self): + assert get_pre_routing_selection({"model": "smart-router", "metadata": {}}) is None + assert get_pre_routing_selection({"model": "smart-router"}) is None + + def test_missing_kwargs_is_a_no_op(self): + """A caller with no kwargs must not raise, and must not leak the selection anywhere.""" + record_pre_routing_selection(None, "tier1") + + assert get_pre_routing_selection({}) is None + + def test_a_non_dict_bucket_is_ignored(self): + kwargs = {"model": "smart-router", "metadata": "not-a-dict"} + record_pre_routing_selection(kwargs, "tier1") + assert get_pre_routing_selection(kwargs) is None + + def test_fallbacks_resolve_against_the_selected_tier(self): + """The lookup the router performs, keyed on the tier rather than the router name.""" + fallbacks = [{"tier1": ["backup-a", "backup-b"]}, {"tier2": ["backup-c"]}] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="tier1")[0] == ["backup-a", "backup-b"] + assert get_fallback_model_group(fallbacks=fallbacks, model_group="smart-router")[0] is None + + +class TestPreRoutingSelectionIsPerHop: + """#38832 review: the buckets also carry whatever the caller sent, and a fallback hop + inherits the previous hop's tier, so a hop must start without a selection.""" + + def test_a_caller_supplied_selection_is_dropped(self): + kwargs = {"model": "plain", "metadata": {"pre_routing_selected_model": "tier1"}} + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + assert "pre_routing_selected_model" not in kwargs["metadata"] + + def test_both_buckets_are_cleared(self): + kwargs = { + "metadata": {"pre_routing_selected_model": "tier1"}, + "litellm_metadata": {"pre_routing_selected_model": "tier2"}, + } + + clear_pre_routing_selection(kwargs) + + assert get_pre_routing_selection(kwargs) is None + + def test_the_rest_of_the_bucket_is_left_alone(self): + kwargs = {"metadata": {"pre_routing_selected_model": "tier1", "tags": ["a"]}} + + clear_pre_routing_selection(kwargs) + + assert kwargs["metadata"] == {"tags": ["a"]} + + def test_clearing_is_a_no_op_without_a_usable_bucket(self): + kwargs = {"model": "plain", "metadata": "not-a-dict"} + + clear_pre_routing_selection(None) + clear_pre_routing_selection(kwargs) + + assert kwargs == {"model": "plain", "metadata": "not-a-dict"} + + def test_a_selection_recorded_after_clearing_is_kept(self): + """Clearing runs before routing, so the hook's own write must survive it.""" + kwargs = {"model": "smart-router", "metadata": {"pre_routing_selected_model": "stale"}} + + clear_pre_routing_selection(kwargs) + record_pre_routing_selection(kwargs, "tier1") + + assert get_pre_routing_selection(kwargs) == "tier1" + + +class TestOrderedFallbackLookupGroups: + def test_tier_first_then_requested_group_deduped(self): + from litellm.router_utils.fallback_event_handlers import ( + PRE_ROUTING_SELECTED_MODEL_KEY, + fallback_lookup_groups, + ) + + kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier1"}} + assert fallback_lookup_groups(kwargs, "smart-router") == ("tier1", "smart-router") + assert fallback_lookup_groups(kwargs, "tier1") == ("tier1",) + assert fallback_lookup_groups({}, "smart-router") == ("smart-router",) + assert fallback_lookup_groups({}, None) == () + + def test_first_resolving_group_wins_and_generic_idx_survives_a_miss(self): + from litellm.router_utils.fallback_event_handlers import ( + get_fallback_model_group_for_lookup_groups, + ) + + fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}, {"*": ["backup-c"]}] + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier1", "smart-router")) == (["backup-a"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "smart-router")) == (["backup-b"], None) + assert get_fallback_model_group_for_lookup_groups(fallbacks, ("tier9", "no-such")) == (["backup-c"], 2) + assert get_fallback_model_group_for_lookup_groups([{"tier1": ["backup-a"]}], ("no", "nope")) == (None, None) diff --git a/tests/test_litellm/router_utils/test_pattern_match_deployments.py b/tests/test_litellm/router_utils/test_pattern_match_deployments.py new file mode 100644 index 00000000000..795d448ef5f --- /dev/null +++ b/tests/test_litellm/router_utils/test_pattern_match_deployments.py @@ -0,0 +1,78 @@ +"""Behavior pins for ``litellm/router_utils/pattern_match_deployments.py``.""" + +from __future__ import annotations + +from litellm.router_utils import pattern_match_deployments +from litellm.router_utils.pattern_match_deployments import PatternMatchRouter + + +def _wildcard_deployment(model_name: str) -> dict: + return {"model_name": model_name, "litellm_params": {"model": model_name}} + + +def _matched_models(matches: list[dict] | None) -> list[str]: + return [deployment["litellm_params"]["model"] for deployment in matches or []] + + +def test_get_pattern_never_resolves_declared_authenticating_providers(monkeypatch): + """Regression: resolving a github_copilot/chatgpt name through ``get_llm_provider`` runs the + provider's OAuth device flow; the auth layer walks every wildcard router on every request, so + a single metadata lookup for an unserved name would block the proxy's event loop.""" + resolution_attempts: list[str] = [] + + def _oauth_tripwire(model, *args, **kwargs): + resolution_attempts.append(model) + raise AssertionError("get_llm_provider would run the OAuth device flow") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _oauth_tripwire) + + unmatched_router = PatternMatchRouter() + unmatched_router.add_pattern("anthropic/*", _wildcard_deployment("anthropic/*")) + assert unmatched_router.get_pattern("github_copilot/gpt-4o") is None + + matched_router = PatternMatchRouter() + matched_router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*")) + assert _matched_models(matched_router.get_pattern("github_copilot/gpt-4o")) == ["github_copilot/gpt-4o"] + assert _matched_models(matched_router.get_pattern("gpt-4o", custom_llm_provider="github_copilot")) == [ + "github_copilot/gpt-4o" + ] + + assert resolution_attempts == [] + + +def test_get_pattern_bare_provider_name_never_matches_that_providers_wildcard(monkeypatch): + """Regression: a bare ``github_copilot`` adopted itself as its provider and retried as + ``github_copilot/github_copilot``, false-matching the wildcard for a name no deployment serves.""" + + def _unknown_provider(model, *args, **kwargs): + raise ValueError(f"unknown provider for {model}") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider) + router = PatternMatchRouter() + router.add_pattern("github_copilot/*", _wildcard_deployment("github_copilot/*")) + assert router.get_pattern("github_copilot") is None + + +def test_get_pattern_missing_model_returns_none(monkeypatch): + """Regression: a request without a model reaches the auth layer's pattern walk as ``None``; the + declared-provider guard raised ``TypeError`` where the old inline resolve swallowed every + resolver error, so the proxy's missing-model 400 became a crash.""" + + def _unknown_provider(model, *args, **kwargs): + raise ValueError(f"unknown provider for {model}") + + monkeypatch.setattr(pattern_match_deployments, "get_llm_provider", _unknown_provider) + router = PatternMatchRouter() + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + assert router.get_pattern(None) is None + + +def test_get_pattern_still_resolves_unqualified_names(monkeypatch): + monkeypatch.setattr( + pattern_match_deployments, + "get_llm_provider", + lambda model, **kwargs: (model, "openai", None, None), + ) + router = PatternMatchRouter() + router.add_pattern("openai/*", _wildcard_deployment("openai/*")) + assert _matched_models(router.get_pattern("gpt-4o")) == ["openai/gpt-4o"] diff --git a/tests/test_litellm/test_dockerfile_apk_repository.py b/tests/test_litellm/test_dockerfile_apk_repository.py new file mode 100644 index 00000000000..cbd772defbf --- /dev/null +++ b/tests/test_litellm/test_dockerfile_apk_repository.py @@ -0,0 +1,52 @@ +""" +Static checks on the root Dockerfile's apk repository configuration. + +The base image (cgr.dev/chainguard/wolfi-base) only configures the +authenticated Chainguard apk repo (https://apk.cgr.dev/chainguard) in +/etc/apk/repositories, which requires a Chainguard enterprise subscription. +Anyone pulling the published litellm image and running `apk add` inside it +hits SSL/auth failures with no fallback repo configured, so nothing can be +installed. See https://github.com/BerriAI/litellm/issues/33518 +""" + +import os +import re + +import pytest + +DOCKERFILE_PATH = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "Dockerfile", +) + + +def _runtime_stage(dockerfile_text: str) -> str: + """Return the contents of the final `FROM ... AS runtime` build stage.""" + match = re.search(r"^FROM .*\bAS runtime\b(.*)\Z", dockerfile_text, re.MULTILINE | re.DOTALL) + assert match, "Dockerfile has no `FROM ... AS runtime` stage" + return match.group(1) + + +@pytest.mark.skipif( + not os.path.exists(DOCKERFILE_PATH), + reason="Dockerfile not present in this checkout", +) +def test_runtime_stage_adds_public_wolfi_repo(): + """The runtime stage must add the public Wolfi apk repo so `apk add` + works for users without a Chainguard enterprise subscription.""" + with open(DOCKERFILE_PATH, "r", encoding="utf-8") as f: + contents = f.read() + + runtime_stage = _runtime_stage(contents) + + assert re.search( + r"echo\s+[\"']?https://packages\.wolfi\.dev/os[\"']?\s*>>\s*/etc/apk/repositories", + runtime_stage, + ), ( + "Runtime stage must append the public Wolfi apk repo " + '(RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories) ' + "so `apk add` works without Chainguard enterprise credentials. " + "See https://github.com/BerriAI/litellm/issues/33518" + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 44c1cdbff06..84f6344be35 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -11595,3 +11595,138 @@ class TestTierParamsTheTargetAccepts: accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) assert accepted == {"reasoning_effort": "max"} + + +class TestPreRoutingTierDrivesFallbacks: + """#38832: a complexity/auto router picks a tier behind the router name, but fallback + lookup stayed on the router name, so the tier's configured chain never ran and a + provider failure on the tier's first hop was returned to the client.""" + + class _TierRouter(litellm.Router): + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + if model == "smart-router": + return PreRoutingHookResponse(model="tier1", messages=messages) + return None + + @classmethod + def _router(cls, fallbacks) -> "litellm.Router": + return cls._TierRouter( + model_list=[ + { + "model_name": "smart-router", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, + }, + { + "model_name": "tier1", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "backup-a", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-a", + }, + }, + { + "model_name": "backup-b", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "from backup-b", + }, + }, + { + "model_name": "failing-backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + { + "model_name": "plain", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-x", + "mock_response": "litellm.RateLimitError", + }, + }, + ], + fallbacks=fallbacks, + num_retries=0, + ) + + @pytest.mark.asyncio + async def test_the_selected_tier_fallback_chain_runs(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="smart-router", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_is_not_used(self): + """The router name has no chain of its own, so nothing should rescue this call.""" + router = self._router([{"tier2": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + @pytest.mark.asyncio + async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self): + """The documented contract: configs keyed on the requested name keep working behind auto-routers.""" + router = self._router([{"smart-router": ["backup-a"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_the_tier_chain_wins_over_the_router_name_chain(self): + router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): + router = self._router([{"tier1": ["backup-a"]}]) + + response = await router.acompletion( + model="tier1", messages=[{"role": "user", "content": "hi"}] + ) + + assert response.choices[0].message.content == "from backup-a" + + @pytest.mark.asyncio + async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self): + """The metadata bucket carries caller-supplied keys, so only the hook may set the tier.""" + router = self._router([{"tier1": ["backup-a"]}]) + + with pytest.raises(litellm.RateLimitError): + await router.acompletion( + model="plain", + messages=[{"role": "user", "content": "hi"}], + metadata={"pre_routing_selected_model": "tier1"}, + ) + + @pytest.mark.asyncio + async def test_each_fallback_hop_resolves_its_own_chain(self): + """The second hop must key off the group it is running, not the tier that failed.""" + router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}]) + + response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) + + assert response.choices[0].message.content == "from backup-b" diff --git a/tests/test_litellm/test_router_order_fallback.py b/tests/test_litellm/test_router_order_fallback.py index 7743cb005d0..fde870e5abe 100644 --- a/tests/test_litellm/test_router_order_fallback.py +++ b/tests/test_litellm/test_router_order_fallback.py @@ -6,12 +6,19 @@ should be tried first, and higher order deployments should be used as fallbacks when lower order deployments fail. """ -from typing import Optional +import json +from typing import Final, Optional +import httpx import pytest +from openai import AsyncOpenAI +import litellm from litellm import Router -from litellm.utils import _get_order_filtered_deployments +from litellm.integrations.custom_logger import CustomLogger +from litellm.router_utils.prompt_caching_cache import PromptCachingCache +from litellm.types.router import RouterRateLimitError +from litellm.utils import _get_deployment_order, _get_order_filtered_deployments # --------------------------------------------------------------------------- # Unit tests for _get_order_filtered_deployments @@ -49,13 +56,22 @@ class TestGetOrderFilteredDeployments: assert len(result) == 1 assert result[0]["model_info"]["id"] == "b" - def test_target_order_no_match_returns_all(self): + def test_target_order_no_match_returns_empty(self): deps = [ self._make_deployment(1, "a"), self._make_deployment(2, "b"), ] result = _get_order_filtered_deployments(deps, target_order=99) - assert len(result) == 2 + assert result == [] + + def test_target_order_no_match_does_not_reselect_lower_order(self): + deps = [ + self._make_deployment(1, "a"), + self._make_deployment(2, "b"), + ] + remaining_after_pre_call = [deps[0]] + result = _get_order_filtered_deployments(remaining_after_pre_call, target_order=2) + assert result == [] def test_no_order_set_returns_all(self): deps = [ @@ -406,35 +422,239 @@ async def test_router_order_fallback_with_hidden_model_group_alias(): assert response._hidden_params["model_id"] == "2" +@pytest.mark.asyncio +async def test_router_order_fallback_does_not_reselect_order_1_when_order_2_is_filtered_out(): + class _DropOrder2(CustomLogger): + async def async_filter_deployments( + self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None + ): + return [d for d in healthy_deployments if _get_deployment_order(d) != 2] + + drop_order_2: Final = _DropOrder2() + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "key", + "mock_response": "litellm.RateLimitError", + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "key", + "mock_response": "success from order 2", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + ) + litellm.callbacks.append(drop_order_2) + try: + with pytest.raises(RouterRateLimitError, match="No deployments available") as exc_info: + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + assert "success from order 2" not in str(exc_info.value) + finally: + litellm.callbacks.remove(drop_order_2) + + +@pytest.mark.asyncio +async def test_router_order_fallback_ignores_prompt_cache_pin_on_target_order(): + messages = [{"role": "user", "content": "word " * 5000}] + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("azure peak load"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "good", + "mock_response": "success from order 2", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + optional_pre_call_checks=["prompt_caching"], + ) + await PromptCachingCache(cache=router.cache).async_add_model_id( + model_id="1", + messages=messages, + tools=None, + ) + response = await router.acompletion(model="test-model", messages=messages) + assert response._hidden_params["model_id"] == "2" + + +@pytest.mark.asyncio +async def test_router_order_fallback_retries_keep_target_order(): + seen_target_orders: Final = [] + + class _RecordTargetOrder(CustomLogger): + async def async_filter_deployments( + self, model, healthy_deployments, messages, request_kwargs=None, parent_otel_span=None + ): + seen_target_orders.append((request_kwargs or {}).get("_target_order")) + return healthy_deployments + + recorder: Final = _RecordTargetOrder() + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("fail order 1"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "gpt-4o", + "api_key": "bad", + "mock_response": Exception("fail order 2"), + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=1, + ) + litellm.callbacks.append(recorder) + try: + with pytest.raises(Exception, match="fail order 2"): + await router.acompletion( + model="test-model", + messages=[{"role": "user", "content": "hi"}], + ) + finally: + litellm.callbacks.remove(recorder) + assert seen_target_orders.count(2) >= 2 + + +@pytest.mark.asyncio +async def test_generic_api_call_strips_target_order_from_provider_kwargs(): + captured: Final = {} + + async def _fake_provider(**provider_kwargs): + captured.update(provider_kwargs) + return "ok" + + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": {"model": "gpt-4o", "api_key": "key", "order": 2}, + "model_info": {"id": "2"}, + }, + ], + ) + response = await router._ageneric_api_call_with_fallbacks_helper( + model="test-model", + original_generic_function=_fake_provider, + _target_order=2, + messages=[{"role": "user", "content": "hi"}], + ) + assert response == "ok" + assert captured["model"] == "gpt-4o" + assert "_target_order" not in captured + + +@pytest.mark.asyncio +async def test_text_completion_order_fallback_hop_does_not_send_target_order_upstream(): + upstream_bodies: Final[list[dict]] = [] + + def _upstream(request: httpx.Request) -> httpx.Response: + upstream_bodies.append(json.loads(request.content)) + return httpx.Response( + 200, + json={ + "id": "cmpl-1", + "object": "text_completion", + "created": 0, + "model": "gpt-3.5-turbo-instruct", + "choices": [{"text": "ok from order 2", "index": 0, "logprobs": None, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + + upstream_client: Final = AsyncOpenAI( + api_key="key", + base_url="http://upstream.test", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(_upstream)), + ) + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_key": "key", + "mock_response": Exception("fail order 1"), + "order": 1, + }, + "model_info": {"id": "1"}, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "text-completion-openai/gpt-3.5-turbo-instruct", + "api_key": "key", + "api_base": "http://upstream.test", + "order": 2, + }, + "model_info": {"id": "2"}, + }, + ], + num_retries=0, + ) + try: + response = await router.atext_completion(model="test-model", prompt="hi", client=upstream_client) + finally: + await upstream_client.close() + + assert response._hidden_params["model_id"] == "2" + assert upstream_bodies + assert all("_target_order" not in body for body in upstream_bodies) + + def test_check_non_standard_fallback_format(): from litellm.router_utils.fallback_event_handlers import ( _check_non_standard_fallback_format, ) # Standard formats - assert ( - _check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) - == False - ) + assert _check_non_standard_fallback_format([{"gpt-3.5-turbo": ["claude-3-haiku"]}]) == False assert _check_non_standard_fallback_format([{"model": ["qwen-backup"]}]) == False - assert ( - _check_non_standard_fallback_format( - [{"model": ["qwen-backup"], "region": ["us-east-1"]}] - ) - == False - ) + assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "region": ["us-east-1"]}]) == False # Non-standard formats assert _check_non_standard_fallback_format([{"model": "qwen-backup"}]) == True assert ( - _check_non_standard_fallback_format( - [{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}] - ) - == True - ) - assert ( - _check_non_standard_fallback_format( - [{"model": ["qwen-backup"], "api_key": "some-key"}] - ) + _check_non_standard_fallback_format([{"model": "qwen-backup", "messages": [{"role": "user", "content": "hi"}]}]) == True ) + assert _check_non_standard_fallback_format([{"model": ["qwen-backup"], "api_key": "some-key"}]) == True diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx index 8c683f230e0..f460b77c2c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.test.tsx @@ -140,6 +140,7 @@ const renderPanel = (canModify = true) => accessToken="token" userRole="Admin" userID="u-admin" + isViewOnly={false} teams={null} createScope={canModify ? "unscoped-ok" : "forbidden"} />, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx index db5120cebce..5b53217f9c1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/AutoRoutersPanel.tsx @@ -21,12 +21,20 @@ interface AutoRoutersPanelProps { accessToken: string; userRole: string; userID: string | null; + isViewOnly: boolean; teams: Team[] | null; /** Owned by the page, which knows how this caller must scope what they create. */ createScope: ModelWriteScope; } -export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createScope }: AutoRoutersPanelProps) { +export function AutoRoutersPanel({ + accessToken, + userRole, + userID, + isViewOnly, + teams, + createScope, +}: AutoRoutersPanelProps) { const canCreate = createScope !== "forbidden"; const { data: deployments, isLoading } = useAutoRouters(); const invalidateAutoRouters = useInvalidateAutoRouters(); @@ -39,8 +47,8 @@ export function AutoRoutersPanel({ accessToken, userRole, userID, teams, createS const [isDeleting, setIsDeleting] = useState(false); const routers = useMemo( - () => toAutoRouterRows(deployments ?? [], { userRole, userID }, teams), - [deployments, userRole, userID, teams], + () => toAutoRouterRows(deployments ?? [], { userRole, userID, isViewOnly }, teams), + [deployments, userRole, userID, isViewOnly, teams], ); const handleCreated = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts index 9944653b638..23585f6c110 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AutoRouters/autoRouterRows.test.ts @@ -5,8 +5,9 @@ import { toAutoRouterRow, toAutoRouterRows } from "./autoRouterRows"; // Existing cases assert resource classification, so they run as a proxy admin: the actor // gate is then a pass-through and canEdit/canDelete still reflect the row itself. -const ADMIN = { userRole: "Admin", userID: "u-admin" }; -const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" }; +const ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false }; +const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false }; +const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true }; const complexityDeployment = { model_name: "tri-tier-router", @@ -224,7 +225,7 @@ describe("autoRouterRows actor gating", () => { { team_id: "team-1", members_with_roles: [{ user_id: "u-team-admin", user_email: "t@t", role: "admin" }] }, ] as never; - const rowIn = (actor: { userRole: string; userID: string }, teamId: string | null) => + const rowIn = (actor: { userRole: string; userID: string; isViewOnly: boolean }, teamId: string | null) => toAutoRouterRow( { ...complexityDeployment, model_info: { id: "cid-1", db_model: true, team_id: teamId } }, 0, @@ -259,4 +260,12 @@ describe("autoRouterRows actor gating", () => { expect(row.canEdit).toBe(true); expect(row.canDelete).toBe(true); }); + + // A proxy_admin_viewer session reads "Admin" through the masquerade, but PATCH and + // DELETE both 403 it, so its rows must not offer the affordances. + it("hides write affordances from a view-only admin session", () => { + const row = rowIn(VIEW_ONLY_ADMIN, null); + expect(row.canEdit).toBe(false); + expect(row.canDelete).toBe(false); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx index a504d75bb63..84a05113177 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.test.tsx @@ -38,8 +38,17 @@ vi.mock("./useModelDashboardData", () => ({ useModelDashboardData: () => ({ availableModelAccessGroups: [], allModelsOnProxy: [], availableModelGroups: [] }), })); -const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false }; -const NON_ADMIN = { accessToken: "at", token: "t", userRole: "Internal User", userId: "u1", premiumUser: false }; +const ADMIN = { accessToken: "at", token: "t", userRole: "Admin", userId: "u1", premiumUser: false, isViewOnly: false }; +const NON_ADMIN = { + accessToken: "at", + token: "t", + userRole: "Internal User", + userId: "u1", + premiumUser: false, + isViewOnly: false, +}; +// A proxy_admin_viewer session: effectiveSessionRole masquerades the role as "Admin". +const VIEW_ONLY_ADMIN = { ...ADMIN, isViewOnly: true }; const renderPage = () => { const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } } }); @@ -99,6 +108,22 @@ describe("ModelsAndEndpointsPage", () => { expect(screen.queryByRole("tab", { name: "Health Status" })).not.toBeInTheDocument(); }); + // POST /model/new 403s a proxy_admin_viewer, so the form's tab must not render for one. + it("hides the Add Model tab for a view-only admin session", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.queryByRole("tab", { name: "Add Model" })).not.toBeInTheDocument(); + expect(screen.getByRole("tab", { name: "All Models" })).toBeInTheDocument(); + }); + + // Read parity: the Auto-Routers list stays reachable for a view-only admin; only the + // create affordance inside it is withheld, which AutoRoutersTabPanel decides. + it("keeps the Auto-Routers tab for a view-only admin session", () => { + mockUseAuthorized.mockReturnValue(VIEW_ONLY_ADMIN); + renderPage(); + expect(screen.getByRole("tab", { name: /Auto-Routers/ })).toBeInTheDocument(); + }); + // Auto-routers are excluded from the All Models table, so this tab is their home: the only // place in the product to list, create, edit or delete one. describe("Auto-Routers tab", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 9ae7dc12f81..34c9d87004e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -80,7 +80,7 @@ const renderPanel = (key: string) => { }; export default function ModelsAndEndpointsPage() { - const { accessToken, userRole, userId: userID, premiumUser } = useAuthorized(); + const { accessToken, userRole, userId: userID, premiumUser, isViewOnly } = useAuthorized(); const { data: teams } = useTeams(); const { data: uiSettings } = useUISettings(); const queryClient = useQueryClient(); @@ -92,7 +92,7 @@ export default function ModelsAndEndpointsPage() { const isInternalUser = userRole && internalUserRoles.includes(userRole); const canCreate = canCreateModels( - { userRole, userID }, + { userRole, userID, isViewOnly }, { teams: teams ?? null, disabledForInternalUsers: @@ -182,6 +182,7 @@ export default function ModelsAndEndpointsPage() { accessToken={accessToken} userID={userID} userRole={userRole} + isViewOnly={isViewOnly} onModelUpdate={invalidateModels} modelAccessGroups={availableModelAccessGroups} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx new file mode 100644 index 00000000000..12f0b95bf13 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.test.tsx @@ -0,0 +1,39 @@ +/* @vitest-environment jsdom */ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import AutoRoutersTabPanel from "./AutoRoutersTabPanel"; + +const panelProps = vi.fn(); +vi.mock("../components/AutoRouters/AutoRoutersPanel", () => ({ + AutoRoutersPanel: (props: Record) => { + panelProps(props); + return
; + }, +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => mockUseAuthorized() })); +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ useTeams: () => ({ data: [] }) })); +vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ + useUISettings: () => ({ data: { values: {} } }), +})); + +const SESSION = { accessToken: "at", userRole: "Admin", userId: "u1", isViewOnly: false }; + +const lastProps = () => panelProps.mock.calls.at(-1)?.[0] as { createScope: string }; + +describe("AutoRoutersTabPanel", () => { + it("grants an unscoped create to a real proxy admin", () => { + mockUseAuthorized.mockReturnValue(SESSION); + render(); + expect(lastProps().createScope).toBe("unscoped-ok"); + }); + + // The masqueraded "Admin" a proxy_admin_viewer session carries: POST /model/new 403s it, + // so the panel must not be told it may create. + it("withholds the create affordance from a view-only admin session", () => { + mockUseAuthorized.mockReturnValue({ ...SESSION, isViewOnly: true }); + render(); + expect(lastProps().createScope).toBe("forbidden"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx index 5f7d56e8e33..69b442da09b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/panels/AutoRoutersTabPanel.tsx @@ -15,13 +15,13 @@ import { AutoRoutersPanel } from "../components/AutoRouters/AutoRoutersPanel"; * Viewer roles reach the list without write affordances. */ export default function AutoRoutersTabPanel() { - const { accessToken, userRole, userId: userID } = useAuthorized(); + const { accessToken, userRole, userId: userID, isViewOnly } = useAuthorized(); const { data: teams } = useTeams(); const { data: uiSettings } = useUISettings(); const isInternalUser = userRole != null && internalUserRoles.includes(userRole); const scope = modelCreationScope( - { userRole, userID }, + { userRole, userID, isViewOnly }, { teams: teams ?? null, disabledForInternalUsers: isInternalUser && uiSettings?.values?.disable_model_add_for_internal_users === true, @@ -33,6 +33,7 @@ export default function AutoRoutersTabPanel() { accessToken={accessToken} userRole={userRole ?? ""} userID={userID ?? null} + isViewOnly={isViewOnly} teams={teams ?? null} createScope={scope} /> diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index ad0f749b189..92951b68fd5 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -82,7 +82,7 @@ const AddModelForm: React.FC = ({ // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test const [connectionTestId, setConnectionTestId] = useState(""); - const { accessToken, userRole, premiumUser, userId } = useAuthorized(); + const { accessToken, userRole, premiumUser, userId, isViewOnly } = useAuthorized(); const { data: providerMetadata, isLoading: isProviderMetadataLoading, @@ -157,7 +157,10 @@ const AddModelForm: React.FC = ({ const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); // Same owner the Auto-Routers tab uses, so the two creation forms cannot disagree about // who has to name a team. This form is only reachable when creation is allowed at all. - const createScope = modelCreationScope({ userRole, userID: userId }, { teams, disabledForInternalUsers: false }); + const createScope = modelCreationScope( + { userRole, userID: userId, isViewOnly }, + { teams, disabledForInternalUsers: false }, + ); const requiresTeamScope = createScope === "team-required"; return ( diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index d8a955b719c..ba380662403 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -101,18 +101,24 @@ vi.mock("./build_complexity_router_config", async (importOriginal) => { }); // A real TeamDropdown fetches teams and renders an antd Select; the wiring under test is -// whether team_id is registered, validated and forwarded, so a plain control stands in. +// whether team_id is registered, validated and forwarded, so a plain control stands in. The +// clear button mirrors the real dropdown's x, which emits null rather than a string. vi.mock("../common_components/team_dropdown", () => ({ - default: ({ value, onChange }: { value?: string; onChange?: (next: string) => void }) => ( - + default: ({ value, onChange }: { value?: string; onChange?: (next: string | null) => void }) => ( + <> + + + ), })); @@ -354,6 +360,25 @@ describe("AddAutoRouterTab", () => { expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); }); + // The shared dropdown emits null on clear while this form's schema wants a string, so the + // form maps null back to "": the user sees the pick-a-team message, not a zod type error. + it("treats a team picked and then cleared like no team at all", async () => { + const user = userEvent.setup(); + vi.mocked(getMissingTiersError).mockReturnValue(null); + + renderWithProviders( + , + ); + + await user.type(screen.getByPlaceholderText(/smart_router/i), "team-scoped-router"); + await user.selectOptions(screen.getByTestId("team-dropdown"), "team-1"); + await user.click(screen.getByTestId("team-dropdown-clear")); + await user.click(screen.getByRole("button", { name: /add auto router/i })); + + expect(await screen.findByText("Please select a team to continue")).toBeInTheDocument(); + expect(handleAddAutoRouterSubmit).not.toHaveBeenCalled(); + }); + it("defaults a new router to session affinity off, matching the backend field default", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 318adcce369..03b7432e4ec 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -548,7 +548,9 @@ const AddAutoRouterTab: React.FC = ({ "Select the team this auto router belongs to. Only keys for this team will be able to call it.", )} > - {({ id, value, onChange }) => } + {({ id, value, onChange }) => ( + onChange(next ?? "")} /> + )} )} diff --git a/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx b/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx index 98e876801ac..cff33546fbc 100644 --- a/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx +++ b/ui/litellm-dashboard/src/components/alerting/alerting_settings.tsx @@ -5,6 +5,7 @@ import React, { useState, useEffect } from "react"; import { alertingSettingsCall, updateConfigFieldSetting } from "../networking"; import DynamicForm from "./dynamic_form"; +import { extractProxyErrorMessage } from "@/lib/http/client"; import { toast } from "@/lib/toast"; interface alertingSettingsItem { field_name: string; @@ -43,7 +44,7 @@ const AlertingSettings: React.FC = ({ accessToken, premiu setAlertingSettings(updatedSettings); }; - const handleSubmit = (formValues: Record) => { + const handleSubmit = async (formValues: Record) => { if (!accessToken) { return; } @@ -64,18 +65,18 @@ const AlertingSettings: React.FC = ({ accessToken, premiu const mergedFormValues = { ...formValues, ...initialFormValues }; const { slack_alerting, ...alertingArgs } = mergedFormValues; try { - updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs); + await updateConfigFieldSetting(accessToken, "alerting_args", alertingArgs); if (typeof slack_alerting === "boolean") { if (slack_alerting == true) { - updateConfigFieldSetting(accessToken, "alerting", ["slack"]); + await updateConfigFieldSetting(accessToken, "alerting", ["slack"]); } else { - updateConfigFieldSetting(accessToken, "alerting", []); + await updateConfigFieldSetting(accessToken, "alerting", []); } } // update value in state toast.success("Wait 10s for proxy to update."); } catch (error) { - // do something + toast.error(extractProxyErrorMessage(error)); } }; diff --git a/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx b/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx index f9b9765f59e..7764052173c 100644 --- a/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/alerting/dynamic_form.integration.test.tsx @@ -38,6 +38,14 @@ const SETTINGS: Setting[] = [ stored_in_db: null, premium_field: false, }, + { + field_name: "daily_spend_per_user_threshold", + field_description: "Daily spend threshold per user", + field_type: "Float", + field_value: 5.5, + stored_in_db: true, + premium_field: false, + }, ]; const renderForm = ( @@ -170,6 +178,19 @@ describe("DynamicForm change notifications", () => { expect(handleInputChange).toHaveBeenCalledWith("daily_report_frequency", 128); }); + it("renders a Float field as a decimal-friendly number input and reports changes as numbers", async () => { + const user = userEvent.setup(); + const { handleInputChange } = renderForm(); + + const input = screen.getByDisplayValue("5.5"); + expect(input).toHaveAttribute("type", "number"); + expect(input).toHaveAttribute("step", "any"); + + await user.type(input, "1"); + + expect(handleInputChange).toHaveBeenCalledWith("daily_spend_per_user_threshold", 5.51); + }); + it("reports a reset with the field name and its row index", async () => { const user = userEvent.setup(); const { handleResetField } = renderForm(); @@ -216,7 +237,7 @@ describe("DynamicForm presentation", () => { expect(screen.getByText("daily_report_frequency")).toBeInTheDocument(); expect(screen.getByText("How often the report runs")).toBeInTheDocument(); - expect(screen.getByText("In DB")).toBeInTheDocument(); + expect(screen.getAllByText("In DB")).toHaveLength(2); expect(screen.getByText("In Config")).toBeInTheDocument(); expect(screen.getByText("Not Set")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx b/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx index be5e76ab0df..42aa58ca0ea 100644 --- a/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx +++ b/ui/litellm-dashboard/src/components/alerting/dynamic_form.tsx @@ -63,11 +63,11 @@ const DynamicForm: React.FC = ({ }; const renderControl = (setting: AlertingSetting) => { - if (setting.field_type === "Integer") { + if (setting.field_type === "Integer" || setting.field_type === "Float") { return ( handleNumericChange(setting, event.target.value)} /> diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx new file mode 100644 index 00000000000..01b8a12d99c --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.test.tsx @@ -0,0 +1,48 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; + +import { chooseSelectOption } from "../../../tests/test-utils"; +import type { Team } from "../key_team_helpers/key_list"; +import TeamDropdown from "./team_dropdown"; + +const TEAMS = [ + { team_id: "team-1", team_alias: "Alpha Team" }, + { team_id: "team-2", team_alias: "Beta Team" }, +] as unknown as Team[]; + +vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({ + useInfiniteTeams: () => ({ + data: { pages: [{ teams: TEAMS }] }, + fetchNextPage: vi.fn(), + hasNextPage: false, + isFetchingNextPage: false, + isLoading: false, + }), +})); + +describe("TeamDropdown", () => { + it("emits the picked team's id and full object", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await chooseSelectOption(user, screen.getByRole("combobox"), /^Beta Team/); + + expect(onChange).toHaveBeenCalledWith("team-2"); + expect(onTeamSelect).toHaveBeenCalledWith(TEAMS[1]); + }); + + it("emits null, never the empty string, when the selection is cleared", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const onTeamSelect = vi.fn(); + render(); + + await user.click(screen.getByRole("button", { name: "Clear" })); + + expect(onChange).toHaveBeenCalledWith(null); + expect(onTeamSelect).toHaveBeenCalledWith(null); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx index 35121f41598..7d385c2a3f7 100644 --- a/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/team_dropdown.tsx @@ -5,7 +5,7 @@ import { Team } from "../key_team_helpers/key_list"; interface TeamDropdownProps { value?: string; - onChange?: (value: string) => void; + onChange?: (value: string | null) => void; /** Callback with the full Team object (or null on clear). */ onTeamSelect?: (team: Team | null) => void; disabled?: boolean; @@ -47,7 +47,7 @@ const TeamDropdown: React.FC = ({ }, [data]); const handleChange = (teamId: string) => { - onChange?.(teamId); + onChange?.(teamId || null); if (onTeamSelect) { onTeamSelect(teamId ? teams.find((t) => t.team_id === teamId) ?? null : null); } diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 742d2593ade..768183907db 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -87,6 +87,7 @@ describe("ModelInfoView", () => { accessToken: "test-token", userID: "123", userRole: "Admin", + isViewOnly: false, onModelUpdate: vi.fn(), modelAccessGroups: ["group1", "group2"], }; @@ -328,6 +329,16 @@ describe("ModelInfoView", () => { }); }); + // A proxy_admin_viewer session reads "Admin" through effectiveSessionRole, but the update + // and delete endpoints 403 it, so the write buttons must not be offered. + it("should disable delete and update buttons for a view-only admin session", async () => { + render(, { wrapper }); + await waitFor(() => { + expect(screen.getByTestId("delete-model-button")).toBeDisabled(); + }); + expect(screen.getByTestId("update-api-key-button")).toBeDisabled(); + }); + it("should disable delete button when model is not a DB model", async () => { const nonDbModelData = { ...defaultModelData, diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index f21e98e084c..35afcdb2985 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -53,6 +53,7 @@ interface ModelInfoViewProps { accessToken: string | null; userID: string | null; userRole: string | null; + isViewOnly: boolean; onModelUpdate?: (updatedModel: any) => void; modelAccessGroups: string[] | null; } @@ -117,6 +118,7 @@ export default function ModelInfoView({ accessToken, userID, userRole, + isViewOnly, onModelUpdate, modelAccessGroups, }: ModelInfoViewProps) { @@ -167,7 +169,7 @@ export default function ModelInfoView({ // Keep modelData variable name for backwards compatibility const modelData = transformedModelData; - const canEditModel = canModifyModel({ userRole, userID }, teams ?? null, { + const canEditModel = canModifyModel({ userRole, userID, isViewOnly }, teams ?? null, { teamId: modelData?.model_info?.team_id, isDbModel: modelData?.model_info?.db_model === true, }); diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 72da01d0918..9cb55b6ed5c 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -293,6 +293,8 @@ const Settings: React.FC = ({ accessToken, userRole, userID, llm_too_slow: "LLM Responses Too Slow", llm_requests_hanging: "LLM Requests Hanging", budget_alerts: "Budget Alerts (API Keys, Users)", + user_spend_thresholds: "User Spend Thresholds (Daily/Monthly)", + user_spend_anomalies: "User Spend Anomaly Detection", db_exceptions: "Database Exceptions (Read/Write)", daily_reports: "Weekly/Monthly Spend Reports", outage_alerts: "Outage Alerts", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7316b4359cc..6e38f3fa15e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22962,7 +22962,7 @@ export interface components { * @description Enum for alert types and management event types * @enum {string} */ - AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted"; + AlertType: "llm_exceptions" | "llm_too_slow" | "llm_requests_hanging" | "budget_alerts" | "spend_reports" | "failed_tracking_spend" | "user_spend_thresholds" | "user_spend_anomalies" | "db_exceptions" | "daily_reports" | "cooldown_deployment" | "new_model_added" | "model_deprecation_warnings" | "outage_alerts" | "region_outage_alerts" | "fallback_reports" | "new_virtual_key_created" | "virtual_key_updated" | "virtual_key_deleted" | "new_team_created" | "team_updated" | "team_deleted" | "new_internal_user_created" | "internal_user_updated" | "internal_user_deleted"; /** AllowedVectorStoreIndexItem */ AllowedVectorStoreIndexItem: { /** Index Name */ diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts index 179a9b7933a..afc2ecd210f 100644 --- a/ui/litellm-dashboard/src/utils/modelPermissions.test.ts +++ b/ui/litellm-dashboard/src/utils/modelPermissions.test.ts @@ -1,14 +1,16 @@ import { describe, expect, it } from "vitest"; import { Team } from "@/components/networking"; -import { canModifyModel, modelCreationScope } from "./modelPermissions"; +import { canCreateModels, canModifyModel, modelCreationScope } from "./modelPermissions"; const teamWhere = (userId: string, role: string, teamId = "team-1"): Team[] => [{ team_id: teamId, members_with_roles: [{ user_id: userId, user_email: "t@test.com", role }] }] as unknown as Team[]; -const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin" }; -const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin" }; -const MEMBER = { userRole: "Internal User", userID: "u-member" }; +const PROXY_ADMIN = { userRole: "Admin", userID: "u-admin", isViewOnly: false }; +const TEAM_ADMIN = { userRole: "Internal User", userID: "u-team-admin", isViewOnly: false }; +const MEMBER = { userRole: "Internal User", userID: "u-member", isViewOnly: false }; +// proxy_admin_viewer sessions: effectiveSessionRole masquerades the role as "Admin". +const VIEW_ONLY_ADMIN = { userRole: "Admin", userID: "u-viewer", isViewOnly: true }; const noLimits = { disabledForInternalUsers: false }; @@ -40,9 +42,24 @@ describe("modelCreationScope", () => { // an unscoped create from them 403s. Treating them as admins here is what let a form submit // a payload the backend always rejected. it("does not treat an org admin as able to create unscoped", () => { - const orgAdmin = { userRole: "org_admin", userID: "u-org" }; + const orgAdmin = { userRole: "org_admin", userID: "u-org", isViewOnly: false }; expect(modelCreationScope(orgAdmin, { teams: teamWhere("u-org", "admin"), ...noLimits })).toBe("team-required"); }); + + // Server-side, POST /model/new 403s the viewer roles, so the "Admin" the masquerade + // reports must not read as a proxy admin here. + it("forbids a view-only admin session despite the masqueraded Admin role", () => { + expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe("forbidden"); + expect(canCreateModels(VIEW_ONLY_ADMIN, { teams: [], ...noLimits })).toBe(false); + }); + + // _check_proxy_admin_viewer_access (route_checks.py) 403s /model/new on the session role + // alone, before the team-scoped carve-out in ModelManagementAuthChecks can run. + it("forbids a view-only admin even when they admin a team", () => { + expect(modelCreationScope(VIEW_ONLY_ADMIN, { teams: teamWhere("u-viewer", "admin"), ...noLimits })).toBe( + "forbidden", + ); + }); }); describe("canModifyModel", () => { @@ -80,6 +97,17 @@ describe("canModifyModel", () => { }); it("does not treat two absent identities as a match", () => { - expect(canModifyModel({ userRole: "Internal User", userID: null }, null, teamRow)).toBe(false); + expect(canModifyModel({ userRole: "Internal User", userID: null, isViewOnly: false }, null, teamRow)).toBe(false); + }); + + // PATCH /model/{id}/update and POST /model/delete 403 the viewer roles like /model/new does. + it("refuses a view-only admin session on a DB row", () => { + expect(canModifyModel(VIEW_ONLY_ADMIN, null, teamRow)).toBe(false); + }); + + // The route RBAC blocks /model/update and /model/delete for the viewer role before the + // team-scoped carve-out runs, so team-admin membership changes nothing here either. + it("refuses a view-only user even when they admin the owning team", () => { + expect(canModifyModel(VIEW_ONLY_ADMIN, teamWhere("u-viewer", "admin"), teamRow)).toBe(false); }); }); diff --git a/ui/litellm-dashboard/src/utils/modelPermissions.ts b/ui/litellm-dashboard/src/utils/modelPermissions.ts index b5914f9d7ea..843d9041026 100644 --- a/ui/litellm-dashboard/src/utils/modelPermissions.ts +++ b/ui/litellm-dashboard/src/utils/modelPermissions.ts @@ -3,20 +3,32 @@ import { Team } from "@/components/networking"; import { isProxyAdminRole, isUserTeamAdminForAnyTeam, isUserTeamAdminForSingleTeam } from "./roles"; /** - * The dashboard's mirror of ModelManagementAuthChecks in - * litellm/proxy/management_endpoints/model_management_endpoints.py. + * The dashboard's mirror of the two server layers that gate model writes: the role-level + * route RBAC (`_check_proxy_admin_viewer_access` in litellm/proxy/auth/route_checks.py), + * which 403s /model/new, /model/update, and /model/delete for every view-only session + * before the endpoint runs, and ModelManagementAuthChecks in + * litellm/proxy/management_endpoints/model_management_endpoints.py behind it. * - * Both questions below are answered there by exactly two inputs: the caller's role, and - * whether the caller admins the team named in `model_info.team_id`. `created_by` is written - * at creation and never read by an auth check, so it is deliberately absent here; gating on - * it hid controls from team admins the API accepts, and showed controls to former team admins - * the API rejects. + * Past that route gate, both questions below are answered by exactly two inputs: the + * caller's role, and whether the caller admins the team named in `model_info.team_id`. + * `created_by` is written at creation and never read by an auth check, so it is deliberately + * absent here; gating on it hid controls from team admins the API accepts, and showed + * controls to former team admins the API rejects. */ export interface ModelActor { userRole: string | null; userID: string | null; + /** + * From useAuthorized(). A proxy_admin_viewer session masquerades as "Admin" in userRole + * (effectiveSessionRole, for read parity), yet every management write 403s it, so the role + * alone cannot answer a write question. + */ + isViewOnly: boolean; } +const isWritableProxyAdmin = ({ userRole, isViewOnly }: ModelActor): boolean => + !isViewOnly && userRole != null && isProxyAdminRole(userRole); + /** How this actor must scope a deployment they create, or that they may not create one. */ export type ModelWriteScope = "forbidden" | "unscoped-ok" | "team-required"; @@ -33,20 +45,25 @@ const isTeamAdminOf = (teams: Team[] | null, userID: string, teamId: string): bo /** * POST /model/new takes a proxy admin unconditionally, or a team admin whose payload names a - * team; an unscoped create from anyone else is a 403. Returning the requirement rather than a - * pair of booleans keeps "may not create" and "may create unscoped" from being confused. + * team; an unscoped create from anyone else is a 403. A view-only session is 403d by the + * route RBAC on its role alone, so team-admin membership cannot rescue it. Returning the + * requirement rather than a pair of booleans keeps "may not create" and "may create + * unscoped" from being confused. */ export const modelCreationScope = ( - { userRole, userID }: ModelActor, + actor: ModelActor, { teams, disabledForInternalUsers }: ModelCreationLimits, ): ModelWriteScope => { - if (userRole != null && isProxyAdminRole(userRole)) { + if (actor.isViewOnly) { + return "forbidden"; + } + if (isWritableProxyAdmin(actor)) { return "unscoped-ok"; } if (disabledForInternalUsers) { return "forbidden"; } - if (userID != null && isUserTeamAdminForAnyTeam(teams, userID)) { + if (actor.userID != null && isUserTeamAdminForAnyTeam(teams, actor.userID)) { return "team-required"; } return "forbidden"; @@ -63,18 +80,18 @@ export interface ModelRowOrigin { /** May this actor edit or delete this specific deployment? */ export const canModifyModel = ( - { userRole, userID }: ModelActor, + actor: ModelActor, teams: Team[] | null, { teamId, isDbModel }: ModelRowOrigin, ): boolean => { - if (!isDbModel) { + if (actor.isViewOnly || !isDbModel) { return false; } - if (userRole != null && isProxyAdminRole(userRole)) { + if (isWritableProxyAdmin(actor)) { return true; } - if (userID == null || teamId == null) { + if (actor.userID == null || teamId == null) { return false; } - return isTeamAdminOf(teams, userID, teamId); + return isTeamAdminOf(teams, actor.userID, teamId); };